(* var-analysis.sml
 *
 * COPYRIGHT (c) 2011 The Diderot Project (http://diderot-language.cs.uchicago.edu)
 * All rights reserved.
 *
 * The variable analysis phase does four 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 instance variables are invariant over all strands, and thus
 *      can be lifted to global scope, and which strand variables are invariant across
 *      the strand execution.
 *
 *   4) eliminate unused and useless variables.
 *)

structure VarAnalysis : sig

    structure IL : SSA

  (* the scope of a variable describes the extent of its use. *)
    datatype scope
      = Global			(* bound globally and used in strands *)
      | StrandParam		(* parameter to strand creation *)
      | StrandConstState	(* state variable that is invariant over strand execution *)
      | StrandState		(* strand state variable *)
      | Local			(* local binding not used in nested scopes *)

    val scopeToString : scope -> string

    val varScope : IL.var -> scope

    val isGlobal : IL.var -> bool

    val optimize : IL.program -> IL.program

  end = struct

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

    datatype scope
      = Global			(* bound globally and used in strands *)
      | StrandParam		(* parameter to strand creation *)
      | StrandConstState	(* state variable that is invariant over strand execution *)
      | StrandState		(* strand state variable *)
      | Local			(* local binding not used in nested scopes *)

    fun scopeToString s = (case s
	   of Global => "Global"
	    | StrandParam => "StrandParam"
	    | StrandConstState => "StrandConstState"
	    | StrandState => "StrandState"
	    | Local => "Local"
	  (* end case *))

  (* nesting depths of scopes *)
    val undefinedScope = 0
    val globalScope = 1		(* global initialization *)
    val initiallyScope = 2	(* initial strand creation *)
    val paramScope = 2		(* strand parameter *)
    val strandScope = 3		(* strand initialization *)
    val methodScope = 4		(* defined/used in a strand method *)
    val methodStateScope = 5	(* 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 : IL.var -> int option, setFn=setUseDepth, clrFn=clrUseDepth, ...} =
	  V.newProp (fn x => raise Fail("no use depth set for " ^ V.toString x))
    fun getUseDepth x = (case peekUseDepth x
	   of NONE => undefinedScope
	    | SOME d => d
	  (* end case *))

  (* property that maps variables to their assigned scope *)
    val {getFn=varScope : IL.var -> scope, setFn=setVarScope, ...} = V.newProp (fn x => Local)

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

    fun isGlobal x = (varScope x = Global)

  (* walk the program determining the binding and maximum use depths of all variables *)
    fun assignDepths (IL.Program{
	  globalInit,
	  initially = IL.Initially{rangeInit, iters, create, ...},
	  strands
	}) = let
	  fun setMaxUse depth x = (case peekUseDepth x
		 of NONE => setUseDepth(x, depth)
		  | SOME d => if (depth > d) then setUseDepth(x, depth) else ()
		(* end case *))
	  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 doCFG (depth, cfg) = IL.CFG.apply (doNode depth) cfg
	  fun doMethod state (IL.Method{stateIn, body, ...}) = let
		val stateInVec = Vector.fromList stateIn
	      (* check to see which state variables are invariant on the path from
	       * the method entry to this point.
	       *)
		fun chk (_, []) = ()
		  | chk (i, x::xs) = let
		      val x' = Vector.sub(stateInVec, i)
		      in
			if V.same(x, x')
			  then ()
			  else setVarying (x', true);
			chk (i+1, xs)
		       end
		fun doMethodNode nd = (case IL.Node.kind nd
		       of IL.EXIT{live, ...} => chk (0, live)
			| _ => doNode methodScope nd
		      (* end case *))
	      (* this function is used to propagate info from stateIn variables to state *)
		fun propInfo ((_, x), x') = (
		      if isVarying x' then setVarying(x, true) else ();
		      setMaxUse (getUseDepth x') x)
		in
		  List.app (fn x => setBindDepth(x, methodStateScope)) stateIn;
		  IL.CFG.apply doMethodNode body;
		(* need to propagate maxUse from stateIn variables to state variables *)
		  ListPair.appEq propInfo (state, stateIn)
		end
	  fun doStrand (IL.Strand{params, state, stateInit, methods, ...}) = (
		List.app (fn x => setBindDepth(x, paramScope)) params;
		doCFG (strandScope, stateInit);
		List.app (doMethod state) methods)
	  in
	    doCFG (globalScope, 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 doNode nd = let
		fun setBindDepth x = let
		      val scope = if getUseDepth x > globalScope then Global else Local
		      in
			setVarScope (x, scope)
		      end
		in
		  List.app setBindDepth (IL.Node.defs nd)
		end
	  val _ = IL.CFG.apply doNode globalInit
	(* create a new exit node that only has those globals with Global scope *)
	  val newExit = let
		val IL.EXIT{kind, live, ...} = IL.Node.kind(IL.CFG.exit globalInit)
		val newLive = List.filter isGlobal live
		in
		  IL.Node.mkEXIT (kind, newLive)
		end
	  in
	    IL.CFG.replaceNode (IL.CFG.exit globalInit, newExit);
	    IL.CFG{
		entry = IL.CFG.entry globalInit,
		exit = newExit
	      }
	  end

  (* Assign variable scopes to variables that are initialized in the initial-strand creation
   * code.  Since the scoping rules for Diderot limit these variables to be used in this
   * scope, all of them are marked as having Local scope,
   *)
    fun assignScopeForInitially (initially as IL.Initially{rangeInit, iters, create, ...}) = let
	  fun setScope x = setVarScope(x, Local)
	  fun doNode nd = List.app setScope (IL.Node.defs nd)
	  in
	    IL.CFG.apply doNode rangeInit;
	    List.app (fn (i, _, _) => setScope i) iters;
	    IL.CFG.apply doNode (#1 create);
	    initially
	  end

    fun rename (env, x) = (case VMap.find (env, x)
	   of NONE => x
	    | SOME x' => x'
	  (* end case *))

  (* Assign variable scopes to variables that are used in a strand *)
    fun assignScopeForStrand (strand as IL.Strand{name, params, state, stateInit, methods}) = let
	(* assign variable scopes for variables that are defined in the stateInit code *)
	  val _ = let
		fun doVar x = if (getUseDepth x > strandScope)
		      then setVarScope(x, StrandState)
		      else setVarScope(x, Local)
		fun doNode nd = List.app doVar (IL.Node.defs nd)
		in
		  IL.CFG.apply doNode stateInit
		end
	(* assign variable scope for the state variables; state variables (other than the
	 * output) that are not used in the strand methods can be reclassified as Local.
	 * Note that this pass may change some of the assignments from the processing of stateInit.
	 *)
	  val filteredState = let
		fun setScope (isOut, x) =
		      if isVarying x
			then (setVarScope(x, StrandState); true)
		      else if isOut orelse (getUseDepth x > strandScope)
			then (setVarScope(x, StrandConstState); true)
			else (
			  LowILCensus.dec x;  (* we are removing the implicit use *)
			  setVarScope(x, Local);
			  false)
		in
		  List.filter setScope state
		end
	(* 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, inits)) = (
		      setVarScope (x, StrandParam);  (* set the scope while we're at it *)
		      if getUseDepth x > strandScope
			then let
			  val x' = V.copy x
			  val init = (x', IL.VAR x)
			  in
			    setUseDepth(x', getUseDepth x);
			    setVarScope (x', StrandConstState);
			    (x::usedParams, (false, x')::shadowVars, init::inits)
			  end
			else (usedParams, shadowVars, inits))
		val (xs, ys, stms) = List.foldr chkParam ([], [], []) params
		in
		  (xs, ys, IL.CFG.mkBlock stms)
		end
	(* extend the stateInit CFG with the shadow-variable initialization *)
	  val stateInit = if IL.CFG.isEmpty initShadowVars
		then stateInit
		else let
		  val initEntry = IL.CFG.entry stateInit
		  val IL.ENTRY{succ as ref fstNd}  = IL.Node.kind initEntry
		  val IL.CFG{entry, exit} = initShadowVars
		  in
		    IL.Node.replaceInEdge {src = initEntry, oldDst = fstNd, dst = entry};
		    IL.Node.replaceOutEdge {oldSrc = initEntry, src = exit, dst = fstNd};
		    stateInit
		  end
	(* rewrite the exit node of the stateInit to only return the strand state variables *)
	  val stateInit = let
		val IL.EXIT{kind, live, ...} = IL.Node.kind(IL.CFG.exit stateInit)
		fun isStrandState x = (case varScope x
		       of StrandConstState => true
			| StrandState => true
			| _ => false
		      (* end case *))
		val live' = List.filter isStrandState live
		val newExit = IL.Node.mkEXIT(kind, (List.map #2 shadowVars) @ live')
		in
		  IL.CFG.replaceNode (IL.CFG.exit stateInit, newExit);
		  IL.CFG{entry = IL.CFG.entry stateInit, exit = newExit}
		end
	(* assign variable scopes for variables in a method and rename uses of the parameters *)
	  fun doMethod (IL.Method{name, stateIn, body}) = let
	      (* filter out any state variables that have been given Local scope and set
	       * the scope for the local copies of the variables.
	       *)
		val filteredStateIn = let
		      fun chk ((_, x), x', xs) = let
			    val scope = varScope x
			    in
			      setVarScope (x', scope);
			      if scope = Local then xs else x'::xs
			    end
		      in
			ListPair.foldrEq chk [] (state, stateIn)
		      end
	      (* generate the extra shadow variables and initialize the environment to map
	       * parameters to their shadows.
	       *)
		val (env, extraVars) = List.foldr
		      (fn (x, (env, xs)) => let
			  val x' = IL.Var.copy x
			  in
			    setUseDepth(x', getUseDepth x);
			    setVarScope (x', StrandConstState);
			    (VMap.insert(env, x, x'), x'::xs)
			  end
			) (VMap.empty, []) usedParams  
	      (* set the variable scope for variables defined in a method *)
		fun setScope x = if getUseDepth x > methodScope
		      then setVarScope (x, StrandState)
		      else setVarScope (x, Local)
	      (* in case the exit node get rewritten, we need to reset it *)
		val exitNd = ref(IL.CFG.exit body)
	      (* do we need to filter the output states of methods? *)
		val outStateNeedsFiltering = List.exists (fn (_, x) => varScope x <> StrandState) state
	      (* rewrite a node if necessary *)
		fun doNode nd = let
		      fun changed x = VMap.inDomain(env, x)
		      fun renameList (env, xs) = List.map (fn x => rename(env, x)) xs
		      in
			case IL.Node.kind nd
			 of IL.JOIN{phis, ...} => let
			      fun doPhi (lhs, rhs) = (
				    setScope lhs;
				    (lhs, renameList(env, 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(env, cond),
					  trueBranch = !trueBranch,
					  falseBranch = !falseBranch
					}
				  in
				    IL.Node.replaceInEdge {src = !pred, oldDst = nd, dst = newNd};
				    IL.Node.replaceOutEdge {oldSrc = nd, src = nd, dst = !trueBranch};
				    IL.Node.replaceOutEdge {oldSrc = nd, src = nd, dst = !falseBranch}
				  end
				else ()
			  | IL.ASSIGN{stm=(lhs, rhs), ...} => let
			      fun replace rhs = IL.CFG.replaceNode(nd, IL.Node.mkASSIGN(lhs, rhs))
			      in
				setScope lhs;
				case rhs
				 of IL.VAR x => if changed x
				      then replace (IL.VAR(rename(env, x)))
				      else ()
				  | IL.LIT _ => ()
				  | IL.OP(rator, args) => if List.exists changed args
				      then replace (IL.OP(rator, renameList(env, args)))
				      else ()
				  | IL.APPLY(f, args) => if List.exists changed args
				      then replace (IL.APPLY(f, renameList(env, args)))
				      else ()
				  | IL.CONS(ty, args) => if List.exists changed args
				      then replace (IL.CONS(ty, renameList(env, args)))
				      else ()
				(* end case *)
			      end
			  | IL.NEW{strand, args, ...} =>
			      if List.exists changed args
				then IL.CFG.replaceNode(
				  nd,
				  IL.Node.mkNEW{strand=strand, args=renameList(env, args)})
				else ()
			  | IL.EXIT{kind, live, ...} =>
			      if outStateNeedsFiltering
				then let (* filter out StradConstState and Local vars *)
				  fun chkVar ((_, x), x', xs) = if (varScope x <> StrandState)
					then xs
					else x' :: xs
				  val live' = ListPair.foldrEq chkVar [] (state, live)
				  val newNd = IL.Node.mkEXIT(kind, live')
				  in
				    if IL.Node.same(nd, !exitNd)
				      then (
					IL.CFG.replaceNode (nd, newNd);
					exitNd := newNd)
				      else ();
				    IL.CFG.replaceNode (nd, newNd)
				  end
				else ()
			  | _ => ()
			(* end case *)
		      end
		val _ = IL.CFG.apply doNode body
		in
		  IL.Method{
		      name = name,
		      stateIn = extraVars @ filteredStateIn,
		      body = IL.CFG{entry = IL.CFG.entry body, exit = !exitNd}
		    }
		end
	  in
	    IL.Strand{
		name = name,
		params = params,
		state = shadowVars @ filteredState,
		stateInit = stateInit,
		methods = List.map doMethod methods
	      }
	  end

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

  end