(**
  We define a pseudo SSA predicate.
  The formalization is similar to the renamedApart property in the LVC framework by Schneider, Smolka and Hack
  http://dblp.org/rec/conf/itp/SchneiderSH15
  Our predicate is not as fully fledged as theirs, but we especially borrow the idea of annotating
  the program with the predicate with the set of free and defined variables
**)
Require Import Coq.QArith.QArith Coq.QArith.Qreals Coq.Reals.Reals.
Require Import Coq.micromega.Psatz.
Require Import Daisy.Infra.RealRationalProps Daisy.Infra.Ltacs.
Require Export Daisy.Commands.

Lemma validVars_add V (e:exp V) Vs n:
  NatSet.Subset (usedVars e) Vs ->
  NatSet.Subset (usedVars e) (NatSet.add n Vs).
Proof.
  induction e; try auto.
  - intros valid_subset.
    hnf. intros a in_singleton.
    specialize (valid_subset a in_singleton).
    rewrite NatSet.add_spec; right; auto.
  - intros vars_binop.
    simpl in *.
    intros a in_empty.
    inversion in_empty.
  - simpl; intros in_vars.
    intros a in_union.
    rewrite NatSet.union_spec in in_union.
    destruct in_union.
    + apply IHe1; try auto.
      hnf; intros.
      apply in_vars.
      rewrite NatSet.union_spec; auto.
    + apply IHe2; try auto.
      hnf; intros.
      apply in_vars.
      rewrite NatSet.union_spec; auto.
Qed.

Lemma validVars_non_stuck (e:exp Q) inVars E:
 NatSet.Subset (usedVars e) inVars ->
  (forall v, NatSet.In v (usedVars e) ->
        exists r, (toREvalEnv E) v = Some (r, M0))%R ->
  exists vR, eval_exp (toREvalEnv E) (toREval (toRExp e)) vR M0.
Proof.
    revert inVars E; induction e;
    intros inVars E vars_valid fVars_live.
  - simpl in *.
    assert (NatSet.In n (NatSet.singleton n)) as in_n by (rewrite NatSet.singleton_spec; auto).
    specialize (fVars_live n in_n).
    destruct fVars_live as [vR E_def].
    exists vR; econstructor; eauto.
  - exists (perturb (Q2R v) 0); constructor.
    simpl (meps M0); rewrite Rabs_R0; rewrite Q2R0_is_0; lra.
  - assert (exists vR, eval_exp (toREvalEnv E) (toREval (toRExp e)) vR M0)
      as eval_e_def by (eapply IHe; eauto).
    destruct eval_e_def as [ve eval_e_def].
    case_eq u; intros; subst.
    + exists (evalUnop Neg ve); constructor; auto.
    + exists (perturb (evalUnop Inv ve) 0); constructor; auto.
      simpl (meps M0); rewrite Q2R0_is_0; rewrite Rabs_R0; lra.
  - repeat rewrite NatSet.subset_spec in *; simpl in *.
    assert (exists vR1, eval_exp (toREvalEnv E) (toREval (toRExp e1)) vR1 M0) as eval_e1_def.
    + eapply IHe1; eauto.
      * hnf; intros.
        apply vars_valid.
        rewrite NatSet.union_spec; auto.
      * intros.
        destruct (fVars_live v) as [vR E_def]; try eauto.
        apply NatSet.union_spec; auto.
    + assert (exists vR2, eval_exp (toREvalEnv E) (toREval (toRExp e2)) vR2 M0) as eval_e2_def.
      * eapply IHe2; eauto.
        { hnf; intros.
        apply vars_valid.
        rewrite NatSet.union_spec; auto. }
        { intros.
          destruct (fVars_live v) as [vR E_def]; try eauto.
          apply NatSet.union_spec; auto. }
      * destruct eval_e1_def as [vR1 eval_e1_def];
          destruct eval_e2_def as [vR2 eval_e2_def].
        exists (perturb (evalBinop b vR1 vR2) 0).
        replace M0 with (computeJoin M0 M0) by auto.
        constructor; auto.
        simpl meps; rewrite Q2R0_is_0. rewrite Rabs_R0; lra.
  - assert (exists vR, eval_exp (toREvalEnv E) (toREval (toRExp e))  vR M0) as eval_r_def by (eapply IHe; eauto).
    destruct eval_r_def as [vr eval_r_def].
    exists vr.
    simpl toREval.
    auto.
