(* high-opt.sml
 *
 * COPYRIGHT (c) 2010 The Diderot Project (http://diderot-language.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

  (********** Counters for statistics **********)
    val cntProbeAdd		= ST.newCounter "high-opt:probe-add"
    val cntProbeScale		= ST.newCounter "high-opt:probe-scale"
    val cntProbeNeg		= ST.newCounter "high-opt:probe-neg"
    val cntDiffField		= ST.newCounter "high-opt:diff-field"
    val cntDiffAdd		= ST.newCounter "high-opt:diff-add"
    val cntDiffScale		= ST.newCounter "high-opt:diff-scale"
    val cntDiffNeg		= ST.newCounter "high-opt:diff-neg"
    val cntUnused		= ST.newCounter "high-opt:unused"
    val firstCounter            = cntProbeAdd
    val lastCounter             = cntUnused

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

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

    fun getRHS x = (case V.binding x
	   of IL.VB_RHS(IL.OP arg) => SOME arg
	    | _ => NONE
	  (* end case *))

  (* optimize the rhs of an assignment, returning NONE if there is no change *)
    fun doRHS (lhs, IL.OP rhs) = (case rhs
	   of (Op.Probe(domTy, rngTy), [f, pos]) => (case getRHS f
		 of SOME(Op.Field _, _) => NONE (* direct probe does not need rewrite *)
		  | SOME(Op.AddField, [f, g]) => raise Fail "Probe(f+g)"
		  | SOME(Op.ScaleField, [s, f]) => raise Fail "Probe(s*f)"
		  | SOME(Op.NegField, [f]) => raise Fail "Probe(-f)"
		  | _ => raise Fail(concat[
			"bogus field binding ", V.toString f, " = ", IL.vbToString(V.binding f)
		      ])
		(* end case *))
	    | (Op.DiffField, [f]) => (case (getRHS f)
		 of SOME(Op.Field dim, [v, h]) => (case getRHS h
		       of SOME(Op.Kernel(kernel, k), []) => let
			    val h' = IL.Var.copy h
			    in
			      ST.tick cntDiffField;
			      decUse f;
			      incUse h'; incUse v;
			      SOME[
				  (h', IL.OP(Op.Kernel(kernel, k+1), [])),
				  (lhs, IL.OP(Op.Field dim, [v, h']))
				]
			    end
			| _ => raise Fail(concat[
			      "bogus kernel binding ", V.toString h, " = ", IL.vbToString(V.binding h)
			    ])
		      (* end case *))
		  | SOME(Op.AddField, [f, g]) => raise Fail "Diff(f+g)"
		  | SOME(Op.ScaleField, [s, f]) => raise Fail "Diff(s*f)"
		  | SOME(Op.NegField, [f]) => raise Fail "Diff(-f)"
		  | _ => NONE
		(* end case *))
	    | _ => NONE
	  (* end case *))
      | doRHS _ = NONE

  (* simplify expressions *)
    fun simplify (nd as IL.ND{kind=IL.ASSIGN{stm=(y, rhs), ...}, ...}) =
	  if (useCount y = 0)
	    then () (* skip unused assignments *)
	    else (case doRHS(y, rhs)
	       of SOME[] => IL.CFG.deleteNode nd
		| SOME assigns => (
		    List.app (fn (y, rhs) => V.setBinding(y, IL.VB_RHS rhs)) assigns;
		    IL.CFG.splice (nd, IL.CFG.mkBlock assigns))
		| NONE => ()
	      (* end case *))
      | simplify _ = ()

  (* reduce the code by removing variables with use counts of 0 *)
    local
      fun checkVar (y, _) = (useCount y > 0) orelse (ST.tick cntUnused; false)
    in
    fun reduce (nd as 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.ASSIGN{stm=(y, rhs), ...} =>
		if (useCount y = 0)
		  then (
		    ST.tick cntUnused;
		    case rhs
		     of IL.VAR x => decUse x
		      | IL.LIT _ => ()
		      | IL.OP(_, xs) => List.app decUse xs
		      | IL.CONS xs => List.app decUse xs
		    (* end case *);
		    IL.CFG.deleteNode nd)
		  else ()
	    | _ => ()
	  (* end case *))
    end (* local *)

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

    fun optimize (prog as IL.Program{globals, globalInit, initially, strands}) = let
	  fun doCFG cfg = (
		loopToFixPt simplify cfg;
		loopToFixPt reduce cfg)
	  val globInitNodes = doCFG globalInit
	  fun doMethod (IL.Method{body, ...}) = doCFG body
	  fun doStrand (IL.Strand{stateInit, methods, ...}) = (
		doCFG stateInit;
		List.app doMethod methods)
	  val _ = List.app doStrand strands
	  in
	    IL.Program{
		globals = globals,
		globalInit = globalInit,
		initially = initially,	(* FIXME: we should optimize this code *)
		strands = strands
	      }
	  end

  end