Module DNNRCBase

Section DNNRCBase.

  Require Import String.
  Require Import List.
  Require Import Arith.
  Require Import EquivDec.
  Require Import Morphisms.

  Require Import Utils BasicRuntime.
  Require Import DData.

  Context {fruntime:foreign_runtime}.
  
  
  Definition var := string.

  Section plug.
    Context {plug_type:Set}.

    Definition coll_bindings := list (string * (list data)).

    Definition bindings_of_coll_bindings (cb:coll_bindings) : bindings :=
      map (fun xy => (fst xy, dcoll (snd xy))) cb.
    
    Class AlgPlug :=
      mkAlgPlug {
          plug_eval : brand_relation_t -> coll_bindings -> plug_type -> option data
          ; plug_normalized :
            forall (h:brand_relation_t) (op:plug_type), forall (constant_env:coll_bindings) (o:data),
                Forall (fun x => data_normalized h (snd x)) (bindings_of_coll_bindings constant_env) ->
                plug_eval h constant_env op = Some o -> data_normalized h o;
        }.

  End plug.
  Global Arguments AlgPlug : clear implicits.
  
  Section GenDNNRCBase.
    Context {A plug_type:Set}.

    Unset Elimination Schemes.

    Inductive dnnrc : Set :=
    | DNNRCVar : A -> var -> dnnrc
    | DNNRCConst : A -> data -> dnnrc
    | DNNRCBinop : A -> binOp -> dnnrc -> dnnrc -> dnnrc
    | DNNRCUnop : A -> unaryOp -> dnnrc -> dnnrc
    | DNNRCLet : A -> var -> dnnrc -> dnnrc -> dnnrc
    | DNNRCFor : A -> var -> dnnrc -> dnnrc -> dnnrc
    | DNNRCIf : A -> dnnrc -> dnnrc -> dnnrc -> dnnrc
    | DNNRCEither : A -> dnnrc -> var -> dnnrc -> var -> dnnrc -> dnnrc
    | DNNRCGroupBy : A -> string -> list string -> dnnrc -> dnnrc
    | DNNRCCollect : A -> dnnrc -> dnnrc
    | DNNRCDispatch : A -> dnnrc -> dnnrc
    | DNNRCAlg : A -> plug_type -> list (string * dnnrc) -> dnnrc.

    Set Elimination Schemes.


    Definition dnnrc_rect (P : dnnrc -> Type)
               (fdvar : forall a, forall v, P (DNNRCVar a v))
               (fdconst : forall a, forall d : data, P (DNNRCConst a d))
               (fdbinop : forall a, forall b, forall n1 n2 : dnnrc, P n1 -> P n2 -> P (DNNRCBinop a b n1 n2))
               (fdunop : forall a, forall u, forall n : dnnrc, P n -> P (DNNRCUnop a u n))
               (fdlet : forall a, forall v, forall n1 n2 : dnnrc, P n1 -> P n2 -> P (DNNRCLet a v n1 n2))
               (fdfor : forall a, forall v, forall n1 n2 : dnnrc, P n1 -> P n2 -> P (DNNRCFor a v n1 n2))
               (fdif : forall a, forall n1 n2 n3 : dnnrc, P n1 -> P n2 -> P n3 -> P (DNNRCIf a n1 n2 n3))
               (fdeither : forall a, forall n0 n1 n2, forall v1 v2,
                       P n0 -> P n1 -> P n2 -> P (DNNRCEither a n0 v1 n1 v2 n2))
               (fdgroupby : forall a, forall g, forall sl, forall n : dnnrc, P n -> P (DNNRCGroupBy a g sl n))
               (fdcollect : forall a, forall n : dnnrc, P n -> P (DNNRCCollect a n))
               (fddispatch : forall a, forall n : dnnrc, P n -> P (DNNRCDispatch a n))
               (fdalg : forall a, forall op:plug_type, forall r : list (string*dnnrc), Forallt (fun n => P (snd n)) r -> P (DNNRCAlg a op r))
      :=
        fix F (n : dnnrc) : P n :=
          match n as n0 return (P n0) with
          | DNNRCVar a v => fdvar a v
          | DNNRCConst a d => fdconst a d
          | DNNRCBinop a b n1 n2 => fdbinop a b n1 n2 (F n1) (F n2)
          | DNNRCUnop a u n => fdunop a u n (F n)
          | DNNRCLet a v n1 n2 => fdlet a v n1 n2 (F n1) (F n2)
          | DNNRCFor a v n1 n2 => fdfor a v n1 n2 (F n1) (F n2)
          | DNNRCIf a n1 n2 n3 => fdif a n1 n2 n3 (F n1) (F n2) (F n3)
          | DNNRCEither a n0 v1 n1 v2 n2 => fdeither a n0 n1 n2 v1 v2 (F n0) (F n1) (F n2)
          | DNNRCGroupBy a g sl n => fdgroupby a g sl n (F n)
          | DNNRCCollect a n => fdcollect a n (F n)
          | DNNRCDispatch a n => fddispatch a n (F n)
          | DNNRCAlg a op r =>
            fdalg a op r ((fix F3 (r : list (string * dnnrc)) : Forallt (fun n => P (snd n)) r :=
                             match r as c0 with
                             | nil => Forallt_nil _
                             | cons n c0 => @Forallt_cons _ _ _ _ (F (snd n)) (F3 c0)
                             end) r)
          end.

    Definition dnnrc_ind (P : dnnrc -> Prop)
               (fdvar : forall a, forall v, P (DNNRCVar a v))
               (fdconst : forall a, forall d : data, P (DNNRCConst a d))
               (fdbinop : forall a, forall b, forall n1 n2 : dnnrc, P n1 -> P n2 -> P (DNNRCBinop a b n1 n2))
               (fdunop : forall a, forall u, forall n : dnnrc, P n -> P (DNNRCUnop a u n))
               (fdlet : forall a, forall v, forall n1 n2 : dnnrc, P n1 -> P n2 -> P (DNNRCLet a v n1 n2))
               (fdfor : forall a, forall v, forall n1 n2 : dnnrc, P n1 -> P n2 -> P (DNNRCFor a v n1 n2))
               (fdif : forall a, forall n1 n2 n3 : dnnrc, P n1 -> P n2 -> P n3 -> P (DNNRCIf a n1 n2 n3))
               (fdeither : forall a, forall n0 n1 n2, forall v1 v2,
                       P n0 -> P n1 -> P n2 -> P (DNNRCEither a n0 v1 n1 v2 n2))
               (fdgroupby : forall a, forall g, forall sl, forall n : dnnrc, P n -> P (DNNRCGroupBy a g sl n))
               (fdcollect : forall a, forall n : dnnrc, P n -> P (DNNRCCollect a n))
               (fddispatch : forall a, forall n : dnnrc, P n -> P (DNNRCDispatch a n))
               (fdalg : forall a, forall op:plug_type, forall r : list (string*dnnrc), Forall (fun n => P (snd n)) r -> P (DNNRCAlg a op r))
      :=
        fix F (n : dnnrc) : P n :=
          match n as n0 return (P n0) with
          | DNNRCVar a v => fdvar a v
          | DNNRCConst a d => fdconst a d
          | DNNRCBinop a b n1 n2 => fdbinop a b n1 n2 (F n1) (F n2)
          | DNNRCUnop a u n => fdunop a u n (F n)
          | DNNRCLet a v n1 n2 => fdlet a v n1 n2 (F n1) (F n2)
          | DNNRCFor a v n1 n2 => fdfor a v n1 n2 (F n1) (F n2)
          | DNNRCIf a n1 n2 n3 => fdif a n1 n2 n3 (F n1) (F n2) (F n3)
          | DNNRCEither a n0 v1 n1 v2 n2 => fdeither a n0 n1 n2 v1 v2 (F n0) (F n1) (F n2)
          | DNNRCGroupBy a g sl n => fdgroupby a g sl n (F n)
          | DNNRCCollect a n => fdcollect a n (F n)
          | DNNRCDispatch a n => fddispatch a n (F n)
          | DNNRCAlg a op r =>
            fdalg a op r ((fix F3 (r : list (string*dnnrc)) : Forall (fun n => P (snd n)) r :=
                             match r as c0 with
                             | nil => Forall_nil _
                             | cons n c0 => @Forall_cons _ _ _ _ (F (snd n)) (F3 c0)
                             end) r)
          end.

    Definition dnnrc_rec (P:dnnrc->Set) := @dnnrc_rect P.

    Lemma dnnrcInd2 (P : dnnrc -> Prop)
          (fdvar : forall a, forall v, P (DNNRCVar a v))
          (fdconst : forall a, forall d : data, P (DNNRCConst a d))
          (fdbinop : forall a, forall b, forall n1 n2 : dnnrc, P (DNNRCBinop a b n1 n2))
          (fdunop : forall a, forall u, forall n : dnnrc, P (DNNRCUnop a u n))
          (fdlet : forall a, forall v, forall n1 n2 : dnnrc, P (DNNRCLet a v n1 n2))
          (fdfor : forall a, forall v, forall n1 n2 : dnnrc, P (DNNRCFor a v n1 n2))
          (fdif : forall a, forall n1 n2 n3 : dnnrc, P (DNNRCIf a n1 n2 n3))
          (fdeither : forall a, forall n0 n1 n2, forall v1 v2,
                  P (DNNRCEither a n0 v1 n1 v2 n2))
          (fdgroupby : forall a, forall g, forall sl, forall n : dnnrc, P n -> P (DNNRCGroupBy a g sl n))
          (fdcollect : forall a, forall n : dnnrc, P (DNNRCCollect a n))
          (fddispatch : forall a, forall n : dnnrc, P (DNNRCDispatch a n))
          (fdalg : forall a, forall op:plug_type, forall r : list (string*dnnrc), Forall (fun n => P (snd n)) r -> P (DNNRCAlg a op r))
