Require Import Omega.
Require Import Wellfounded.

Inductive Exp : Set :=
  | elit : nat -> Exp
  | evar : nat -> Exp
  | eadd : Exp -> Exp -> Exp.

Fixpoint eval (env : nat -> nat) (e : Exp) :=
  match e with
  | elit n => n
  | evar n => env n
  | eadd a b => eval env a + eval env b
  end.

Inductive RPNOp :=
  | rpnlit : nat -> RPNOp
  | rpnvar : nat -> RPNOp
  | rpnadd : RPNOp.

Fixpoint rpn (e : Exp) :=
  match e with
  | elit n => cons (rpnlit n) nil
  | evar n => cons (rpnvar n) nil
  | eadd a b => app (rpn b) (app (rpn a) (cons rpnadd nil))
  end.

Definition _fmap A B (f : A -> B) (o : option A) :=
  match o with
  | Some x => Some (f x)
  | None   => None
  end.
Notation fmap := (_fmap _ _).

Fixpoint rpn_eval' (env : nat -> nat) (rpn : list RPNOp) :=
  match rpn with
  | nil => Some nil
  | cons (rpnlit n) a => fmap (cons n) (rpn_eval' env a)
  | cons (rpnvar n) a => fmap (cons (env n)) (rpn_eval' env a)
  | cons (rpnadd) a => match rpn_eval' env a with
    | Some (cons x (cons y stack)) => Some (cons (x+y) stack)
    | _                            => None
    end
  end.

Function rpn_eval (env : nat -> nat) (rpn : list RPNOp) :=
  match rpn_eval' env (List.rev rpn) with
  | Some (cons x nil) => Some x
  | _                 => None
  end.

Fixpoint exprsize e :=
  match e with
  | elit _ => 1
  | evar _ => 1
  | eadd a b => 1 + exprsize a + exprsize b
  end.

Lemma correctness_partial_compilation e :
  forall env rest results,
  Some results = rpn_eval' env rest ->
  Some (eval env e :: results)%list = rpn_eval' env (app (List.rev (rpn e)) rest).
Proof.
induction e using (well_founded_induction
  (wf_inverse_image _ nat _ exprsize PeanoNat.Nat.lt_wf_0)).
intros.
destruct e; simpl; try rewrite <- H0; try reflexivity.
rewrite (List.rev_app_distr (rpn e2)).
rewrite (List.rev_app_distr (rpn e1)).
simpl.
(* Only eadd is left now *)
rewrite (List.app_assoc_reverse (List.rev (rpn e1)) (List.rev (rpn e2)) rest).
assert (Some (eval env e1 :: eval env e2 :: results)%list
    = rpn_eval' env (app (List.rev (rpn e1)) (app (List.rev (rpn e2)) rest))) as recstep.
  apply H; try apply H; simpl; try assumption; omega.
rewrite <- recstep.
reflexivity.
Qed.

Theorem correctness e env : Some (eval env e) = rpn_eval env (rpn e).
Proof.
assert (Some (eval env e :: nil)%list = rpn_eval' env (List.rev (rpn e))) as outsourcing.
rewrite <- (List.app_nil_r (List.rev (rpn e))).
apply correctness_partial_compilation.
simpl. reflexivity.
unfold rpn_eval.
rewrite <- outsourcing.
reflexivity.
Qed.