Module Qcert.NNRSimp.Optim.NNRSimpUnflatten

Require Import String.
Require Import List.
Require Import ListSet.
Require Import Arith.
Require Import Bool.
Require Import Equivalence.
Require Import Morphisms.
Require Import Setoid.
Require Import Decidable.
Require Import EquivDec.
Require Import Program.
Require Import Utils.
Require Import DataRuntime.
Require Import NNRSimpRuntime.

Section NNRSimpUnflatten.
  Local Open Scope nnrs_imp.

  Context {fruntime:foreign_runtime}.

  Section unflatten.

    Section def.

      Fixpoint nnrs_imp_expr_unflatten_init (e:nnrs_imp_expr) : option nnrs_imp_expr
        := match e with
           | NNRSimpBinop OpBagUnion e₁ e₂ =>
             lift2 (NNRSimpBinop OpBagUnion)
                   (nnrs_imp_expr_unflatten_init e₁)
                   (nnrs_imp_expr_unflatten_init e₂)
           | NNRSimpUnop OpBag (NNRSimpConst (dcoll l)) =>
             Some (NNRSimpConst (dcoll l))
           | NNRSimpUnop OpBag (NNRSimpUnop OpBag e₁) =>
             Some (NNRSimpUnop OpBag e₁)
           | NNRSimpConst (dcoll nil) =>
             Some (NNRSimpConst (dcoll nil))
           | NNRSimpConst (dcoll ((dcoll l)::nil)) =>
             Some (NNRSimpConst (dcoll l))
           | _ => None
           end.

      Definition nnrs_imp_expr_unflatten_init_option (oe:option nnrs_imp_expr)
        : option (option nnrs_imp_expr)
        := match oe with
           | Some e => lift Some (nnrs_imp_expr_unflatten_init e)
           | None => Some None
           end.

      Fixpoint nnrs_imp_expr_unflatten_read (e:nnrs_imp_expr) (v:var) : option nnrs_imp_expr
        := match e with
           | NNRSimpVar x =>
             if x == v
             then None
             else Some (NNRSimpVar x)
           | NNRSimpUnop OpFlatten (NNRSimpVar x) =>
             if x == v
             then Some (NNRSimpVar x)
             else Some (NNRSimpUnop OpFlatten (NNRSimpVar x) )
           | NNRSimpGetConstant x =>
             Some (NNRSimpGetConstant x)
           | NNRSimpConst d =>
             Some (NNRSimpConst d)
           | NNRSimpUnop u e₁ =>
             lift (NNRSimpUnop u)
                  (nnrs_imp_expr_unflatten_read e₁ v)
           | NNRSimpBinop b e₁ e₂ =>
             lift2 (NNRSimpBinop b)
                   (nnrs_imp_expr_unflatten_read e₁ v)
                   (nnrs_imp_expr_unflatten_read e₂ v)
           | NNRSimpGroupBy g sl e₁ =>
             lift (NNRSimpGroupBy g sl) (nnrs_imp_expr_unflatten_read e₁ v)
           end.

      Definition nnrs_imp_expr_unflatten_read_option (oe:option nnrs_imp_expr) (v:var)
        : option (option nnrs_imp_expr)
        := match oe with
           | Some e => lift Some (nnrs_imp_expr_unflatten_read e v)
           | None => Some None
           end.

      Fixpoint nnrs_imp_expr_unflatten_write (e:nnrs_imp_expr) (v:var) : option nnrs_imp_expr
        := match e with
           | NNRSimpVar x =>
             if x == v
             then Some (NNRSimpVar x)
             else None
           | NNRSimpBinop OpBagUnion e₁ e₂ =>
             lift2 (NNRSimpBinop OpBagUnion)
                   (nnrs_imp_expr_unflatten_write e₁ v)
                   (nnrs_imp_expr_unflatten_write e₂ v)
           | NNRSimpUnop OpBag e₁ =>
             if in_dec string_dec v (nnrs_imp_expr_free_vars e₁)
             then None
             else Some e₁
           | NNRSimpConst (dcoll nil) =>
             Some (NNRSimpConst (dcoll nil))
           | NNRSimpConst (dcoll ((dcoll l)::nil)) =>
             Some (NNRSimpConst (dcoll l))
           | _ => None
           end.

      Definition lift3 {A B C D} (f:A->B->C->D) x y z
        := match x, y, z with
           | Some x, Some y, Some z => Some (f x y z)
           | _, _, _ => None
           end.
      
      Fixpoint nnrs_imp_stmt_unflatten (s:nnrs_imp_stmt) (v:var) : option nnrs_imp_stmt
        := match s with
           | NNRSimpSkip =>
             Some NNRSimpSkip
           | NNRSimpSeq s₁ s₂ =>
             lift2 NNRSimpSeq
                   (nnrs_imp_stmt_unflatten s₁ v)
                   (nnrs_imp_stmt_unflatten s₂ v)
           | NNRSimpAssign x e =>
             lift (NNRSimpAssign x)
                  (if x == v
                   then nnrs_imp_expr_unflatten_write e v
                   else nnrs_imp_expr_unflatten_read e v)
           | NNRSimpLet x oe s₁ =>
             lift2 (NNRSimpLet x)
                   (nnrs_imp_expr_unflatten_read_option oe v)
                   (if x == v then Some s₁ else nnrs_imp_stmt_unflatten s₁ v)
           | NNRSimpFor x e s₁ =>
             lift2 (NNRSimpFor x)
                   (nnrs_imp_expr_unflatten_read e v)
                   (if x == v then Some s₁ else nnrs_imp_stmt_unflatten s₁ v)
           | NNRSimpIf e s₁ s₂ =>
             lift3 NNRSimpIf
                   (nnrs_imp_expr_unflatten_read e v)
                   (nnrs_imp_stmt_unflatten s₁ v)
                   (nnrs_imp_stmt_unflatten s₂ v)
           | NNRSimpEither e x₁ s₁ x₂ s₂ =>
             lift3 (fun e s₁ s₂ => NNRSimpEither e x₁ s₁ x₂ s₂)
                   (nnrs_imp_expr_unflatten_read e v)
                   (if x₁ == v then Some s₁ else nnrs_imp_stmt_unflatten s₁ v)
                   (if x₂ == v then Some s₂ else nnrs_imp_stmt_unflatten s₂ v)
           end.

      Section safe.

        Fixpoint nnrs_imp_stmt_read_out (s:nnrs_imp_stmt) (v:var) : bool
          := match s with
             | NNRSimpSkip => false
             | NNRSimpSeq s₁ s₂ =>
               nnrs_imp_stmt_read_out s₁ v
               || nnrs_imp_stmt_read_out s₂ v
             | NNRSimpAssign x e =>
               (x <>b v) && nnrs_imp_expr_may_use e v
             | NNRSimpLet x oe s₁ =>
               match oe with
               | Some e => nnrs_imp_expr_may_use e v
               | None => false
               end
               || ((x <>b v) && nnrs_imp_stmt_read_out s₁ v)
             | NNRSimpFor x e s₁ =>
               nnrs_imp_expr_may_use e v
               || ((x <>b v) && nnrs_imp_stmt_read_out s₁ v)
             | NNRSimpIf e s₁ s₂ =>
               nnrs_imp_expr_may_use e v
               || nnrs_imp_stmt_read_out s₁ v
               || nnrs_imp_stmt_read_out s₂ v
             | NNRSimpEither e x₁ s₁ x₂ s₂ =>
               nnrs_imp_expr_may_use e v
               || ((x₁ <>b v) && nnrs_imp_stmt_read_out s₁ v)
               || ((x₂ <>b v) && nnrs_imp_stmt_read_out s₂ v)
             end.
        
        Definition nnrs_imp_stmt_unflatten_safe s v : option nnrs_imp_stmt
          := if nnrs_imp_stmt_read_out s v
             then nnrs_imp_stmt_unflatten s v
             else None.

      End safe.

    End def.

    Section ident.
      
      Lemma nnrs_imp_expr_unflatten_read_nfree_id {e v}:
        ~ In v (nnrs_imp_expr_free_vars e) ->
        nnrs_imp_expr_unflatten_read e v = Some e.

      Lemma nnrs_imp_stmt_unflatten_nfree_id {s v}:
        ~ In v (nnrs_imp_stmt_free_vars s) ->
        nnrs_imp_stmt_unflatten s v = Some s.

    End ident.
    Section eval.

      Definition flattens_to (d d':data)
        := match d, d' with
           | dcoll l, dcoll l' => oflatten l = Some l'
           | _, _ => False
           end.

      Definition flattens_to_op (d d':data)
        : flattens_to d d'
          <-> lift_oncoll (fun l => (lift dcoll (oflatten l))) d = Some d'.

      Lemma nnrs_imp_expr_unflatten_init_eval e e' h σc σ :
        nnrs_imp_expr_unflatten_init e = Some e' ->
        lift2P flattens_to
               (nnrs_imp_expr_eval h σc σ e)
               (nnrs_imp_expr_eval h σc σ e').

      Lemma nnrs_imp_expr_unflatten_read_eval e e' v h σc σ d d' :
        nnrs_imp_expr_unflatten_read e v = Some e' ->
        lookup equiv_dec σ v = Some d ->
        lift2P flattens_to d d' ->
        nnrs_imp_expr_eval h σc σ e =
        nnrs_imp_expr_eval h σc (update_first equiv_dec σ v d') e'.

    End eval.

    
    Section free_vars.
      
      Lemma nnrs_imp_expr_unflatten_init_free_vars e e' :
        nnrs_imp_expr_unflatten_init e = Some e' ->
        nnrs_imp_expr_free_vars e = nnrs_imp_expr_free_vars e'.
      
      Lemma nnrs_imp_expr_unflatten_read_free_vars e e' v :
        nnrs_imp_expr_unflatten_read e v = Some e' ->
        nnrs_imp_expr_free_vars e = nnrs_imp_expr_free_vars e'.

      Lemma nnrs_imp_expr_unflatten_write_free_vars e e' v :
        nnrs_imp_expr_unflatten_write e v = Some e' ->
        nnrs_imp_expr_free_vars e = nnrs_imp_expr_free_vars e'.

      Lemma nnrs_imp_stmt_unflatten_free_vars s s' v :
        nnrs_imp_stmt_unflatten s v = Some s' ->
        nnrs_imp_stmt_free_vars s = nnrs_imp_stmt_free_vars s'.

    End free_vars.

    Section usage.

      Lemma nnrs_imp_stmt_unflatten_var_usage {s s' v}:
        nnrs_imp_stmt_unflatten s v = Some s' ->
        forall v',
          nnrs_imp_stmt_var_usage s v' = nnrs_imp_stmt_var_usage s' v'.

      Lemma nnrs_imp_stmt_unflatten_var_usage_sub {s x s'} :
        nnrs_imp_stmt_unflatten s x = Some s' ->
        stmt_var_usage_sub s s'.

    End usage.
    
  End unflatten.
  
End NNRSimpUnflatten.