Qed.

Inductive ssaPrg (V:Type): (cmd V) -> (NatSet.t) -> (NatSet.t) -> Prop :=
 ssaLet m x e s inVars Vterm:
   NatSet.Subset (usedVars e) inVars->
   NatSet.mem x inVars = false ->
   ssaPrg s (NatSet.add x inVars) Vterm ->
   ssaPrg (Let m x e s) inVars Vterm
|ssaRet e inVars Vterm:
   NatSet.Subset (usedVars e) inVars ->
   NatSet.Equal inVars Vterm ->
   ssaPrg (Ret e) inVars Vterm.


Lemma ssa_subset_freeVars V (f:cmd V) inVars outVars:
  ssaPrg f inVars outVars ->
  NatSet.Subset (Commands.freeVars f) inVars.
Proof.
  intros ssa_f; induction ssa_f.
  - simpl in *. hnf; intros a in_fVars.
    rewrite NatSet.remove_spec, NatSet.union_spec in in_fVars.
    destruct in_fVars as [in_union not_eq].
    destruct in_union; try auto.
    specialize (IHssa_f a H1).
    rewrite NatSet.add_spec in IHssa_f.
    destruct IHssa_f; subst; try auto.
    exfalso; apply not_eq; auto.
  - hnf; intros.
    simpl in H1.
    apply H; auto.
Qed.


Lemma ssa_equal_set V (f:cmd V) inVars inVars' outVars:
  NatSet.Equal inVars' inVars ->
  ssaPrg f inVars outVars ->
  ssaPrg f inVars' outVars.
