(* normalize.sml
 *
 * This code is part of the Diderot Project (http://diderot-language.cs.uchicago.edu)
 *
 * COPYRIGHT (c) 2015 The University of Chicago
 * All rights reserved.
 *)

structure Normalize : sig

    val transform : HighIR.program -> HighIR.program

  end = struct

    structure IR = HighIR
    structure Op = HighOps
    structure V = IR.Var
    structure ST = Stats

  (********** Counters for statistics **********)
    val cntUnused               = ST.newCounter "high-opt:unused"

    structure UnusedElim = UnusedElimFn (
        structure IR = IR
        val cntUnused = cntUnused)

    fun useCount (IR.V{useCnt, ...}) = !useCnt

  (* adjust a variable's use count *)
    fun incUse (IR.V{useCnt, ...}) = (useCnt := !useCnt + 1)
    fun decUse (IR.V{useCnt, ...}) = (useCnt := !useCnt - 1)
    fun use x = (incUse x; x)

    fun getEinApp x = (case V.getDef x
	   of IR.EINAPP(e, arg) => SOME(e, arg)
	    | _ => NONE
	  (* end case *))

 (* doNormalize : EIN -> EIN
  * Orders EIN, normalizes it, then cleans the summation
  *)
    fun doNormalize e' = let
          val ordered = Reorder.transform e'
          in
	    case NormalizeEin.transform ordered
	     of NONE => ordered
              | SOME e => EinSums.clean e
	    (* end case *)
          end

    fun nameCnt e = String.concat[V.toString e, "(", Int.toString(useCount e), ")"]
    fun nameCnts e = String.concatWithMap "" nameCnt e

  (* rewriteEin : EIN.Params*int*int*EIN*HighIR Vars* rhs list* EIN  *HighIR Var
   *	-> int*EIN*int*rhs list
   * Orders EIN, normalizes it, then cleans the summation orig-original EIN 
   *)
    fun rewriteEin (params, place, changed, newE, newArgs, done, newEinApp, orig, lhs) = (
	  case List.nth(params, place)
	   of Ein.TEN(false, _) => (changed, orig, place+1, done@[newEinApp])
	    | _ => let
		val rtnArgs = done @ newArgs
		val (c, subst) = Apply.apply(orig, place, newE)
		in
		  if c
		    then (true, subst, place + length newArgs, rtnArgs)
		    else (
		      incUse lhs; List.app incUse newArgs; decUse newEinApp;
		      (true, subst, place + length newArgs, rtnArgs))
		end
	  (* end case *))

  (* doRHS: HighIR.var * rhs -> (var * rhs) list option
   * Looks at each argument to the original EINAPP.
   * If it is another EIN APP calls rewriteEin to do application
   * "place"-The Param-id for the EIN operator.
   * Keeps track of the place of the argument in substitution.
   *)
    fun doRHS (lhs, IR.EINAPP(ein, args)) = let
          fun rewrite (false, _, _, [], _) = NONE
            | rewrite (true, orig, _, [], args') =
		SOME[(lhs, IR.EINAPP(doNormalize orig, args'))]
            | rewrite (changed, orig, place, e::es, args') = (case getEinApp e
                 of NONE => rewrite(changed, orig, place+1, es, args'@[e])
                  | SOME(newE, newA) => let
                      val Ein.EIN{params, index, body} = orig
                      val (changed, e', place', done') =
			    rewriteEin (params, place, changed, newE, newA, args', e, orig, lhs)
                      in
			rewrite(changed, e', place', es, done')
                      end
		(* end case *))
          in
	    rewrite (false, ein, 0, args, [])
          end
      | doRHS _ = NONE

    structure Rewrite = RewriteFn (
      struct
        structure IR = IR
        val doAssign = doRHS

        fun doMAssign _ = NONE
        val elimUnusedVars = UnusedElim.reduce
      end)

    structure Promote = PromoteFn (IR)

    val transform = Promote.transform o Rewrite.transform

(*DEBUG*
fun transform prog = let
      val prog = Rewrite.transform prog
      val _ = HighPP.output(Log.logFile(), "AFTER REWRITE", prog)
      val prog = Promote.transform prog
      val _ = HighPP.output(Log.logFile(), "AFTER PROMOTE", prog)
      in
	prog
      end
*DEBUG*)

  end