Module Qcert.NNRCMR.Optim.NNRCMROptimizer

Require Import String.
Require Import List ListSet.
Require Import Arith.
Require Import Equivalence.
Require Import Morphisms.
Require Import Setoid.
Require Import EquivDec.
Require Import Program.
Require Import Utils.
Require Import DataSystem.
Require Import cNNRCSystem.
Require Import NNRCOptim.
Require Import OptimizerLogger.
Require Import OptimizerStep.
Require Import NNRCMR.
Require Import NNRCMRRewrite.
Require Import ForeignReduceOps.

Section NNRCMROptimizer.
  Definition trew_nnrcmr
             {fruntime:foreign_runtime} {fredop:foreign_reduce_op} {logger:optimizer_logger string nnrc}
             (l: nnrcmr) :=
    let inputs_loc := l.(mr_inputs_loc) in
    let chain :=
        List.map
          (fun mr =>
             let map :=
                 match mr.(mr_map) with
                 | MapDist (x, n) => MapDist (x, nnrc_optim_top_default n)
                 | MapDistFlatten (x, n) => MapDistFlatten (x, nnrc_optim_top_default n)
                 | MapScalar (x, n) => MapScalar (x, nnrc_optim_top_default n)
                 end
             in
             let reduce :=
                 match mr.(mr_reduce) with
                 | RedId => RedId
                 | RedCollect (x, n) => RedCollect (x, nnrc_optim_top_default n)
                 | RedOp op => RedOp op
                 | RedSingleton => RedSingleton
                 end
             in
             mkMR mr.(mr_input) mr.(mr_output) map reduce)
          l.(mr_chain)
    in
    let last :=
        let '((params, n), args) := l.(mr_last) in
        ((params, nnrc_optim_top_default n), args)
    in
    mkMRChain
      inputs_loc
      chain
      last.

  Definition nnrcmr_optim_top
             {fruntime:foreign_runtime} {fredop:foreign_reduce_op} {logger:optimizer_logger string nnrc}
             q :=
    let q := trew_nnrcmr (mr_optimize q) in
    trew_nnrcmr q.
  
End NNRCMROptimizer.