Require Import Omega Wellfounded Coq.Lists.List.
Import ListNotations.

Inductive Exp (v : Type) :=
  | elit : nat -> Exp v
  | evar : v   -> Exp v
  | eadd : Exp v -> Exp v -> Exp v.

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

Inductive RPNOp (v : Type) :=
  | rpnlit : nat -> RPNOp v
  | rpnvar : v   -> RPNOp v
  | rpnadd : RPNOp v.

Fixpoint rpn (v : Type) (e : Exp v) : list (RPNOp v) :=
  match e with
  | elit _ n => [rpnlit v n]
  | evar _ n => [rpnvar v n]
  | eadd _ a b => rpn v b ++ rpn v a ++ [rpnadd v]
  end.

Definition fmap (a : Type) (b : Type) (f : a -> b) (o : option a) : option b :=
  match o with
  | Some x => Some (f x)
  | None   => None
  end.
Notation "x <$> y" := (fmap _ _ x y) (at level 65, left associativity).

Definition bind (a : Type) (b : Type) (f : a -> option b) (o : option a) : option b :=
  match o with
  | Some x => f x
  | None   => None
  end.
Notation "x >>= f" := (bind _ _ f x) (at level 55, left associativity).

Fixpoint pn_eval (v : Type) (env : v -> nat) (rpn : list (RPNOp v)) : option (list nat) :=
  match rpn with
  | []              => Some []
  | rpnlit _ n :: a => cons n <$> pn_eval v env a
  | rpnvar _ n :: a => cons (env n) <$> pn_eval v env a
  | rpnadd _   :: a => pn_eval v env a >>= fun stack => match stack with
    | x :: y :: stack => Some (x+y :: stack)
    | _               => None
    end
  end.

Function rpn_eval (v : Type) (env : v -> nat) (rpn : list (RPNOp v)) : option nat :=
  match pn_eval v env (rev rpn) with
  | Some [x] => Some x
  | _        => None
  end.

Lemma correctness_partial_compilation (v : Type) (e : Exp v) :
  forall (env : v -> nat) (rest : list (RPNOp v)) (results : list nat),
  Some results = pn_eval v env rest ->
  Some (eval v env e :: results) = pn_eval v env (rev (rpn v e) ++ rest).
Proof.

Fixpoint exprsize (v : Type) (e : Exp v) : nat :=
  match e with
  | elit _ _ => 1
  | evar _ _ => 1
  | eadd _ a b => 1 + exprsize v a + exprsize v b
  end.
induction e using (well_founded_induction
  (wf_inverse_image _ nat _ (exprsize v) PeanoNat.Nat.lt_wf_0)).

intros.
destruct e; simpl; try rewrite <- H0; try reflexivity.
rewrite (rev_app_distr (rpn v e2)).
rewrite (rev_app_distr (rpn v e1)).
simpl.
rewrite (app_assoc_reverse (rev (rpn v e1)) (rev (rpn v e2)) rest).
assert (Some (eval v env e1 :: eval v env e2 :: results)
    = pn_eval v env (rev (rpn v e1) ++ rev (rpn v e2) ++ rest)) as recstep.
  apply H; try apply H; simpl; try assumption; omega.
rewrite <- recstep.
reflexivity.
Qed.

Theorem correctness (v : Type) (e : Exp v) (env : v -> nat) :
  Some (eval v env e) = rpn_eval v env (rpn v e).
Proof.
assert (Some [eval v env e] = pn_eval v env (rev (rpn v e))) as outsourcing.
rewrite <- (app_nil_r (rev (rpn v e))).
apply correctness_partial_compilation.
simpl. reflexivity.
unfold rpn_eval.
rewrite <- outsourcing.
reflexivity.
Qed.