(* lift.sml
 *
 * COPYRIGHT (c) 2010 The Diderot Project (http://diderot.cs.uchicago.edu)
 * All rights reserved.
 *
 * Lift field operations to global scope and split global initialization to
 * pre and post image loading phases.
 *
 * NOTE: this process can be streamlined as follows:
 *	1) identify the static variables
 *	2) evaluate eagerly, with the proviso that if the lhs is not static, then
 *	   only evaluate if the rhs is static and the operation is supported
 *	3) reduce code
 * This evaluation process could be extended to the body of the actors too.
 *)

structure Lift : sig

    val transform : Simple.program -> Simple.program

  end = struct

    structure BV = BasisVars
    structure S = Simple
    structure VSet = Var.Set
    structure VMap = Var.Map

  (* the kinds of things a variable in Simple AST can be bound to *)
    datatype var_binding
      = RHS of S.exp
      | Param

  (* identify the image load operations and lift them and their antecedents; in terms of BTA,
   * this phase is essentially determining what must be static in order to get the image
   * info needed for the rest of the compile.
   *)
    fun liftLoads block = let
	(* analysis to compute the set of static variables *)
	  fun mkStatic (env, statics, x) = if VSet.member(statics, x)
		then statics
		else let
		  val statics = VSet.add(statics, x)
		  in
		    case VMap.find(env, x)
		     of SOME(S.E_Var y) => mkStatic (env, statics, y)
		      | SOME(S.E_Tuple ys) => mkStatics (env, statics, ys)
		      | SOME(S.E_Apply(_, _, ys, _)) => mkStatics (env, statics, ys)
		      | SOME(S.E_Cons ys) => mkStatics (env, statics, ys)
		      | SOME(S.E_Input(_, _, SOME y)) => mkStatic (env, statics, y)
		      | SOME _ => statics
		      | NONE => raise Fail(concat["variable ", Var.uniqueNameOf x, " has no binding"])
		    (* end case *)
		  end
	  and mkStatics (env, statics, xs) =
		List.foldl (fn (x, statics) => mkStatic(env, statics, x)) statics xs
	  fun doBlock (env, statics, S.Block stms) = let
		fun doStmts (env, statics, []) = statics
		  | doStmts (env, statics, stm::stms) = let
		      val (env, statics) = doStmt (env, statics, stm)
		      in
			doStmts (env, statics, stms)
		      end
		in
		  doStmts (env, statics, stms)
		end
	  and doStmt (env, statics, stm) = (case stm
		 of S.S_Assign(x, e) => let
		      val env = VMap.insert(env, x, e)
		      in
			case e
			 of S.E_Apply(f, _, xs, _) =>
			      if Var.same(f, BV.fn_load)
				then (env, mkStatic(env, statics, x))
				else (env, statics)
			  | _ => (env, statics)
			(* end case *)
		      end
		  | S.S_IfThenElse(x, b1, b2) => let
		      val statics1 = doBlock (env, statics, b1)
		      val statics2 = doBlock (env, statics, b2)
		      val n = VSet.numItems statics
		      in
			if ((n <> VSet.numItems statics1) orelse (n <> VSet.numItems statics2))
			  then (env, mkStatic(env, statics, x))
			  else (env, statics)
		      end
		  | _ => (env, statics)
		(* end case *))
	  val statics = doBlock (VMap.empty, VSet.empty, block)
	(* lift out the static code *)
	  fun doBlock (S.Block stms) = let
		fun doStmts ([], staticStms) = S.Block(List.rev staticStms)
		  | doStmts (stm::stms, staticStms) = (case doStmt stm
		       of SOME stm => doStmts (stms, stm::staticStms)
			| NONE => doStmts (stms, staticStms)
		      (* end case *))
		in
		  doStmts (stms, [])
		end
	  and doStmt stm = (case stm
		 of S.S_Assign(x, e) => if VSet.member(statics, x)
		      then SOME stm
		      else NONE
		  | S.S_IfThenElse(x, b1, b2) => if VSet.member(statics, x)
		      then SOME(S.S_IfThenElse(x, doBlock b1, doBlock b2))
		      else NONE
		  | _ => NONE
		(* end case *))
	  val staticBlock = doBlock block
	  in
Log.msg "**** static variables: ";
VSet.app (fn x => Log.msg(" "^Var.uniqueNameOf x)) statics;
Log.msg "\n";
	    staticBlock
	  end

  (* given values for the static variables; reduce the static initialization code *)
    fun reduce (env, blk) = let
	  fun doBlock (S.Block stms) =
		List.foldr (fn (stm, stms) => doStmt stm @ stms) [] stms
	  and doStmt stm = (case stm
		 of S.S_Assign(x, e) => (case Var.Map.find(env, x)
		       of SOME v => let
			    val rhs = (case v
				   of (Eval.BV b) => S.E_Lit(Literal.Bool b)
				    | (Eval.SV s) => S.E_Lit(Literal.String s)
				    | (Eval.IV i) => S.E_Lit(Literal.Int i)
				    | (Eval.TV _) => e
				    | (Eval.FV fld) => S.E_Field fld
				    | (Eval.Img info) => S.E_LoadImage info
				    | (Eval.KV h) => e
				  (* end case *))
			    in
			      [S.S_Assign(x, rhs)]
			    end
			| NONE => [stm]
		      (* end case *))
		  | S.S_IfThenElse(x, b1, b2) => (case Var.Map.find(env, x)
		       of SOME(Eval.BV b) => if b then doBlock b1 else doBlock b2
			| NONE => [stm]
		      (* end case *))
		  | _ => [stm]
		(* end case *))
	  in
	    S.Block(doBlock blk)
	  end

    fun transform (prog as S.Program{globals, globalInit, actors}) = let
	  val staticInit = liftLoads globalInit
	  val staticEnv = Eval.evalStatics staticInit
	  val globalInit = reduce (staticEnv, globalInit)
	  in
	    S.Program{
		globals = globals,
		globalInit = globalInit,
		actors = actors
	      }
	  end

  end