(* low-to-tree-fn.sml
 *
 * COPYRIGHT (c) 2011 The Diderot Project (http://diderot-language.cs.uchicago.edu)
 * All rights reserved.
 *
 * This module translates the LowIL representation of a program (i.e., a pure CFG) to
 * a block-structured AST with nested expressions.
 *
 * NOTE: this translation is pretty dumb about variable coalescing (i.e., it doesn't do any).
 *)

functor LowToTreeFn (Target : sig

  (* tests for whether various expression forms can appear inline *)
    val inlineCons : int -> bool	(* can n'th-order tensor construction appear inline *)
    val inlineMatrixExp : bool		(* can matrix-valued expressions appear inline? *)

  end) : sig

    val translate : LowIL.program -> TreeIL.program

  end = struct

    structure IL = LowIL
    structure Ty = LowILTypes
    structure V = LowIL.Var
    structure Op = LowOps
    structure Nd = LowIL.Node
    structure CFG = LowIL.CFG
    structure T = TreeIL
    structure VA = VarAnalysis

  (* create new tree IL variables *)
    local
      fun newVar (name, kind, ty) = T.V{
	      name = name,
	      id = Stamp.new(),
	      kind = kind,
	      ty = ty
	    }
      val cnt = ref 0
      fun genName prefix = let
	    val n = !cnt
	    in
	      cnt := n+1;
	      String.concat[prefix, "_", Int.toString n]
	    end
    in
    fun newGlobal x = newVar ("G_" ^ V.name x, T.VK_Global, V.ty x)
    fun newStateVar (strand, x) =
	  newVar (concat[Atom.toString strand, "_", V.name x], T.VK_State strand, V.ty x)
    fun newParam x = newVar (genName("p_" ^ V.name x), T.VK_Local, V.ty x)
    fun newLocal x = newVar (genName("l_" ^ V.name x), T.VK_Local, V.ty x)
    fun newIter x = newVar (genName("i_" ^ V.name x), T.VK_Local, V.ty x)
    end

    fun mkBlock stms = T.Block{locals=[], body=stms}
    fun mkIf (x, stms, []) = T.S_IfThen(x, mkBlock stms)
      | mkIf (x, stms1, stms2) = T.S_IfThenElse(x, mkBlock stms1, mkBlock stms2)

  (* an environment that tracks bindings of variables to target expressions and the list
   * of locals that have been defined.
   *)
    local
      structure VT = V.Tbl
      fun decCount (IL.V{useCnt, ...}) = let
	    val n = !useCnt - 1
	    in
	      useCnt := n;  (n <= 0)
	    end
      datatype target_binding
	= GLOB of T.var		(* variable is global *)
	| TREE of T.exp		(* variable bound to target expression tree *)
	| DEF of T.exp		(* either a target variable or constant for a defined variable *)
      datatype env = E of {
	  tbl : target_binding VT.hash_table,
	  locals : T.var list
	}
    in
    fun newEnv () = E{tbl = VT.mkTable (512, Fail "tbl"), locals=[]}

  (* use a variable.  If it is a pending expression, we remove it from the table *)
    fun useVar (E{tbl, ...}) x = (case VT.find tbl x
	   of SOME(GLOB x') => T.E_Var x'
	    | SOME(TREE e) => (
(*print(concat["useVar ", V.toString x, " ==> TREE\n"]);*)
		ignore(VT.remove tbl x);
		e)
	    | SOME(DEF e) => (
(*print(concat["useVar ", V.toString x, " ==> DEF; use count = ", Int.toString(V.useCount x), "\n"]);*)
	      (* if this is the last use of x, then remove it from the table *)
		if (decCount x) then ignore(VT.remove tbl x) else ();
		e)
	    | NONE => (
print "*** dump environment\n";
VT.appi (fn (x, _) => print(concat["  ", IL.Var.toString x, "\n"])) tbl;
print "***\n";
raise Fail(concat ["useVar(", V.toString x, ")"])
)
	  (* end case *))

  (* record a local variable *)
    fun addLocal (E{tbl, locals}, x) = E{tbl=tbl, locals=x::locals}

    fun global (E{tbl, ...}, x, x') = VT.insert tbl (x, GLOB x')

  (* insert a pending expression into the table.  Note that x should only be used once! *)
    fun insert (env as E{tbl, ...}, x, exp) = (
	  VT.insert tbl (x, TREE exp);
	  env)

    fun rename (env as E{tbl, ...}, x, x') = (
	  VT.insert tbl (x, DEF(T.E_Var x'));
	  env)

    fun peekGlobal (E{tbl, ...}, x) = (case VT.find tbl x
	   of SOME(GLOB x') => SOME x'
	    | _ => NONE
	  (* end case *))

    fun bindLocal (env, lhs, rhs) = if (V.useCount lhs = 1)
	  then (insert(env, lhs, rhs), [])
	  else let
	    val t = newLocal lhs
	    in
	      (rename(addLocal(env, t), lhs, t), [T.S_Assign(t, rhs)])
	    end

    fun bind (env, lhs, rhs) = (case peekGlobal (env, lhs)
	   of SOME x => (env, [T.S_Assign(x, rhs)])
	    | NONE => bindLocal (env, lhs, rhs)
	  (* end case *))

  (* set the definition of a variable, where the RHS is either a literal constant or a variable *)
    fun bindSimple (env as E{tbl, ...}, lhs, rhs) = (
	  case peekGlobal (env, lhs)
	   of SOME x => (env, [T.S_Assign(x, rhs)])
	    | NONE => (VT.insert tbl (lhs, DEF rhs); (env, []))
	  (* end case *))

  (* at the end of a block, we need to assign any pending expressions to locals.  The
   * blkStms list and the resulting statement list are in reverse order.
   *)
    fun flushPending (E{tbl, locals}, blkStms) = let
	  fun doVar (x, TREE e, (locals, stms)) = let
		val t = newLocal x
		in
		  VT.insert tbl (x, DEF(T.E_Var t));
		  (t::locals, T.S_Assign(t, e)::stms)
		end
	    | doVar (_, _, acc) = acc
	  val (locals, stms) = VT.foldi doVar (locals, blkStms) tbl
	  in
	    (E{tbl=tbl, locals=locals}, stms)
	  end

    fun doPhi ((lhs, rhs), (env, predBlks : T.stm list list)) = let
	(* t will be the variable in the continuation of the JOIN *)
	  val t = newLocal lhs
	  val predBlks = ListPair.map
		(fn (x, stms) => T.S_Assign(t, useVar env x)::stms)
		  (rhs, predBlks)
	  in
	    (rename (addLocal(env, t), lhs, t), predBlks)
	  end

    fun endScope (E{locals, ...}, stms) = T.Block{
	    locals = List.rev locals,
	    body = stms
	  }

    end

  (* Certain IL operators cannot be compiled to inline expressions.  Return
   * false for those and true for all others.
   *)
    fun isInlineOp rator = let
	  fun chkTensorTy (Ty.TensorTy[]) = true
	    | chkTensorTy (Ty.TensorTy[_]) = true
	    | chkTensorTy (Ty.TensorTy _) = Target.inlineMatrixExp
	    | chkTensorTy _ = true
	  in
	   case rator
	     of Op.LoadVoxels(_, 1) => true
	      | Op.LoadVoxels _ => false
	      | Op.Add ty => chkTensorTy ty
	      | Op.Sub ty => chkTensorTy ty
	      | Op.Neg ty => chkTensorTy ty
	      | Op.Scale ty => chkTensorTy ty
	      | Op.MulMatMat _ => Target.inlineMatrixExp
	      | Op.Identity _ => Target.inlineMatrixExp
	      | Op.Zero _ => Target.inlineMatrixExp
	      | Op.TensorToWorldSpace(_, ty) => chkTensorTy ty
	      | _ => true
	    (* end case *)
	  end

  (* translate a LowIL assignment to a list of zero or more target statements *)
    fun doAssign (env, (lhs, rhs)) = let
	  fun doLHS () = (case peekGlobal(env, lhs)
		 of SOME lhs' => (env, lhs')
		  | NONE => let
		      val t = newLocal lhs
		      in
			(rename (addLocal(env, t), lhs, t), t)
		      end
		(* end case *))
	(* for expressions that are going to be compiled to a call statement *)
	  fun assignExp (env, exp) = let
	      (* operations that return matrices may not be supported inline *)
		val (env, t) = doLHS()
		in
		  (env, [T.S_Assign(t, exp)])
		end
	  in
	    case rhs
	     of IL.VAR x => bindSimple (env, lhs, useVar env x)
	      | IL.LIT lit => bindSimple (env, lhs, T.E_Lit lit)
	      | IL.OP(Op.LoadImage info, [a]) => let
		  val (env, t) = doLHS()
		  in
		    (env, [T.S_LoadImage(t, ImageInfo.dim info, useVar env a)])
		  end
	      | IL.OP(Op.Input(ty, name), []) => let
		  val (env, t) = doLHS()
		  in
		    (env, [T.S_Input(t, name, NONE)])
		  end
	      | IL.OP(Op.InputWithDefault(ty, name), [a]) => let
		  val (env, t) = doLHS()
		  in
		    (env, [T.S_Input(t, name, SOME(useVar env a))])
		  end
	      | IL.OP(rator, args) => let
		  val exp = T.E_Op(rator, List.map (useVar env) args)
		  in
		    if isInlineOp rator
		      then bind (env, lhs, exp)
		      else assignExp (env, exp)
		  end
	      | IL.APPLY(f, args) =>
		  bind (env, lhs, T.E_Apply(f, List.map (useVar env) args))
	      | IL.CONS(ty, args) => let
		  val inline = (case ty
			 of Ty.IVecTy _ => true
			  | Ty.TensorTy dd => Target.inlineCons(List.length dd)
			(* end case *))
		  val exp = T.E_Cons(ty, List.map (useVar env) args)
		  in
		    if inline
		      then bind (env, lhs, exp)
		      else assignExp (env, exp)
		  end
	    (* end case *)
	  end

  (* In order to reconstruct the block-structure from the CFG, we keep a stack of open ifs.
   * the items on this stack distinguish between when we are processing the then and else
   * branches of the if.
   *)
    datatype open_if
    (* working on the "then" branch.  The fields are statments that preceed the if, the condition,
     * and the else-branch node.
     *)
      = THEN_BR of T.stm list * T.exp * IL.node
    (* working on the "else" branch.  The fields are statments that preceed the if, the condition,
     * the "then" branch statements, and the node that terminated the "then" branch (will be
     * a JOIN, DIE, or STABILIZE).
     *)
      | ELSE_BR of T.stm list * T.exp * T.stm list * IL.node_kind

    fun trCFG (env, prefix, finish, cfg) = let
	  fun join (env, [], _, IL.JOIN _) = raise Fail "JOIN with no open if"
	    | join (env, [], stms, _) = endScope (env, prefix @ List.rev stms)
	    | join (env, THEN_BR(stms1, cond, elseBr)::stk, thenBlk, k) = let
		val (env, thenBlk) = flushPending (env, thenBlk)
		in
		  doNode (env, ELSE_BR(stms1, cond, thenBlk, k)::stk, [], elseBr)
		end
	    | join (env, ELSE_BR(stms, cond, thenBlk, k1)::stk, elseBlk, k2) = let
		val (env, elseBlk) = flushPending (env, elseBlk)
		in
		  case (k1, k2)
		   of (IL.JOIN{phis, succ, ...}, IL.JOIN _) => let
			val (env, [thenBlk, elseBlk]) =
			      List.foldl doPhi (env, [thenBlk, elseBlk]) (!phis)
			val stm = mkIf(cond, List.rev thenBlk, List.rev elseBlk)
			in
			  doNode (env, stk, stm::stms, !succ)
			end
		    | (IL.JOIN{phis, succ, ...}, _) => let
			val (env, [thenBlk]) = List.foldl doPhi (env, [thenBlk]) (!phis)
			val stm = mkIf(cond, List.rev thenBlk, List.rev elseBlk)
			in
			  doNode (env, stk, stm::stms, !succ)
			end
		    | (_, IL.JOIN{phis, succ, ...}) => let
			val (env, [elseBlk]) = List.foldl doPhi (env, [elseBlk]) (!phis)
			val stm = mkIf(cond, List.rev thenBlk, List.rev elseBlk)
			in
			  doNode (env, stk, stm::stms, !succ)
			end
		    | (_, _) => raise Fail "no path to exit unimplemented" (* FIXME *)
		  (* end case *)
		end
	  and doNode (env, ifStk : open_if list, stms, nd) = (
		case Nd.kind nd
		 of IL.NULL => raise Fail "unexpected NULL"
		  | IL.ENTRY{succ} => doNode (env, ifStk, stms, !succ)
		  | k as IL.JOIN{phis, succ, ...} => join (env, ifStk, stms, k)
		  | IL.COND{cond, trueBranch, falseBranch, ...} => let
		      val cond = useVar env cond
		      val (env, stms) = flushPending (env, stms)
		      in
			doNode (env, THEN_BR(stms, cond, !falseBranch)::ifStk, [], !trueBranch)
		      end
		  | IL.COM {text, succ, ...} =>
		      doNode (env, ifStk, T.S_Comment text :: stms, !succ)
		  | IL.ASSIGN{stm, succ, ...} => let
		      val (env, stms') = doAssign (env, stm)
		      in
			doNode (env, ifStk, stms' @ stms, !succ)
		      end
		  | IL.NEW{strand, args, succ, ...} => raise Fail "NEW unimplemented"
		  | k as IL.EXIT{kind, live, ...} => (case kind
		       of ExitKind.FRAGMENT =>
			    endScope (env, prefix @ List.revAppend(stms, finish env))
			| ExitKind.RETURN => let
			    val suffix = finish env @ [T.S_Exit(List.map (useVar env) live)]
			    in
			      endScope (env, prefix @ List.revAppend(stms, suffix))
			    end
			| ExitKind.ACTIVE => let
			    val suffix = finish env @ [T.S_Active(List.map (useVar env) live)]
			    in
			      endScope (env, prefix @ List.revAppend(stms, suffix))
			    end
			| ExitKind.STABILIZE => let
			    val stms = T.S_Stabilize(List.map (useVar env) live) :: stms
			    in
			      join (env, ifStk, stms, k)
			    end
			| ExitKind.DIE => join (env, ifStk, T.S_Die :: stms, k)
		      (* end case *))
		(* end case *))
	  in
	    doNode (env, [], [], CFG.entry cfg)
	  end

    fun trInitially (env, IL.Initially{isArray, rangeInit, iters, create=(createInit, strand, args)}) =
	  let
	  val iterPrefix = trCFG (env, [], fn _ => [], rangeInit)
	  fun cvtIter ((param, lo, hi), (env, iters)) = let
		val param' = newIter param
		val env = rename (env, param, param')
		in
		  (env, (param', useVar env lo, useVar env hi)::iters)
		end
	  val (env, iters) = List.foldr cvtIter (env, []) iters
	  val createPrefix = trCFG (env, [], fn _ => [], createInit)
	  in {
	    isArray = isArray,
	    iterPrefix = iterPrefix,
	    iters = iters,
	    createPrefix = createPrefix,
	    strand = strand,
	    args = List.map (useVar env) args
	  } end

    fun trMethod (env, stateVars) (IL.Method{name, stateIn, body}) = let
	  fun bindStateVar (x, T.SV{var, ...}, (env, stms)) = let
		val (env, stms') = bindLocal(env, x, T.E_Var var)
		in
		  (env, stms' @ stms)
		end
	  val (env, stms) = ListPair.foldrEq bindStateVar (env, []) (stateIn, stateVars)
	  in
	    T.Method{name = name, body = trCFG (env, stms, fn _ => [], body)}
	  end

    fun trStrand env (IL.Strand{name, params, state, stateInit, methods}) = let
	  val params' = List.map newParam params
	  val env = ListPair.foldlEq (fn (x, x', env) => rename(env, x, x')) env (params, params')
	  val stateVars = let
		fun cvtVar (isOut, x) = T.SV{
			varying = (case VA.varScope x
			   of VA.StrandConstState => false
			    | VA.StrandState => true
			    | s => raise Fail(concat[
				  "state variable ", IL.Var.toString x,
				  " has bogus scope annotation ", VA.scopeToString s
				])
			  (* end case *)),
			output = isOut, var = newStateVar(name, x)
		      }
		in
		  List.map cvtVar state
		end
	  in
	    T.Strand{
		name = name,
		params = params',
		state = stateVars,
		stateInit = trCFG (env, [], fn _ => [], stateInit),
		methods = List.map (trMethod(env, stateVars)) methods
	      }
	  end

    fun translate prog = let
	(* first we do a variable analysis pass on the Low IL *)
	  val prog as IL.Program{globalInit, initially, strands} = VA.optimize prog
	  val _ = ( (* DEBUG *)
		LowPP.output (Log.logFile(), "LowIL after variable analysis", prog);
		if CheckLowIL.check ("after LowIL variable analysis", prog)
		  then raise Fail "bogus Low IL after variable analysis"
		  else ())
	  val env = newEnv()
	  val globals = List.map
		(fn x => let val x' = newGlobal x in global(env, x, x'); x' end)
		  (IL.CFG.liveAtExit globalInit)
	  in
	    T.Program{
		globals = globals,
		globalInit = trCFG (env, [], fn _ => [], globalInit),
		strands = List.map (trStrand env) strands,
		initially = trInitially (env, initially)
	      }
	  end

  end