(* high-opt.sml
 *
 * COPYRIGHT (c) 2010 The Diderot Project (http://diderot.cs.uchicago.edu)
 * All rights reserved.
 *
 * Optimization of the HighIL representation of Diderot terms.  The main
 * task of this phase is to statically resolve field definitions.
 *)

structure HighOptimizer : sig

    val optimize : HighIL.program -> HighIL.program

  end = struct

    structure IL = HighIL
    structure Op = HighOps
    structure V = IL.Var
    structure ST = Stats
    structure F = FieldDef

    structure Census = CensusFn (IL)

  (********** Counters for statistics **********)
    val cntConstConvolve	= ST.newCounter "high-opt:const-convolve"
    val cntConstField		= ST.newCounter "high-opt:const-field"
    val cntConstDiff		= ST.newCounter "high-opt:const-diff"
    val cntUnused		= ST.newCounter "high-opt:unused"
    val firstCounter            = cntConstConvolve
    val lastCounter             = cntUnused

    datatype binding
      = Unknown
      | OP of Op.rator * IL.var list
      | PHI of IL.var list

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

  (* decrement a variable's use count *)
    fun decUse (IL.V{useCnt, ...}) = (useCnt := !useCnt - 1)

    val {getFn=getBinding, setFn=setBinding, clrFn=clrBinding, ...}
	  = V.newProp (fn _ => Unknown)

  (* optimize the rhs of an assignment, returning NONE if there is no change *)
    fun doRHS rhs = (case rhs
	   of IL.OP(Op.Convolve, [v, h]) => (case (getBinding v, getBinding h)
		 of (OP(Op.LoadImage v', []), OP(Op.Kernel h', [])) => (
		      ST.tick cntConstConvolve;
		      decUse v; decUse h;
		      SOME(IL.OP(Op.Field(F.convolve(v', h')), [])))
		  | _ => raise Fail "non-constant Convolve"
		(* end case *))
	    | IL.OP(Op.AddField, [f, g]) => (case (getBinding f, getBinding g)
		 of (OP(Op.Field f', []), OP(Op.Field g', [])) => (
		      ST.tick cntConstField;
		      decUse f; decUse g;
		      SOME(IL.OP(Op.Field(F.SUM(f', g')), [])))
		  | _ => NONE
		(* end case *))
	    | IL.OP(Op.NegField, [f]) => (case (getBinding f)
		 of OP(Op.Field f', []) => (
		      ST.tick cntConstField;
		      decUse f;
		      SOME(IL.OP(Op.Field(F.neg f'), [])))
		  | _ => NONE
		(* end case *))
	    | IL.OP(Op.DiffField, [f]) => (case (getBinding f)
		 of OP(Op.Field f', []) => (
		      ST.tick cntConstDiff;
		      decUse f;
		      SOME(IL.OP(Op.Field(F.diff f'), [])))
		  | _ => raise Fail "non-constant DiffField"
		(* end case *))
	    | _ => NONE
	  (* end case *))

    fun setBindings nodes = let
	  fun doNode (IL.JOIN{phis, ...}) =
		List.app (fn (y, xs) => setBinding(y, PHI xs)) (!phis)
	    | doNode (IL.BLOCK{body, ...}) =
		List.app (fn (y, IL.OP rhs) => setBinding(y, OP rhs) | _ => ()) (!body)
	    | doNode _ = ()
	  in
	    List.app doNode nodes
	  end

  (* simplify expressions *)
    fun simplify nodes = let
	  fun doAssign (y, rhs) = (case doRHS rhs
		 of SOME(rhs' as IL.OP t) => (setBinding(y, OP t); (y, rhs'))
		  | SOME rhs' => (setBinding(y, Unknown); (y, rhs'))
		  | NONE => (y, rhs)
		(* end case *))
	  fun doNode (IL.ND{kind=IL.BLOCK{body, ...}, ...}) = body := List.map doAssign (!body)
	    | doNode _ = ()
	  in
	    List.app doNode nodes
	  end

  (* reduce the code by removing variables with use counts of 0 *)
    fun reduce nodes = let
	  fun checkVar (y, _) = (useCount y > 0) orelse (ST.tick cntUnused; false)
	  fun doNode (IL.ND{kind, ...}) = (case kind
		 of IL.JOIN{phis, ...} => let
		      fun doVar (y, xs) = if (useCount y = 0)
			    then (
			      ST.tick cntUnused;
			      List.app decUse xs;
			      false)
			    else true
		      in
			phis := List.filter doVar (!phis)
		      end
		  | IL.BLOCK{body, ...} => let
		    (* check for unused lhs variables in reverse order *)
		      fun doAssigns [] = []
			| doAssigns ((y, rhs)::r) = let
			    val r = doAssigns r
			    in
			      if (useCount y = 0)
				then (
				  ST.tick cntUnused;
				  case rhs
				   of IL.VAR x => decUse x
				    | IL.OP(_, xs) => List.app decUse xs
				    | IL.CONS xs => List.app decUse xs
				  (* end case *);
				  r)
				else (y, rhs)::r
			    end
		      in
			body := doAssigns (!body)
		      end
		  | _ => ()
		(* end case *))
	  in
	    List.app doNode (List.rev nodes)
	  end

    fun clearBindings nodes = let
	  fun doNode (IL.ND{kind, ...}) = (case kind
		 of IL.JOIN{phis, ...} => List.app (fn (y, xs) => clrBinding y) (!phis)
		  | IL.BLOCK{body, ...} => List.app (fn (y, _) => clrBinding y) (!body)
		  | _ => ()
		(* end case *))
	  in
	    List.app doNode nodes
	  end

    fun loopToFixPt f prog = let
	  fun loop (n, prog) = let
		val () = f prog
		val n' = Stats.sum{from=firstCounter, to=lastCounter}
		in
		  if (n = n') then () else loop(n', prog)
		end
	  in
	    loop (Stats.sum{from=firstCounter, to=lastCounter}, prog)
	  end

    fun optimize (prog as IL.Program{globals, globalInit, actors}) = let
	  val _ = Census.init prog
	  fun doStmt stm = let
		val nodes = IL.sortNodes stm
		in
		  loopToFixPt simplify nodes;
		  loopToFixPt reduce nodes;
		  nodes
		end
	  val globInitNodes = doStmt globalInit
	  fun doMethod (IL.Method{body, ...}) = let
		val nodes = IL.sortNodes body
		in
		  loopToFixPt simplify nodes;
		  loopToFixPt reduce nodes;
		  clearBindings nodes
		end
	  fun doActor (IL.Actor{stateInit, methods, ...}) = let
		val nodes = IL.sortNodes stateInit
		in
		  loopToFixPt simplify nodes;
		  loopToFixPt reduce nodes;
		  List.app doMethod methods;
		  clearBindings nodes
		end
	  val _ = List.app doActor actors
	(* now we are done with the program body, so we can clean up the global initialization *)
	  val _ = clearBindings globInitNodes
	  in
	    IL.Program{
		globals = globals,
		globalInit = globalInit,
		actors = actors
	      }
	  end

  end