Module NNRCMRtoNNRC

Require Import Arith.
Require Import ZArith.
Require Import String.
Require Import List.
Require Import EquivDec.
Require Import Utils.
Require Import BasicRuntime.
Require Import NNRCRuntime.
Require Import NNRCMRRuntime.
Require Import ForeignToReduceOps.

Local Open Scope list_scope.





  

  
  

  


Section NNRCMRtoNNRC.

  Context {fruntime:foreign_runtime}.
  Context {fredop:foreign_reduce_op}.
  Context {ftoredop:foreign_to_reduce_op}.

  Context (h:brand_relation_t).

  Definition gen_apply_fun (f: var * nnrc) (e2: nnrc) : nnrc :=
    let (x, e1) := f in
    NNRCLet x e2
           e1.

  Fixpoint gen_apply_fun_n (f: list var * nnrc) (eff: list (var * dlocalization)) : option nnrc :=
    let (form, e) := f in
    lift (List.fold_right
            (fun t k =>
               let '(x, (y, loc)) := t in
               NNRCLet x (NNRCVar y)
                      k)
            e)
         (zip form eff).


  Definition nnrc_of_mr_map (input: var) (mr_map: map_fun) : nnrc :=
    match mr_map with
    | MapDist (x, n) =>
      NNRCFor x (NNRCVar input) n
    | MapDistFlatten (x, n) =>
      let res_map :=
          NNRCFor x (NNRCVar input) n
      in
      NNRCUnop AFlatten res_map
    | MapScalar (x, n) =>
      NNRCFor x (NNRCVar input) n
    end.

  Definition nnrc_of_mr (m:mr) : option nnrc :=
    match (m.(mr_map), m.(mr_reduce)) with
    | (MapScalar (x, NNRCUnop AColl n), RedSingleton) =>
      Some (gen_apply_fun (x, n) (NNRCVar m.(mr_input)))
    | (_, RedSingleton) =>
      None
    | (map, RedId) =>
      Some (nnrc_of_mr_map (mr_input m) map)
    | (map, RedCollect red_fun) =>
      let map_expr := nnrc_of_mr_map (mr_input m) map in
      Some (gen_apply_fun red_fun map_expr)
    | (map, RedOp op) =>
      match op with
      | RedOpForeign frop =>
        let map_expr := nnrc_of_mr_map (mr_input m) map in
        lift (fun op => NNRCUnop op map_expr) (foreign_to_reduce_op_to_unary_op op)
      end
    end.

  Fixpoint nnrc_of_mr_chain (outputs: list var) (l: list mr) (k: nnrc) : option nnrc :=
    match l with
    | nil => Some k
    | mr :: l =>
      match (nnrc_of_mr mr, nnrc_of_mr_chain (mr_output mr :: outputs) l k) with
      | (Some n, Some k) =>
        Some (NNRCLet
                (mr_output mr)
                (if in_dec equiv_dec (mr_output mr) outputs then
                   NNRCBinop AUnion (NNRCVar (mr_output mr)) n
                 else n)
                k)
      | _ => None
      end
    end.

  Definition nnrc_of_nnrcmr (l: nnrcmr) : option nnrc :=
    let (last_fun, last_args) := mr_last l in
    let k := gen_apply_fun_n last_fun last_args in
    olift (nnrc_of_mr_chain nil (mr_chain l)) k.

End NNRCMRtoNNRC.