Proof.
  intros set_eq ssa_f.
  revert set_eq; revert inVars'.
  induction ssa_f.
  - constructor.
    + rewrite set_eq; auto.
    + case_eq (NatSet.mem x inVars'); intros case_mem; try auto.
      rewrite NatSet.mem_spec in case_mem.
      rewrite set_eq in case_mem.
      rewrite <- NatSet.mem_spec in case_mem.
      rewrite case_mem in *; congruence.
    + apply IHssa_f; auto.
      apply NatSetProps.Dec.F.add_m; auto.
  - constructor.
    + rewrite set_eq; auto.
    + rewrite set_eq; auto.
Qed.


Fixpoint validSSA (f:cmd Q) (inVars:NatSet.t) :=
  match f with
  |Let m x e g =>
    andb (andb (negb (NatSet.mem x inVars)) (NatSet.subset (usedVars e) inVars)) (validSSA g (NatSet.add x inVars))
  |Ret e => NatSet.subset (usedVars e) inVars
  end.

Lemma validSSA_sound f inVars:
  validSSA f inVars = true ->
  exists outVars, ssaPrg f inVars outVars.
Proof.
  revert inVars; induction f.
  - intros inVars validSSA_let.
    simpl in *.
    andb_to_prop validSSA_let.
    specialize (IHf (NatSet.add n inVars) R).
    destruct IHf as [outVars IHf].
    exists outVars.
    constructor; eauto.
    + rewrite <- NatSet.subset_spec; auto.
    + rewrite negb_true_iff in L0. auto.
  - intros inVars validSSA_ret.
    simpl in *.
    exists inVars.
    constructor; auto.
    + rewrite <- NatSet.subset_spec; auto.
    + hnf; intros; split; auto.
Qed.

Lemma ssa_shadowing_free m x y v v' e f inVars outVars E:
  ssaPrg (Let m x (toREval (toRExp e)) (toREvalCmd (toRCmd f))) inVars outVars ->
  NatSet.In y inVars ->
  eval_exp E (toREval (toRExp e)) v M0 ->
  forall E n, updEnv x m v (updEnv y m v' E) n = updEnv y m v' (updEnv x m v E) n.
Proof.
  intros ssa_let y_free eval_e.
  inversion ssa_let; subst.
  intros E' n; unfold updEnv.
  case_eq (n =? x).
  - intros n_eq_x.
    rewrite Nat.eqb_eq in n_eq_x.
    case_eq (n =? y); try auto.
    intros n_eq_y.
    rewrite Nat.eqb_eq in n_eq_y.
    rewrite n_eq_x in n_eq_y.
    rewrite <- n_eq_y in y_free.
    rewrite <- NatSet.mem_spec in y_free.
    rewrite y_free in H6. inversion H6.
  - intros n_neq_x.
    case_eq (n =? y); auto.
Qed.


Lemma shadowing_free_rewriting_exp e v E1 E2:
(*  (forall n, exists v m,
        E1 n = Some (v,m) ->
        exists m', (E2 n = Some (v,m') /\ (meps m) == (meps m'))) ->*)
  (forall n, E1 n = E2 n)->
  eval_exp E1 (toREval e) v M0 <->
  eval_exp E2 (toREval e) v M0.
Proof.
  revert v E1 E2.
  induction e; intros v' E1 E2 agree_on_vars.
  - split; intros eval_Var.
    + inversion eval_Var; subst.
      (*assert (m1 = M0) by (apply ifisMorePreciseM0; auto); subst.*)
      rewrite agree_on_vars in H2.
      eapply Var_load; eauto.
    + inversion eval_Var; subst.
      rewrite <- agree_on_vars in H2.
      eapply Var_load; eauto.
  - split; intros eval_Const; inversion eval_Const; subst; econstructor; auto.
  - split; intros eval_Unop; inversion eval_Unop; subst; econstructor; try auto.
    + erewrite IHe; eauto.
    + erewrite IHe; eauto.
    + erewrite <- IHe; eauto.
    + erewrite <- IHe; eauto.
  - split; intros eval_Binop; inversion eval_Binop; subst; econstructor; eauto;
    destruct m1; destruct m2; inversion H2;
      try (erewrite IHe1; eauto);
      try (erewrite IHe2; eauto); auto.
  - split; intros eval_Downcast; simpl; simpl in eval_Downcast; erewrite IHe; eauto.
Qed.

Lemma shadowing_free_rewriting_cmd f E1 E2 vR:
  (forall n, E1 n = E2 n) ->
  bstep (toREvalCmd f) E1 vR M0 <->
  bstep (toREvalCmd f) E2 vR M0.
Proof.
revert E1 E2 vR.
  induction f; intros E1 E2 vR agree_on_vars.
  - split; intros bstep_Let; inversion bstep_Let; subst.
    + erewrite shadowing_free_rewriting_exp in H6; auto.
      econstructor; eauto.
      rewrite <- IHf.
      apply H7.
      intros n'; unfold updEnv.
      case_eq (n' =? n); auto.
    + erewrite <- shadowing_free_rewriting_exp in H6; auto.
      econstructor; eauto.
      rewrite IHf.
      apply H7.
      intros n'; unfold updEnv.
      case_eq (n' =? n); auto.
  - split; intros bstep_Ret; inversion bstep_Ret; subst.
    + erewrite shadowing_free_rewriting_exp in H0; auto.
      constructor; auto.
    + erewrite <- shadowing_free_rewriting_exp in H0; auto.
      constructor; auto.
Qed.

Lemma dummy_bind_ok e v x v' inVars E:
  NatSet.Subset (usedVars e) inVars ->
  NatSet.mem x inVars = false ->
  eval_exp E (toREval (toRExp e)) v M0 ->
  eval_exp (updEnv x M0 v' E) (toREval (toRExp e)) v M0.
Proof.
  revert v x v' inVars E.
  induction e; intros v1 x v2 inVars E valid_vars x_not_free eval_e.
  - inversion eval_e; subst.
    assert ((updEnv x M0 v2 E) n = E n).
    + unfold updEnv.
      case_eq (n =? x); try auto.
      intros n_eq_x.
      rewrite Nat.eqb_eq in n_eq_x.
      subst; simpl in *.
      hnf in valid_vars.
      assert (NatSet.mem x (NatSet.singleton x) = true)
        as in_singleton by (rewrite NatSet.mem_spec, NatSet.singleton_spec; auto).
      rewrite NatSet.mem_spec in *.
      specialize (valid_vars x in_singleton).
      rewrite <- NatSet.mem_spec in valid_vars.
      rewrite valid_vars in *; congruence.
    + econstructor.
      eauto.
      rewrite H; auto. 
  - inversion eval_e; subst; constructor; auto.
  - inversion eval_e; subst; econstructor; eauto.
  - simpl in valid_vars.
    inversion eval_e; subst; econstructor; eauto;
    destruct m1; destruct m2; inversion H2;
      subst.
    + eapply IHe1; eauto.
      hnf; intros a in_e1.
      apply valid_vars;
        rewrite NatSet.union_spec; auto.
    + eapply IHe2; eauto.
      hnf; intros a in_e2.
      apply valid_vars;
        rewrite NatSet.union_spec; auto.
  - apply (IHe v1 x v2 inVars E); auto.
Qed.

Fixpoint exp_subst (e:exp Q) x e_new :=
  match e with
  |Var _ m v => if v =? x then e_new else Var Q m v
  |Unop op e1 => Unop op (exp_subst e1 x e_new)
  |Binop op e1 e2 => Binop op (exp_subst e1 x e_new) (exp_subst e2 x e_new)
  |Downcast m e1 => Downcast m (exp_subst e1 x e_new)
  | e => e
  end.

Lemma exp_subst_correct e e_sub E x v v_res:
  eval_exp E (toREval (toRExp e_sub)) v M0 ->
  eval_exp (updEnv x M0 v E) (toREval (toRExp e)) v_res M0 <->
  eval_exp E (toREval (toRExp (exp_subst e x e_sub))) v_res M0.
Proof.
  intros eval_e.
  revert v_res.
  induction e.
  - intros v_res; split; [intros eval_upd | intros eval_subst].
    + unfold updEnv in eval_upd. simpl in *.
      inversion eval_upd; subst.
      case_eq (n =? x); intros; try auto.
      * rewrite H in H2.
        inversion H2; subst; auto.
      * rewrite H in H2. 
        eapply Var_load; eauto.
    + simpl in eval_subst.
      case_eq (n =? x); intros n_eq_case;
        rewrite n_eq_case in eval_subst.
      * simpl.
        assert (updEnv x M0 v E n = Some (v_res, M0)).
        { unfold updEnv; rewrite n_eq_case.
          f_equal. assert (v = v_res) by (eapply meps_0_deterministic; eauto). rewrite H. auto. }
        { econstructor; eauto. } 
      * simpl. inversion eval_subst; subst.
        assert (E n = updEnv x M0 v E n).
        { unfold updEnv; rewrite n_eq_case; reflexivity. }
        { rewrite H in H2. eapply Var_load; eauto. } (*unfold updEnv in *. rewrite n_eq_case in *. unfold toREvalEnv. apply H4. }*)
  - intros v_res; split; [intros eval_upd | intros eval_subst].
    + inversion eval_upd; constructor; auto.
    + inversion eval_subst; constructor; auto.
  - split; [intros eval_upd | intros eval_subst].
    + inversion eval_upd; econstructor; auto;
        rewrite <- IHe; auto.
    + inversion eval_subst; constructor; try auto;
        rewrite IHe; auto.
  - intros v_res; split; [intros eval_upd | intros eval_subst].
    + inversion eval_upd; econstructor; auto;
      destruct m1; destruct m2; inversion H2.
      * rewrite <- IHe1; auto.
      * rewrite <- IHe2; auto.
    + inversion eval_subst; econstructor; auto;
        destruct m1; destruct m2; inversion H2.
      * rewrite IHe1; auto.
      * rewrite IHe2; auto.
  - split; [intros eval_upd | intros eval_subst].
    + rewrite <- IHe; auto.
    + rewrite IHe; auto.
Qed.

Fixpoint map_subst (f:cmd Q) x e :=
  match f with
  |Let m y e_y g => Let m y (exp_subst e_y x e) (map_subst g x e)
  |Ret e_r => Ret (exp_subst e_r x e)
  end.


(* Lemma updEnv_comp_toREval (n:nat) (v:R): *)
(*   forall E var, updEnv n M0 v (toREvalEnv E) var = toREvalEnv (updEnv n M0 v E) var. *)
(* Proof. *)
(*   intros. *)
(*   unfold updEnv; unfold toREvalEnv. *)
(*   case_eq (var =? n); intros; auto. *)
(* Qed. *)

(* Lemma bli (e:exp Q) (n:nat) (v vR:R) (E:env): *)
(*   eval_exp (toREvalEnv (updEnv n M0 v E)) (toREval (toRExp e)) vR M0 <-> *)
(*   eval_exp (updEnv n M0 v (toREvalEnv E)) (toREval (toRExp e)) vR M0. *)
(* Proof. *)
(*   revert e n v vR E. *)
(*   induction e; split; intros. *)
(*   - inversion H; subst; econstructor; eauto. *)
(*     rewrite updEnv_comp_toREval; auto. *)
(*   - inversion H; subst; econstructor; eauto. *)
(*     rewrite <- updEnv_comp_toREval; auto. *)
(*   - inversion H; subst; econstructor; eauto. *)
(*   - inversion H; subst; econstructor; eauto. *)
(*   - inversion H; subst; econstructor; eauto; apply IHe; auto. *)
(*   - inversion H; subst; econstructor; eauto; apply IHe; auto. *)
(*   - inversion H; subst; econstructor; eauto; *)
(*       assert (m1 = M0) by (apply (ifM0isJoin_l M0 m1 m2); auto); *)
(*       assert (m2 = M0) by (apply (ifM0isJoin_r M0 m1 m2); auto); subst. *)
(*     apply IHe1; auto. *)
(*     apply IHe2; auto. *)
(*   - inversion H; subst; econstructor; eauto; *)
(*       assert (m1 = M0) by (apply (ifM0isJoin_l M0 m1 m2); auto); *)
(*       assert (m2 = M0) by (apply (ifM0isJoin_r M0 m1 m2); auto); subst. *)
(*     apply IHe1; auto. *)
(*     apply IHe2; auto. *)
(*   - apply IHe; auto. *)
(*   - apply IHe; auto. *)
(* Qed. *)

(* Lemma bla (c: cmd Q) (n:nat) (v vR:R) (E:env): *)
(*   bstep (toREvalCmd (toRCmd c)) (updEnv n M0 v (toREvalEnv E)) vR M0 <-> *)
(*   bstep (toREvalCmd (toRCmd c)) (toREvalEnv (updEnv n M0 v E)) vR M0. *)
(* Proof. *)
(*   revert c n v vR E. *)
(*   induction c; split; intros. *)
(*   - inversion H; subst. *)
(*     econstructor; eauto. *)
(*     apply bli; eauto. *)
(*     apply IHc; auto. *)
(* Admitted. *)
(*     remember (updEnv n0 M0 v (toREvalEnv E)) as E'. *)

(*     replace E' with (toREvalEnv E') in H9. *)
(*     + apply IHc in H9. *)
(*       rewrite HeqE' in H9. *)
(*       replace E with (toREvalEnv E) by admit. *)
(*       apply H9. *)
(* Admitted. *)


Lemma stepwise_substitution x e v f E vR inVars outVars:
  ssaPrg (toREvalCmd (toRCmd f)) inVars outVars ->
  NatSet.In x inVars ->
  NatSet.Subset (usedVars e) inVars ->
  eval_exp E (toREval (toRExp e)) v M0 ->
  bstep (toREvalCmd (toRCmd f)) (updEnv x M0 v E) vR M0 <->
  bstep (toREvalCmd (toRCmd (map_subst f x e))) E vR M0.
Proof.
  revert E e x vR inVars outVars.
  induction f; intros E e0 x vR inVars outVars ssa_f x_free valid_vars_e eval_e.
  - split; [ intros bstep_let | intros bstep_subst].
    + inversion bstep_let; subst.
      inversion ssa_f; subst.
      assert (forall E var, updEnv n M0 v0 (updEnv x M0 v E) var = updEnv x M0 v (updEnv n M0 v0 E) var).
      * eapply ssa_shadowing_free.
        apply ssa_f.
        apply x_free.
        apply H6.
      * erewrite (shadowing_free_rewriting_cmd _ _ _ _) in H7; try eauto.
        simpl in *.
        econstructor; eauto.
        { rewrite <- exp_subst_correct; eauto. }
        { rewrite <- IHf; eauto.
          rewrite NatSet.add_spec; right; auto.
          apply validVars_add; auto.
          eapply dummy_bind_ok; eauto. }
    + inversion bstep_subst; subst.
      simpl in *.
      inversion ssa_f; subst.
      econstructor; try auto.
      rewrite exp_subst_correct; eauto.
      rewrite <- IHf in H7; eauto.
      * rewrite <- shadowing_free_rewriting_cmd in H7; eauto.
        eapply ssa_shadowing_free; eauto.
        rewrite <- exp_subst_correct in H6; eauto.
      * rewrite NatSet.add_spec; auto.
      * apply validVars_add; auto.
      * eapply dummy_bind_ok; eauto.
  - split; [intros bstep_let | intros bstep_subst].
    + inversion bstep_let; subst.
      simpl in *.
      rewrite exp_subst_correct in H0; eauto.
      constructor. assumption.
    + inversion bstep_subst; subst.
      simpl in *.
      rewrite <- exp_subst_correct in H0; eauto.
      econstructor; eauto.
Qed.

Fixpoint let_subst (f:cmd Q) :=
  match f with
  | Let m x e1 g =>
    match (let_subst g) with
    |Some e' =>  Some (exp_subst e' x e1)
    |None => None
    end
  | Ret e1 =>  Some e1
  end.

Theorem let_free_form f E vR inVars outVars e:
  ssaPrg f inVars outVars ->
  bstep (toREvalCmd (toRCmd f)) E vR M0 ->
  let_subst f = Some e ->
  bstep (toREvalCmd (toRCmd (Ret e))) E vR M0.
Proof.
    revert E vR inVars outVars e;
    induction f;
    intros E vR inVars outVars e_subst ssa_f bstep_f subst_step.
  - simpl.
    inversion bstep_f; subst.
    inversion ssa_f; subst.
    simpl in subst_step.
    case_eq (let_subst f).
    + intros f_subst subst_f_eq.
      specialize (IHf (updEnv n M0 v E) vR (NatSet.add n inVars) outVars f_subst H9 H7 subst_f_eq).
      rewrite subst_f_eq in subst_step.
      inversion IHf; subst.
      constructor.
      inversion subst_step.
      subst.
      rewrite <- exp_subst_correct; eauto.
    + intros subst_eq; rewrite subst_eq in subst_step; inversion subst_step.
  - inversion bstep_f; subst.
    constructor.
    simpl in *.
    inversion subst_step; subst.
    assumption.
Qed.