(* 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 = IL.Op
    structure V = IL.Var
    structure ST = Stats

    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            = cntUnusedStmt
    val lastCounter             = cntUnusedCFun

    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.LoadImage v', []), (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.Field f', []), (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.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.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, _) => setBinding(y, PHI xs)) (!phis)
	    | doNode (IL.BLOCK{body, ...}) =
		List.app (fn (y, IL.OP rhs) => setBinding(y, RHS 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, RHS t); (y, rhs'))
		  | SOME rhs' => (setBinding(y, Unknown); (y, rhs'))
		  | NONE => (y, rhs)
		(* end case *))
	  fun doNode (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.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
	    | doNode (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
		in
		  body := doAssigns (!body)
		end
	    | doNode _ = ()
	  in
	    List.app doNode (List.rev nodes)
	  end

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

    fun loopToFixPt f = let
	  fun loop (n, prog) = let
		val prog = f prog
		val n' = Stats.sum{from=firstCounter, last=lastCounter}
		in
		  if (n = n') then prog else loop(n', prog)
		end
	  in
	    loop (Stats.sum{from=firstCounter, last=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