Require Import Omega.
Require Import Wellfounded.

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

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

Inductive RPN :=
  | rpnnil : RPN
  | rpnlit : nat -> RPN -> RPN
  | rpnvar : nat -> RPN -> RPN
  | rpnadd : RPN -> RPN.

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

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

Fixpoint rpn_eval' (env : nat -> nat) (rpn : RPN) :=
  match rpn with
  | rpnnil => Some nil
  | rpnlit n a => match rpn_eval' env a with
    | Some stack => Some (cons n stack) (* TODO: use functor *)
    | None       => None
    end
  | rpnvar n a => match rpn_eval' env a with
    | Some stack => Some (cons (env n) stack)
    | None       => None
    end
  | 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 : append_rpn (append_rpn a b) c = append_rpn a (append_rpn b c).
Proof.
induction a; simpl; try rewrite IHa; reflexivity.
Qed.

Lemma free_append' e1 :
  forall env rest rest_results,
  Some rest_results = rpn_eval' env rest ->
  Some (eval env e1 :: rest_results)%list = rpn_eval' env (append_rpn (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 :: rest_results)%list = rpn_eval' env (append_rpn (rpn e1_1) (append_rpn (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 : append_rpn 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 (append_rpn (rpn e) rpnnil = rpn e).
apply append_nil_at_end.
rewrite <- H.
apply free_append'.
simpl. reflexivity.
unfold rpn_eval.
rewrite <- H.
reflexivity.
Qed.