: forall n, P n.

    Definition dnnrc_annotation_get (d:dnnrc) : A
      := match d with
         | DNNRCVar a _ => a
         | DNNRCConst a _ => a
         | DNNRCBinop a _ _ _ => a
         | DNNRCUnop a _ _ => a
         | DNNRCLet a _ _ _ => a
         | DNNRCFor a _ _ _ => a
         | DNNRCIf a _ _ _ => a
         | DNNRCEither a _ _ _ _ _ => a
         | DNNRCGroupBy a _ _ _ => a
         | DNNRCCollect a _ => a
         | DNNRCDispatch a _ => a
         | DNNRCAlg a _ _ => a
         end.

    Definition dnnrc_annotation_update (f:A->A) (d:dnnrc) : dnnrc
      := match d with
         | DNNRCVar a v => DNNRCVar (f a) v
         | DNNRCConst a c => DNNRCConst (f a) c
         | DNNRCBinop a b d₁ d₂ => DNNRCBinop (f a) b d₁ d₂
         | DNNRCUnop a u d₁ => DNNRCUnop (f a) u d₁
         | DNNRCLet a x d₁ d₂ => DNNRCLet (f a) x d₁ d₂
         | DNNRCFor a x d₁ d₂ => DNNRCFor (f a) x d₁ d₂
         | DNNRCIf a d₀ d₁ d₂ => DNNRCIf (f a) d₀ d₁ d₂
         | DNNRCEither a d₀ x₁ d₁ x₂ d₂ => DNNRCEither (f a) d₀ x₁ d₁ x₂ d₂
         | DNNRCGroupBy a g sl d₁ => DNNRCGroupBy (f a) g sl d₁
         | DNNRCCollect a d₀ => DNNRCCollect (f a) d₀
         | DNNRCDispatch a d₀ => DNNRCDispatch (f a) d₀
         | DNNRCAlg a p args => DNNRCAlg (f a) p args
         end .

    Context (h:brand_relation_t).
    Fixpoint dnnrc_eval `{plug: AlgPlug plug_type} (env:dbindings) (e:dnnrc) : option ddata :=
      match e with
      | DNNRCVar _ x =>
        lookup equiv_dec env x
      | DNNRCConst _ d =>
        Some (Dlocal (normalize_data h d))
      | DNNRCBinop _ bop e1 e2 =>
        olift2 (fun d1 d2 => lift Dlocal (fun_of_binop h bop d1 d2))
               (olift checkLocal (dnnrc_eval env e1)) (olift checkLocal (dnnrc_eval env e2))
      | DNNRCUnop _ uop e1 =>
        olift (fun d1 => lift Dlocal (fun_of_unaryop h uop d1)) (olift checkLocal (dnnrc_eval env e1))
      | DNNRCLet _ x e1 e2 =>
        match dnnrc_eval env e1 with
        | Some d => dnnrc_eval ((x,d)::env) e2
        | _ => None
        end
      | DNNRCFor _ x e1 e2 =>
        match dnnrc_eval env e1 with
        | Some (Ddistr c1) =>
          let inner_eval d1 :=
              let env' := (x,Dlocal d1) :: env in olift checkLocal (dnnrc_eval env' e2)
          in
          lift (fun coll => Ddistr coll) (rmap inner_eval c1)
        | Some (Dlocal (dcoll c1)) =>
          let inner_eval d1 :=
              let env' := (x,Dlocal d1) :: env in olift checkLocal (dnnrc_eval env' e2)
          in
          lift (fun coll => Dlocal (dcoll coll)) (rmap inner_eval c1)
        | Some (Dlocal _) => None
        | None => None
        end
      | DNNRCIf _ e1 e2 e3 =>
        let aux_if d :=
            match d with
            | dbool b =>
              if b then dnnrc_eval env e2 else dnnrc_eval env e3
            | _ => None
            end
        in olift aux_if (olift checkLocal (dnnrc_eval env e1))
      | DNNRCEither _ ed xl el xr er =>
        match olift checkLocal (dnnrc_eval env ed) with
        | Some (dleft dl) =>
          dnnrc_eval ((xl,Dlocal dl)::env) el
        | Some (dright dr) =>
          dnnrc_eval ((xr,Dlocal dr)::env) er
        | _ => None
        end
      | DNNRCGroupBy _ g sl e1 =>
        None
      | DNNRCCollect _ e1 =>
        let collected := olift checkDistrAndCollect (dnnrc_eval env e1) in
        lift Dlocal collected
      | DNNRCDispatch _ e1 =>
        match olift checkLocal (dnnrc_eval env e1) with
        | Some (dcoll c1) =>
          Some (Ddistr c1)
        | _ => None
        end
      | DNNRCAlg _ plug nnrcbindings =>
        match listo_to_olist (map (fun x =>
               match dnnrc_eval env (snd x) with
               | Some (Ddistr coll) => Some (fst x, coll)
               | _ => None
               end) nnrcbindings) with
        | Some args =>
          match plug_eval h args plug with
          | Some (dcoll c) => Some (Ddistr c)
          | _ => None
          end
        | None => None
        end
      end.

    Require Import DDataNorm.

    Lemma Forall_dcoll_map_lift l:
      Forall (fun x : string * (list data) => data_normalized h (dcoll (snd x))) l ->
      Forall (fun x : string * data => data_normalized h (snd x))
             (map (fun xy : string * list data => (fst xy, dcoll (snd xy))) l).

    Lemma Forall_dcoll_map_lift2 l:
      Forall (fun x : string * data => data_normalized h (snd x))
             (map (fun xy : string * list data => (fst xy, dcoll (snd xy))) l) ->
      Forall (fun x : string * (list data) => data_normalized h (dcoll (snd x))) l.

    Lemma Forall_dcoll_map_lift3 l:
      Forall (fun x : string * data => data_normalized h (snd x))
             (map (fun xy : string * list data => (fst xy, dcoll (snd xy))) l) <->
      Forall (fun x : string * (list data) => data_normalized h (dcoll (snd x))) l.

    Lemma dnnrc_alg_bindings_normalized {plug:AlgPlug plug_type} denv l r :
      Forall (ddata_normalized h) (map snd denv) ->
      Forall
        (fun n : string * dnnrc =>
           forall (denv : dbindings) (o : ddata),
             dnnrc_eval denv (snd n) = Some o ->
             Forall (ddata_normalized h) (map snd denv) -> ddata_normalized h o) r ->
      rmap
         (fun x : string * dnnrc =>
          match dnnrc_eval denv (snd x) with
          | Some (Dlocal _) => None
          | Some (Ddistr coll) => Some (fst x, coll)
          | None => None
          end) r = Some l ->
      Forall (fun x : string * (list data) => data_normalized h (dcoll (snd x))) l.

    Lemma dnnrc_eval_normalized {plug:AlgPlug plug_type} denv e {o} :
      dnnrc_eval denv e = Some o ->
      Forall (ddata_normalized h) (map snd denv) ->
      ddata_normalized h o.

    Corollary dnnrc_eval_normalized_local {plug:AlgPlug plug_type} denv e {d} :
      dnnrc_eval denv e = Some (Dlocal d) ->
      Forall (ddata_normalized h) (map snd denv) ->
      data_normalized h d.
         
  End GenDNNRCBase.

  Section NraEnvPlug.
    Require Import cNRAEnv.
    
    Definition nraenv_eval h aconstant_env op :=
      let aenv := drec nil in
      let aid := dcoll ((drec nil)::nil) in
      nraenv_core_eval h (bindings_of_coll_bindings aconstant_env) op aenv aid.

    Lemma nraenv_eval_normalized h :
      forall op:nraenv_core, forall (constant_env:coll_bindings) (o:data),
      Forall (fun x => data_normalized h (snd x)) (bindings_of_coll_bindings constant_env) ->
      nraenv_eval h constant_env op = Some o ->
      data_normalized h o.

    Global Program Instance NRAEnvPlug : (@AlgPlug nraenv_core) :=
      mkAlgPlug nraenv_eval nraenv_eval_normalized.

    Definition dnnrc_nraenv_core {A} := @dnnrc A nraenv_core.

  End NraEnvPlug.

End DNNRCBase.


Tactic Notation "dnnrc_cases" tactic(first) ident(c) :=
  first;
  [ Case_aux c "DNNRCVar"%string
  | Case_aux c "DNNRCConst"%string
  | Case_aux c "DNNRCBinop"%string
  | Case_aux c "DNNRCUnop"%string
  | Case_aux c "DNNRCLet"%string
  | Case_aux c "DNNRCFor"%string
  | Case_aux c "DNNRCIf"%string
  | Case_aux c "DNNRCEither"%string
  | Case_aux c "DNNRCGroupBy"%string
  | Case_aux c "DNNRCCollect"%string
  | Case_aux c "DNNRCDispatch"%string
  | Case_aux c "DNNRCAlg"%string ].