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 RPN :=
  | rpnnil : RPN
  | rpnlit : nat -> RPN -> RPN
  | rpnvar : nat -> RPN -> RPN
  | rpnadd : RPN -> RPN.

Fixpoint rpn_app (a : RPN) (b : RPN) :=
  match a with
  | rpnnil     => b
  | rpnlit n a => rpnlit n (rpn_app a b)
  | rpnvar n a => rpnvar n (rpn_app a b)
  | rpnadd a   => rpnadd (rpn_app a b)
  end.

Fixpoint rpn (e : Exp) :=
  match e with
  | elit n => rpnlit n rpnnil
  | evar n => rpnvar n rpnnil
  | eadd a b => rpnadd (rpn_app (rpn a) (rpn b))
  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 : RPN) :=
  match rpn with
  | rpnnil => Some nil
  | rpnlit n a => fmap (cons n) (rpn_eval' env a)
  | rpnvar n a => fmap (cons (env n)) (rpn_eval' env a)
  | 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 : RPN) :=
  match rpn_eval' env 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 exprsize_left  e1 e2 : exprsize e1 < exprsize (eadd e1 e2).
Proof. simpl. omega. Qed.
Lemma exprsize_right e1 e2 : exprsize e1 < exprsize (eadd e2 e1).
Proof. simpl. omega. Qed.

Lemma append_assoc a b c : rpn_app (rpn_app a b) c = rpn_app a (rpn_app b c).
Proof. induction a; simpl; try rewrite IHa; reflexivity. Qed.

Lemma free_append' e1 :
  forall env rest results,
  Some results = rpn_eval' env rest ->
  Some (eval env e1 :: results)%list = rpn_eval' env (rpn_app (rpn e1) rest).
Proof.
induction e1 using (well_founded_induction
  (wf_inverse_image _ nat _ exprsize PeanoNat.Nat.lt_wf_0)).
intros.
destruct e1; simpl; try rewrite <- H0; try reflexivity.
(* Only eadd is left now *)
rewrite append_assoc.
assert (Some (eval env e1_1 :: eval env e1_2 :: results)%list
    = rpn_eval' env (rpn_app (rpn e1_1) (rpn_app (rpn e1_2) rest))).
  apply H. apply exprsize_left.
  apply H. apply exprsize_right.
  assumption.
rewrite <- H1.
reflexivity.
Qed.

Lemma append_nil_at_end r : rpn_app r rpnnil = r.
Proof.
induction r; simpl; try rewrite IHr; reflexivity.
Qed.

Lemma correctness e env : Some (eval env e) = rpn_eval env (rpn e).
Proof.
assert (Some (eval env e :: nil)%list = rpn_eval' env (rpn e)).
assert (rpn_app (rpn e) rpnnil = rpn e).
apply append_nil_at_end.
rewrite <- H.
apply free_append'.
simpl. reflexivity.
unfold rpn_eval.
rewrite <- H.
reflexivity.
Qed.