Module UnLoad

Require Import
  Coqlib Utf8 Util
  AST Integers Floats
  Csharpminor.

Require Import
  AssocList
  Psatz.

Section LEAST_LARGER.

Local Open Scope positive_scope.

Definition least_larger (l: list ident) : ident :=
  match l with
  | nil => 1
  | i :: l' => List.fold_left (λ a j, if a <=? j then Pos.succ j else a) l' (Pos.succ i)
  end.

Lemma least_larger_larger l :
  ∀ x, In x l → x < least_larger l.

Lemma least_larger_least l :
  ∀ x, (∀ y, In y l → y < x) →
       least_larger l <= x.

End LEAST_LARGER.

Module ST.
  Definition t S A : Type := S → A * S.

  Instance mon {S} : monad (t S) :=
    { bind A B m f := λ s, let '(a, s') := m(s) in f a s';
      ret A a := λ s, (a, s) }.

  Definition get {S} : t S S := λ s, (s, s).
  Definition put {S} s : t S _ := λ _, (tt, s).

  Definition upd {S} (f: S → S) : t S _ := λ s, (tt, f s).

  Definition left {S₁ S₂ A} (m: t S₁ A) : t (S₁ * S₂) A := λ s, let '(a, s') := m (fst s) in (a, (s', snd s)).
  Definition right {S₁ S₂ A} (m: t S₂ A) : t (S₁ * S₂) A := λ s, let '(a, s') := m (snd s) in (a, (fst s, s')).

End ST.

Section EXPR_EQ.

Definition constant_beq (c d: constant) : bool :=
  match c, d with
  | Ointconst i, Ointconst j => Int.eq i j
  | Ofloatconst f, Ofloatconst g => Float.eq_dec f g
  | Osingleconst f, Osingleconst g => Float32.eq_dec f g
  | Olongconst i, Olongconst j => Int64.eq i j
  | _, _ => false
  end.

Lemma constant_beq_eq c d :
  constant_beq c d = true ↔ c = d.

Instance CmpEqDec : EqDec comparison.

Instance UopEqDec : EqDec unary_operation.

Instance BopEqDec : EqDec binary_operation.

Fixpoint expr_beq (x y: expr) {struct x} : bool :=
  match x, y with
  | Evar i, Evar j
  | Eaddrof i, Eaddrof j => Pos.eqb i j
  | Econst c, Econst d => constant_beq c d
  | Eunop o a, Eunop p b =>
    eq_dec o p && expr_beq a b
  | Ebinop o a b, Ebinop p c d =>
    eq_dec o p && expr_beq a c && expr_beq b d
  | Eload κ a, Eload κ' b =>
    eq_dec κ κ' && expr_beq a b
  | _, _ => false
  end.

Lemma expr_beq_eq:
  ∀ x y,
  expr_beq x y = true ↔ x = y.

End EXPR_EQ.

Instance ExprEqDec : EqDec expr := eq_dec_of_beq expr_beq expr_beq_eq.

Section EXTRACTION.

Fixpoint unload (e: expr) : ST.t (list expr) (list (expr * ident) → expr) :=
  match e with
  | Eunop o e₁ => do x₁ <- unload e₁; ret (λ l, Eunop o (x₁ l))
  | Ebinop o e₁ e₂ => do x₁ <- unload e₁; do x₂ <- unload e₂; ret (λ l, Ebinop o (x₁ l) (x₂ l))
  | Eload _ _ =>
    do l <- ST.get;
    do _ <- if List.in_dec eq_dec e l
           then ret tt
           else ST.put (e::l);
    ret (λ l, Evar (match assoc e l with Some t => t | None => xH end))
  | Evar _
  | Eaddrof _
  | Econst _
    => ret (λ _, e)
  end.

End EXTRACTION.

Section TRANSFORM.

Variable unload : expr → (list (expr * ident) → expr) * list expr.

Definition next_temp : ST.t ident ident :=
  do i <- ST.get;
  do _ <- ST.put (Pos.succ i);
  ret i.

Definition next_assoc (e: expr) : ST.t (list (expr * ident) * ident) ident :=
  do t <- ST.right next_temp;
  do _ <- ST.left (ST.upd (λ s, (e, t) :: s));
  ret t.

Fixpoint expand' (base: list (expr * ident) → stmt) (l: list expr) : ST.t (list (expr * ident) * ident) stmt :=
  match l with
  | nil => do x <- ST.left ST.get; ret (base x)
  | e :: l' => do t <- next_assoc e; do s <- expand' base l'; ret (Sseq (Sset t e) s)
  end.

Definition expand (base: list (expr * ident) → stmt) (l: list expr) :
ST.t ident stmt :=
  λ i, let '(s, (_, i')) := expand' base l (nil, i) in (s, i').

Fixpoint transform_stmt (s: stmt) : ST.t ident stmt :=
  match s with
  | Sseq s₁ s₂ =>
    do s₁' <- transform_stmt s₁;
    do s₂' <- transform_stmt s₂;
    ret (Sseq s₁' s₂')
  | Sifthenelse e s₁ s₂ =>
    do s₁' <- transform_stmt s₁;
    do s₂' <- transform_stmt s₂;
    ret (Sifthenelse e s₁' s₂')
  | Sloop s₁ =>
    do s₁' <- transform_stmt s₁;
    ret (Sloop s₁')
  | Sblock s₁ =>
    do s₁' <- transform_stmt s₁;
    ret (Sblock s₁')
  | Sreturn (Some e) =>
    let '(e', l) := unload e in
    expand (λ q, Sreturn (Some (e' q))) l
  | Sset x e =>
    let '(e', l) := unload e in
    expand (λ q, Sset x (e' q)) l
  | Sstore κ a e =>
    let '(e', l) := unload e in
    expand (λ q, Sstore κ a (e' q)) l
  | Sskip
  | Sexit _
  | Slabel _ _
  | Sgoto _
  | Sreturn None
  | Sswitch _ _ _
  | Scall _ _ _ _
  | Sbuiltin _ _ _
    => ret s
  end.

End TRANSFORM.

Definition all_ident acc (first next: ident) : list ident :=
  let num := (Pos.succ next - first)%positive in
  Pos.peano_rect (λ _, list ident)
                 acc (λ p l, Pos.pred (p + first) :: l) num.


Definition doit (p: program) : program :=
  transform_program
    (λ f,
     match f with
     | Internal f =>
       let t₀ := Pos.max (least_larger f.(fn_temps)) (least_larger f.(fn_params)) in
       let '(body, t) := transform_stmt (λ e, unload e nil) f.(fn_body) t₀ in
       (Internal (mkfunction f.(fn_sig) f.(fn_params) f.(fn_vars) (all_ident f.(fn_temps) t₀ t) body))
     | External _ => f
     end)
    p.