(* translate.sml
 *
 * COPYRIGHT (c) 2010 The Diderot Project (http://diderot-language.cs.uchicago.edu)
 * All rights reserved.
 *
 * Translate Simple-AST code into the IL representation.  This translation is based on the
 * algorithm described in
 *
 *	Single-pass generation of static single assignment form for structured languages
 *	ACM TOPLAS, Nov. 1994
 *	by Brandis and MossenBock.
 *)

structure Translate : sig

    val translate : Simple.program -> HighIL.program

  end = struct

    structure S = Simple
    structure Ty = Types
    structure VMap = Var.Map
    structure VSet = Var.Set
    structure IL = HighIL
    structure DstTy = HighILTypes

    type env = IL.var VMap.map

    fun lookup env x = (case VMap.find (env, x)
	   of SOME x' => x'
	    | NONE => raise Fail(concat[
		  "no binding for ", Var.uniqueNameOf x, " in environment"
		])
	  (* end case *))

    fun cvtTy ty = (case TypeUtil.prune ty
	   of Ty.T_Bool => DstTy.BoolTy
	    | Ty.T_Int => DstTy.IntTy
	    | Ty.T_String => DstTy.StringTy
	    | Ty.T_Kernel _ => DstTy.KernelTy
	    | Ty.T_Tensor(Ty.Shape dd) => let
		fun cvtDim (Ty.DimConst 1) = NONE
		  | cvtDim (Ty.DimConst d) = SOME d
		in
		  DstTy.TensorTy(List.mapPartial cvtDim dd)
		end
	    | Ty.T_Image _ => DstTy.ImageTy
	    | Ty.T_Field _ => DstTy.FieldTy
	    | ty => raise Fail("cvtTy: unexpected " ^ TypeUtil.toString ty)
	  (* end case *))

  (* create a new instance of a variable *)
    fun newVar x = IL.Var.new (Var.nameOf x, cvtTy(Var.monoTypeOf x))

  (* a pending-join node tracks the phi nodes needed to join the assignments that flow into
   * the join node.
   *)
    datatype join = JOIN of {
	arity : int,			(* number of predecessors *)
	nd : IL.node,			(* the CFG node for this pending join *)
	phiMap : IL.phi VMap.map ref	(* a mapping from Simple AST variables that are assigned *)
					(* to their phi nodes. *)
      }

  (* create a new pending-join node *)
    fun newJoin arity = JOIN{arity = arity, nd = IL.Node.mkJOIN [], phiMap = ref VMap.empty}

  (* record an assignment to the IL variable dstVar (corresponding to the Simple AST variable
   * srcVar) in the current pending-join node.  The predIndex specifies which path into the
   * JOIN node this assignment occurs on.
   *)
    fun recordAssign (env, JOIN{arity, phiMap, ...}, srcVar, dstVar, predIndex) = let
	  val m = !phiMap
	  val m'= (case VMap.find (m, srcVar)
		 of NONE => let
		      val dstVar' = newVar srcVar
		      val dfltVar = lookup env srcVar
		      val rhs = List.tabulate (arity, fn i => if (i = predIndex) then dstVar else dfltVar)
		      in
			VMap.insert (m, srcVar, (dstVar', rhs))
		      end
		  | SOME(lhs, rhs) => let
		      fun update (i, l as x::r) = if (i = predIndex)
			    then dstVar::r
			    else x::update(i+1, r)
			| update _ = raise Fail "invalid predecessor index"
		      in
			VMap.insert (m, srcVar, (lhs, update(0, rhs)))
		      end
		(* end case *))
	  in
	    phiMap := m'
	  end

  (* complete a pending join operation by filling in the phi nodes from the phi map and
   * updating the environment.
   *)
    fun completeJoin (env, JOIN{nd, phiMap, ...}) = let
	  val IL.ND{kind=IL.JOIN{phis, ...}, ...} = nd
	  fun doVar (srcVar, phi as (dstVar, _), (env, phis)) =
		(VMap.insert (env, srcVar, dstVar), phi::phis)
	  val (env, phis') = VMap.foldli doVar (env, []) (!phiMap)
	  in
	    phis := phis';
	    env
	  end

  (* expression translation *)
    fun cvtExp (env, lhs, exp) = (case exp
	   of S.E_Var x => [(lhs, IL.VAR(lookup env x))]
	    | S.E_Lit lit => [(lhs, IL.LIT lit)]
	    | S.E_Tuple xs => raise Fail "E_Tuple not implemeted"
	    | S.E_Apply(f, tyArgs, args, ty) => let
		val args' = List.map (lookup env) args
		in
		  TranslateBasis.translate (lhs, f, tyArgs, args')
		end
	    | S.E_Cons args => [(lhs, IL.CONS(List.map (lookup env) args))]
	    | S.E_Slice(x, indices, ty) => let
		val x = lookup env x
		val mask = List.map isSome indices
		fun cvt NONE = NONE
		  | cvt (SOME x) = SOME(lookup env x)
		val indices = List.mapPartial cvt indices
		in
		  if List.all (fn b => b) mask
		    then [(lhs, IL.OP(HighOps.Subscript(IL.Var.ty x), x::indices))]
		    else [(lhs, IL.OP(HighOps.Slice(IL.Var.ty lhs, mask), x::indices))]
		end
	    | S.E_Input(_, name, NONE) =>
		[(lhs, IL.OP(HighOps.Input(IL.Var.ty lhs, name), []))]
	    | S.E_Input(_, name, SOME dflt) =>
		[(lhs, IL.OP(HighOps.InputWithDefault(IL.Var.ty lhs, name), [lookup env dflt]))]
	    | S.E_Field fld => [(lhs, IL.OP(HighOps.Field fld, []))]
	    | S.E_LoadImage info => [(lhs, IL.OP(HighOps.LoadImage info, []))]
	  (* end case *))

    fun cvtBlock (env, blk, k : (env * IL.node * IL.node) -> IL.cfg) = raise Fail "cvtBlock"

(*
  (* convert a Simple AST block to an IL statement.  We return the statement that represents the
   * block, plus the environment mapping Simple AST variables to their current SSA representations
   * and the set of Simple AST variables that were assigned to in the block.
   *)
    fun cvtBlock (env, S.Block stms, optExit) = let
	  fun toBlock (env, assigned, [], assignments) = let
		val stm = IL.Stmt.mkBLOCK(List.rev assignments, optExit)
		in
		  (stm, IL.Stmt.tail stm, env, assigned)
		end
	    | toBlock (env, assigned, S.S_Assign(x, e)::rest, assignments) = let
		val x' = newVar x
		val stms = cvtExp(env, x', e)
		val assigned = VSet.add(assigned, x)
		val env = VMap.insert(env, x, x')
		in
		  toBlock (env, assigned, rest, stms@assignments)
		end
	    | toBlock (env, assigned, stms, assignments) = let
		val (next, last, env, assigned) = toStmt (env, assigned, stms)
		val blk = IL.Stmt.mkBLOCK(List.rev assignments, SOME next)
		in
		  IL.Node.addEdge (IL.Stmt.tail blk, IL.Stmt.entry next);
		  (blk, last, env, assigned)
		end
	  and toStmt (env, assigned, []) = let
	      (* this case only occurs for the empty else arm of an if-then-else statement *)
		val stm = IL.Stmt.mkBLOCK([], optExit)
		in
		  (stm, IL.Stmt.tail stm, env, assigned)
		end
	    | toStmt (env, assigned, stms as stm::rest) = (case stm
		 of S.S_Assign _ => toBlock (env, assigned, stms, [])
		  | S.S_IfThenElse(x, b1, b2) => let
		      val x' = lookup env x
		      val (s1, last1, env1, assigned1) = cvtBlock(env, b1, NONE)
		      val (s2, last2, env2, assigned2) = cvtBlock(env, b2, NONE)
		      val assigned = VSet.union(assigned1, assigned2)
(* PROBLEM: what about variables that are assigned for the first time in one branch
 * and not the other?  This situation should only occur for variables who's scope is
 * the branch of the if.  Short-term solution is to ignore variables that are defined
 * in only one branch.
 *)
		      val (env, phis) = let
			    fun mkPhi (x, (env, phis)) = (
				  case (VMap.find(env1, x), VMap.find(env2, x))
				   of (SOME x1, SOME x2) => let
					val x' = newVar x
					in
					  (VMap.insert(env, x, x'), (x', [x1, x2])::phis)
					end
				    | _ => (env, phis)
				  (* end case *))
			    in
			      VSet.foldl mkPhi (env, []) assigned
			    end
		      in
			case rest
			 of [] => let
			      val join = IL.Stmt.mkJOIN (phis, optExit)
			      val joinNd = IL.Stmt.entry join
			      val stm = IL.Stmt.mkIF(x', s1, s2, SOME join)
			      in
				IL.Node.addEdge (last2, joinNd);
				IL.Node.addEdge (last1, joinNd);
				(stm, joinNd, env, assigned)
			      end
			  | _ => let
			      val (next, last, env, assigned) = toStmt (env, assigned, rest)
			      val join = IL.Stmt.mkJOIN (phis, SOME next)
			      val joinNd = IL.Stmt.entry join
			      val stm = IL.Stmt.mkIF(x', s1, s2, SOME join)
			      in
				IL.Node.addEdge (last2, joinNd);
				IL.Node.addEdge (last1, joinNd);
				IL.Node.addEdge (joinNd, IL.Stmt.entry next);
				(stm, last, env, assigned)
			      end
			(* end case *)
		      end
		  | S.S_New(name, xs) => let
		      val xs' = List.map (lookup env) xs
		      in
			case rest
			 of [] => let
			      val stm = IL.Stmt.mkNEW(name, xs', optExit)
			      in
				(stm, IL.Stmt.tail stm, env, assigned)
			      end
			  | _ => let
			      val (next, last, env, assigned) = toStmt (env, assigned, rest)
			      val stm = IL.Stmt.mkNEW(name, xs', SOME next)
			      in
				IL.Node.addEdge (IL.Stmt.tail stm, IL.Stmt.entry next);
				(stm, last, env, assigned)
			      end
		      end
		  | S.S_Die => let
		      val stm = IL.Stmt.mkDIE()
		      in
			(stm, IL.Stmt.tail stm, env, assigned)
		      end
		  | S.S_Stabilize => let
		      val stm = IL.Stmt.mkSTABILIZE()
		      in
			(stm, IL.Stmt.tail stm, env, assigned)
		      end
		(* end case *))
	  in
	    toStmt (env, VSet.empty, stms)
	  end
*)

    fun cvtTopLevelBlock (env, blk) = let
	  fun finish (env, firstNd, lastNd) = let
		val entry = IL.Node.mkENTRY ()
		val exit = IL.Node.mkEXIT ()
		in
		  IL.Node.addEdge (entry, firstNd);
		(* NOTE: this addEdge could fail if all control paths end in DIE or STABILIZE,
		 * so we wrap it in a handler
		 *)
		  IL.Node.addEdge (lastNd, exit) handle _ => ();
		  IL.CFG{entry = ref entry, exit = ref exit}
		end
	  in
	    cvtBlock (env, blk, finish)
	  end

  (* generate fresh SSA variables and add them to the environment *)
    fun freshVars (env, xs) = let
	  fun cvtVar (x, (env, xs)) = let
		val x' = newVar x
		in
		  (VMap.insert(env, x, x'), x'::xs)
		end
	  val (env, xs) = List.foldl cvtVar (env, []) xs
	  in
	    (env, List.rev xs)
	  end

    fun translate (S.Program{globals, globalInit, actors}) = let
	  val (globalInit, env) = cvtTopLevelBlock (VMap.empty, globalInit)
	(* get the SSA names for the globals and a reduced environment that just defines
	 * the globals.
	 *)
	  val (env, globs) = let
		val lookup = lookup env
		fun cvtVar (x, (env, globs)) = let
		      val x' = lookup x
		      in
			(VMap.insert(env, x, x'), x'::globs)
		      end
		val (env, globs) = List.foldl cvtVar (VMap.empty, []) globals
		in
		  (env, List.rev globs)
		end
	  fun cvtActor (S.Actor{name, params, state, stateInit, methods}) = let
		val (env, params) = let
		      fun cvtParam (x, (env, xs)) = let
			    val x' = newVar x
			    in
			      (VMap.insert(env, x, x'), x'::xs)
			    end
		      val (env, params) = List.foldl cvtParam (env, []) params
		      in
			(env, List.rev params)
		      end
		val (stateInit, env) = cvtTopLevelBlock (env, stateInit)
		val state' = List.map (lookup env) state
		fun cvtMethod (S.Method(name, blk)) = let
		    (* allocate fresh variables for the state variables *)
		      val (env, stateIn) = freshVars (env, state)
		      val (body, env) = cvtTopLevelBlock (env, blk)
		      val stateOut = List.map (lookup env) state
		      in
			IL.Method{name=name, stateIn=stateIn, stateOut=stateOut, body=body}
		      end
		in
		  IL.Actor{
		      name = name,
		      params = params,
		      state = state',
		      stateInit = stateInit,
		      methods = List.map cvtMethod methods
		    }
		end
	  in
	    IL.Program{
		globals = globs,
		globalInit = globalInit,
		actors = List.map cvtActor actors
	      }
	  end

  end