(* var-analysis.sml
 *
 * COPYRIGHT (c) 2011 The Diderot Project (http://diderot-language.cs.uchicago.edu)
 * All rights reserved.
 *
 * The variable analysis phase does the following things:
 *
 *   1) determine which variables that are bound at global level are used in strand
 *      or method scopes.
 *
 *   2) determine which strand parameters are used in methods (and thus have to be
 *      shadowed by instance variables).
 *
 *   3) determine which strand state variables are invariant across the strand
 *      execution.
 *
 *   4) eliminate unused and useless variables.
 *
 * TODO
 *
 *   5) determine which state variables are invariant over all strands and thus
 *      can be lifted to global scope
 *
 *   6) determine which state variables are not used and thus can be eliminated.
 *)

structure VarAnalysis : sig

    structure IL : SSA

    val optimize : IL.program -> IL.program

  (* returns true if the state variable is invariant over strand execution *)
    val isVarying : IL.state_var -> bool

  end = struct

    structure IL = LowIL
    structure V = IL.Var
    structure StV = IL.StateVar
    structure VMap = V.Map

  (* nesting depths of scopes *)
    val undefinedScope = 0
    val globalScope = 1		(* global initialization *)
    val inputScope = 2		(* global inputs *)
    val initiallyScope = 3	(* initial strand creation *)
    val paramScope = 3		(* strand parameter *)
    val strandScope = 4		(* strand initialization *)
    val methodScope = 5		(* defined/used in a strand method *)
    val methodStateScope = 6	(* state variable in a strand method *)

  (* property that maps variables to the depth where they are bound *)
    val {getFn=getBindDepth : IL.var -> int, setFn=setBindDepth, clrFn=clrBindDepth, ...} =
	  V.newProp (fn x => raise Fail("no binding depth set for " ^ V.toString x))

  (* property that maps variables to the maximum depth where they are used *)
    val {peekFn=peekUseDepth, getFn=getUseDepth, setFn=setUseDepth, clrFn=clrUseDepth, ...} =
	  V.newProp (fn x => undefinedScope)

  (* boolean flag on strand-state variables that vary across method executions *)
    val {getFn=isVarying, setFn=setVarying} = StV.newFlag ()

  (* walk the program determining the binding and maximum use depths of all variables *)
    fun assignDepths (IL.Program{
	  props,
	  globalInit,
	  initially = IL.Initially{rangeInit, iters, create, ...},
	  strands
	}) = let
	  fun setMaxUse depth x = if (depth > getUseDepth x)
                then setUseDepth(x, depth)
                else ()
	  fun doNode depth nd = (case IL.Node.kind nd
		 of IL.EXIT{kind, live, ...} => (case kind
		       of ExitKind.FRAGMENT => List.app (setMaxUse depth) live
			| ExitKind.RETURN => List.app (setMaxUse depth) live
			| _ => () (* we don't count ACTIVE, etc. as uses *)
		      (* end case *))
		  | _ => (
		      List.app (setMaxUse depth) (IL.Node.uses nd);
		      List.app (fn x => setBindDepth(x, depth)) (IL.Node.defs nd))
		(* end case *))
	  fun doGlobalInitNode nd = (case IL.Node.kind nd
		 of IL.ASSIGN{stm=(x, IL.OP(LowOps.Input _, _)), ...} => setMaxUse inputScope x
		  | _ => doNode globalScope nd
		(* end case *))
	  fun doCFG (depth, cfg) = IL.CFG.apply (doNode depth) cfg
	  fun assignDepthsForMethod state (IL.Method{body, ...}) = let
	      (* check to see which state variables are invariant *)
		fun doMethodNode nd = (case IL.Node.kind nd
		       of IL.SAVE{lhs, rhs, ...} => let
                            fun varying () = (
                                  setVarying(lhs, true);
                                  setMaxUse methodScope rhs)
                            in
                            (* check to see if the state variable is varying on the path to
                             * this node.
                             *)
                              case V.binding rhs
                               of IL.VB_RHS(IL.STATE x) =>
                                    if StV.same(x, lhs)
                                    (* we don't count non-varying save uses of variables as setting
                                     * the scope of a variable.  This will allow us to determine
                                     * which state variables are actually used in a method, since
                                     * those will be bound to variables with methodScope.
                                     *)
                                      then setMaxUse strandScope rhs
                                      else varying()
                                | _ => varying()
                              (* end case *)
                            end
			| _ => doNode methodScope nd
		      (* end case *))
		in
		  IL.CFG.apply doMethodNode body
		end
	(* assign use depths for a strand.  The state variables are a bit tricky.  We first
	 * propagate information from the stateInit and methods up to the state-variable
	 * list.  Then we copy the accumulated use-depth information back down to the
	 * corresponding variables in the stateInit and methods.
	 *)
	  fun doStrand (IL.Strand{params, state, stateInit, methods, ...}) = (
		List.app (fn x => setBindDepth(x, paramScope)) params;
	      (* assign depths for the stateInit code *)
		doCFG (strandScope, stateInit);
              (* examine the methods *)
		List.app (assignDepthsForMethod state) methods)
	  in
	    IL.CFG.apply doGlobalInitNode globalInit;
	  (* do initially code *)
	    doCFG (initiallyScope, rangeInit);
	    List.app
	      (fn (i, min, max) => (
		  setBindDepth (i, initiallyScope);
		  setMaxUse initiallyScope min;
		  setMaxUse initiallyScope max)
	      ) iters;
	    doCFG (initiallyScope, #1 create);
	    List.app (setMaxUse initiallyScope) (#3 create);
	  (* do strands *)
	    List.app doStrand strands
	  end

  (* Assign variable scopes to variables that are initialized in the global initialization.
   * For the global initialization, variables that have a maxUseDepth > globalScope are
   * really global.  The others are local.
   *)
    fun assignScopeForGlobalInit globalInit = let
          fun isGlobal x = (getUseDepth x > globalScope)
	  in
	  (* create a new exit node that only has those globals with Global scope *)
	    IL.CFG.updateExit (globalInit, fn live => List.filter isGlobal live)
	  end

  (* insert a sequence of statements (initCFG) between the entry node of cfg and its
   * successor.
   *)
    fun prependCFG (initCFG, cfg) =
	  if IL.CFG.isEmpty initCFG
            then cfg
            else let
              val initEntry = IL.CFG.entry cfg
              val IL.ENTRY{succ as ref fstNd}  = IL.Node.kind initEntry
              val IL.CFG{entry, exit} = initCFG
              in
                IL.Node.replaceInEdge {src = initEntry, oldDst = fstNd, dst = entry};
                IL.Node.replaceOutEdge {oldSrc = initEntry, src = exit, dst = fstNd};
                cfg
              end

  (* Assign variable scopes to variables that are used in a strand *)
    fun assignScopeForStrand (strand as IL.Strand{name, params, state, stateInit, methods}) = let
(* FIXME: strand state variables that are only used in the stateInit code can be replaced by local vars *)
	(* for any parameter that is used inside methods, we need to introduce a
	 * shadow state variable.  Here we compute the list of parameters that are
	 * used inside strand methods, the new state variables that we create to
	 * shadow them, and the statements to initialize the shadow variables.
	 *)
	  val (usedParams, shadowVars, initShadowVars) = let
		fun chkParam (x, (usedParams, shadowVars, initCFG)) = (
		      if getUseDepth x > strandScope
			then let
			  val x' = StV.new(false, V.name x, V.ty x)
			  val initCFG = IL.CFG.appendNode (initCFG, IL.Node.mkSAVE(x', x))
			  in
			    LowILCensus.inc x;
			    (x::usedParams, x'::shadowVars, initCFG)
			  end
			else (usedParams, shadowVars, initCFG))
		val (xs, ys, initCFG) = List.foldr chkParam ([], [], IL.CFG.empty) params
		in
		  (xs, ys, initCFG)
		end
	(* prepend the stateInit CFG with the shadow-variable initialization *)
	  val stateInit = prependCFG(initShadowVars, stateInit)
	(* assign variable scopes for variables in a method and rename uses of the parameters *)
	  fun doMethod (IL.Method{name, body}) = let
              (* create local names for the parameters and the code to read them out from the
               * shadow state variables.
               *)
                val (env, loadCFG) = let
                      fun f (x::xs, y::ys, env, loadCFG) = let
                            val x' = V.copy x
                            val env = VMap.insert(env, x, x')
                            val loadCFG = IL.CFG.appendNode (loadCFG, IL.Node.mkASSIGN(x', IL.STATE y))
                            in
                              f (xs, ys, env, loadCFG)
                            end
                        | f ([], [], env, loadCFG) = (env, loadCFG)
                      in
                        f (usedParams, shadowVars, VMap.empty, IL.CFG.empty)
                      end
              (* rename used parameters *)
		fun rename x = (case VMap.find (env, x)
		       of NONE => x
			| SOME x' => (
			  (* replacing x with x', so adjust census counts *)
			    LowILCensus.dec x; LowILCensus.inc x';
			    x')
		      (* end case *))
              (* process method-body nodes.  We rename parameters to the local variable that
               * is read out from the shadow state variable and we remove save nodes for non-varying
               * state variables.
               *)
		fun doNode nd = let
		      fun changed x = VMap.inDomain(env, x)
		      val renameList = List.map rename
		      in
			case IL.Node.kind nd
			 of IL.JOIN{phis, ...} => let
			      fun doPhi (lhs, rhs) = (lhs, renameList rhs)
			      in
				phis := List.map doPhi (!phis)
			      end
			  | IL.COND{pred, cond, trueBranch, falseBranch} =>
			      if changed cond
				then let
				  val newNd = IL.Node.mkCOND {
					  cond = rename cond,
					  trueBranch = !trueBranch,
					  falseBranch = !falseBranch
					}
				  in
				    IL.Node.replaceInEdge {src = !pred, oldDst = nd, dst = newNd};
				    IL.Node.replaceOutEdge {oldSrc = nd, src = newNd, dst = !trueBranch};
				    IL.Node.replaceOutEdge {oldSrc = nd, src = newNd, dst = !falseBranch}
				  end
				else ()
			  | IL.ASSIGN{stm=(lhs, rhs), ...} => let
			      fun replace rhs = IL.CFG.replaceNode(nd, IL.Node.mkASSIGN(lhs, rhs))
			      in
				case rhs
				 of IL.STATE x => ()
                                  | IL.VAR x => if changed x
				      then replace (IL.VAR(rename x))
				      else ()
				  | IL.LIT _ => ()
				  | IL.OP(rator, args) => if List.exists changed args
				      then replace (IL.OP(rator, renameList args))
				      else ()
				  | IL.APPLY(f, args) => if List.exists changed args
				      then replace (IL.APPLY(f, renameList args))
				      else ()
				  | IL.CONS(ty, args) => if List.exists changed args
				      then replace (IL.CONS(ty, renameList args))
				      else ()
				(* end case *)
			      end
                          | IL.MASSIGN{stm=(lhs, rator, args), ...} =>
                              if List.exists changed args
                                then IL.CFG.replaceNode(nd, IL.Node.mkMASSIGN(lhs, rator, renameList args))
                                else ()
			  | IL.NEW{strand, args, ...} =>
			      if List.exists changed args
				then IL.CFG.replaceNode(
				  nd,
				  IL.Node.mkNEW{strand=strand, args=renameList args})
				else ()
			  | IL.SAVE{lhs, rhs, ...} => if (isVarying lhs)
                              then ()
                              else IL.CFG.deleteNode nd
			  | _ => ()
			(* end case *)
		      end
		val _ = IL.CFG.apply doNode body
		in
		  IL.Method{
		      name = name,
		      body = prependCFG (loadCFG, body)
		    }
		end
	  in
	    IL.Strand{
		name = name,
		params = params,
		state = shadowVars @ state,
		stateInit = stateInit,
		methods = List.map doMethod methods
	      }
	  end

    fun optimize prog = let
	  val IL.Program{props, globalInit, initially, strands} = prog
	(* first we compute binding and use depths *)
	  val _ = assignDepths prog
        (* then rewrite the code *)
	  val globalInit = assignScopeForGlobalInit globalInit
	  val strands = List.map assignScopeForStrand strands
	  in
	    IL.Program{props=props, globalInit=globalInit, initially=initially, strands=strands}
	  end

  end