(* low-to-tree.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).
 *)

structure LowToTree : 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

  (* 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 (V.name x, T.VK_Global, V.ty x)
    fun newLocal x = newVar (genName(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
	}
      fun peekGlobal (E{tbl, ...}, x) = (case VT.find tbl x
	     of SOME(GLOB x') => SOME x'
	      | _ => NONE
	    (* end case *))
    in
    fun newEnv () = E{tbl = VT.mkTable (512, Fail "tbl"), locals=[]}

    fun newScope (E{tbl, ...}) = E{tbl=tbl, locals=[]}

  (* use a variable.  If it is a pending expression, we decrement its use count *)
    fun useVar (E{tbl, ...}) x = (case VT.find tbl x
	   of SOME(GLOB x') => T.E_Var x'
	    | SOME(TREE e) => (
		if (decCount x) then ignore(VT.remove tbl x) else ();
		e)
	    | SOME(DEF e) => e
	    | NONE => raise Fail "useVar"
	  (* 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')

    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 bind (env, lhs, rhs) = (case peekGlobal (env, lhs)
	   of SOME x => (rename(env, lhs, x), [T.S_Assign(x, rhs)])
	    | NONE => if (V.useCount lhs = 1)
		then (insert(env, lhs, rhs), [])
		else let
		  val t = newLocal lhs
		  in
		    (rename(env, lhs, t), [T.S_Assign(t, rhs)])
		  end
	  (* end case *))

    fun setDef (env, lhs, rhs) = (case peekGlobal (env, lhs)
	   of SOME x => (rename(env, lhs, x), [T.S_Assign(x, rhs)])
	    | NONE => (insert(env, lhs, rhs), [])
	  (* 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 endBlock (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
	  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
	  val t = newLocal lhs
	  val predBlks = ListPair.map
		(fn (x, stms) => T.S_Assign(t, useVar env x)::stms)
		  (rhs, predBlks)
	  in
	    (rename (env, lhs, t), predBlks)
	  end

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

    end

  (* translate a LowIL assignment to a list of zero or more target statements *)
    fun doAssign (env, (lhs, rhs)) = (case rhs
	   of IL.VAR x => setDef (env, lhs, useVar env x)
	    | IL.LIT lit => setDef (env, lhs, T.E_Lit lit)
	    | IL.OP(Op.LoadImage info, [a]) => let
		val t = newLocal lhs
		in
		  (rename(env, lhs, t), [T.S_LoadImage(t, ImageInfo.dim info, useVar env a)])
		end
	    | IL.OP(Op.Input(ty, name), []) => let
		val t = newLocal lhs
		in
		  (rename(env, lhs, t), [T.S_Input(t, name, NONE)])
		end
	    | IL.OP(Op.InputWithDefault(ty, name), [a]) => let
		val t = newLocal lhs
		in
		  (rename(env, lhs, t), [T.S_Input(t, name, SOME(useVar env a))])
		end
	    | IL.OP(rator, args) =>
		bind (env, lhs, T.E_Op(rator, List.map (useVar env) args))
	    | IL.CONS args => let
	      (* we give cons expressions names, since not all targets support them inline *)
		val t = newLocal lhs
		val rhs = List.map (useVar env) args
		in
		  (rename(env, lhs, t), [T.S_Cons(t, rhs)])
		end
	  (* end case *))

    fun gen (env, cfg) = let
	  fun doNode (env, ifCont : T.stm list * IL.node_kind -> T.block, stms, nd) = (
		case Nd.kind nd
		 of IL.NULL => raise Fail "unexpected NULL"
		  | IL.ENTRY{succ} => doNode (env, ifCont, stms, !succ)
		  | IL.JOIN{phis, succ, ...} => ifCont (stms, Nd.kind nd)
		  | IL.COND{cond, trueBranch, falseBranch, ...} => let
		      fun kThen (stms', _) = let
			    val (env, thenBlk) = endBlock (env, stms')
			    fun kElse (stms', IL.JOIN{phis, succ, ...}) = let
				  val (env, elseBlk) = endBlock (env, stms')
				  val (env, [thenBlk, elseBlk]) =
					List.foldl doPhi (env, [thenBlk, elseBlk]) (!phis)
				  val stm = mkIf(useVar env cond, List.rev thenBlk, List.rev elseBlk)
				  in
				    doNode (env, ifCont, stm::stms, !succ)
				  end
			    in
			      doNode (env, kElse, [], !falseBranch)
			    end
		      in
			doNode (env, kThen, [], !trueBranch)
		      end
		  | IL.COM {text, succ, ...} =>
		      doNode (env, ifCont, T.S_Comment text :: stms, !succ)
		  | IL.ASSIGN{stm, succ, ...} => let
		      val (env, stms') = doAssign (env, stm)
		      in
			doNode (env, ifCont, stms' @ stms, !succ)
		      end
		  | IL.NEW{strand, args, succ, ...} => raise Fail "NEW unimplemented"
		  | IL.DIE _ =>
		      mkBlock (List.rev (T.S_Die :: stms))
		  | IL.STABILIZE _ =>
		      mkBlock (List.rev stms)
		  | IL.EXIT _ => endScope (env, List.rev (T.S_Stabilize :: stms))
		(* end case *))
	  in
	    doNode (env, fn _ => raise Fail "bogus ifCont at JOIN node", [], CFG.entry cfg)
	  end

    fun translate (IL.Program{globals, globalInit, strands}) = let
	  val env = newEnv()
	  val globals = List.map
		(fn x => let val x' = newGlobal x in global(env, x, x'); x' end)
		  globals
	  in
	    T.Program{
		globals = globals,
		globalInit = gen (env, globalInit),
		strands = [] (* FIXME *)
	      }
	  end

  end