(* codegen-fn.sml
 *
 * COPYRIGHT (c) 2010 The Diderot Project (http://diderot-language.cs.uchicago.edu)
 * All rights reserved.
 *
 * Generic support for translating LowIL code to the target representation.  We
 * assume that the LowIL has first been run through the splitting pass to match
 * the target's vector widths.
 *)

functor CodeGenFn (T : TARGET) : sig

    val generate : string * LowIL.program -> unit

  end = struct

    structure IL = TreeIL
    structure Ty = IL.Ty
    structure Op = IL.Op
    structure V = IL.Var

  (* convert LowIL types to T types *)
    fun cvtTy ty = (case ty
	   of Ty.BoolTy => T.boolTy
	    | Ty.StringTy => T.stringTy
	    | Ty.IVecTy 1 => T.intTy
	    | Ty.IVecTy n => T.ivecTy n		(* FIXME: what about vector splits? *)
	    | Ty.VecTy 1 => T.realTy
	    | Ty.VecTy n => T.vecTy n		(* FIXME: what about vector splits? *)
	    | Ty.AddrTy info => T.imageDataTy info
	    | Ty.ImageTy info => T.imageTy info
	  (* end case *))

    fun addBindings (env, xs, ys) =
	  ListPair.foldlEq (fn (x, y, env) => V.Map.insert(env, x, y)) env (xs, ys)

    fun lookup (env, x) = (case V.Map.find (env, x)
	   of SOME x' => x'
	    | NONE => raise Fail(concat["lookup(_, ", V.name x, ")"])
	  (* end case *))

    fun trExp (env, e) = (case e
	   of IL.E_Var x => (case V.kind x
		 of IL.VK_Global => T.Expr.global(lookup(env, x))
		  | IL.VK_State strand => T.Expr.getState(lookup(env, x))
		  | IL.VK_Local => T.Expr.var(lookup(env, x))
		(* end case *))
	    | IL.E_Lit(Literal.Int n) => T.Expr.intLit n
	    | IL.E_Lit(Literal.Bool b) => T.Expr.boolLit b
	    | IL.E_Lit(Literal.Float f) => T.Expr.floatLit f
	    | IL.E_Lit(Literal.String s) => T.Expr.stringLit s
	    | IL.E_Op(rator, args) => (case (rator, trExps(env, args))
		 of (Op.Add ty, [a, b])		=> T.Expr.add(a, b)
		  | (Op.Sub ty, [a, b])		=> T.Expr.sub(a, b)
		  | (Op.Mul ty, [a, b])		=> T.Expr.mul(a, b)
		  | (Op.Div ty, [a, b])		=> T.Expr.divide(a, b)
		  | (Op.Neg ty, [a])		=> T.Expr.neg a
		  | (Op.LT ty, [a, b])		=> T.Expr.lt(a, b)
		  | (Op.LTE ty, [a, b])		=> T.Expr.lte(a, b)
		  | (Op.EQ ty, [a, b])		=> T.Expr.equ(a, b)
		  | (Op.NEQ ty, [a, b])		=> T.Expr.neq(a, b)
		  | (Op.GT ty, [a, b])		=> T.Expr.gt(a, b)
		  | (Op.GTE ty, [a, b])		=> T.Expr.gte(a, b)
		  | (Op.Not, [a])		=> T.Expr.not a
		  | (Op.Max, [a, b])		=> T.Expr.max(a, b)
		  | (Op.Min, [a, b])		=> T.Expr.min(a, b)
		  | (Op.Sin, [a])		=> T.Expr.sin a
		  | (Op.Cos, [a])		=> T.Expr.cos a
		  | (Op.Pow, [a, b])		=> T.Expr.pow(a, b)
		  | (Op.Dot d, [a, b])		=> T.Expr.dot(a, b)
		  | (Op.Cross, [a, b])		=> T.Expr.cross(a, b)
		  | (Op.Select(ty, i), [a])	=> T.Expr.select(i, a)
		  | (Op.Norm d, [a])		=> T.Expr.length a
		  | (Op.Scale d, [a, b])	=> T.Expr.mul(a, b)
		  | (Op.InvScale d, [a, b])	=> T.Expr.divide(a, b)
		  | (Op.CL, _)			=> raise Fail "CL unimplemented"
		  | (Op.PrincipleEvec ty, _)	=> raise Fail "PrincipleEvec unimplemented"
(*
		  | (Op.Subscript ty,
*)
		  | (Op.Ceiling d, [a])		=> T.Expr.ceil a
		  | (Op.Floor d, [a])		=> T.Expr.floor a
		  | (Op.Round d, [a])		=> T.Expr.round a
		  | (Op.Trunc d, [a])		=> T.Expr.trunc a
		  | (Op.IntToReal, [a])		=> T.Expr.toReal a
		  | (Op.RealToInt d, [a])	=> T.Expr.toInt a
		  | (Op.ImageAddress d, [a])	=> T.Expr.imageAddr a
		  | (Op.LoadVoxels(info, 1), [a]) => T.Expr.getImgData a
		  | (Op.LoadVoxels _, [a])	=> raise Fail "impossible"
		  | (Op.PosToImgSpace d, [v, x]) => T.Expr.posToImgSpace(v, x)
		  | (Op.GradToWorldSpace d, [v, x]) => T.Expr.intLit 0 (* FIXME *)
		  | (Op.LoadImage info, [a])	=> raise Fail "impossible"
		  | (Op.Inside(d, s), [x, v])	=> T.Expr.inside(x, v, s)
		  | (Op.Input(ty, name), [])	=> raise Fail "impossible"
		  | (Op.InputWithDefault(ty, name), [a]) => T.Expr.intLit 0 (* FIXME *)
		  | _ => raise Fail(concat[
			"incorrect number of arguments for ", Op.toString rator
		      ])
		(* end case *))
	  (* end case *))

    and trExps (env, exps) = List.map (fn exp => trExp(env, exp)) exps

    fun trBody (stateVars, env, body) = let
	  fun trStmt (env, stm) = (case stm
		 of IL.S_Comment text => [T.Stmt.comment text]
		  | IL.S_Assign(x, exp) => (case V.kind x
		       of IL.VK_Global => [T.Stmt.assign(lookup(env, x), trExp(env, exp))]
			| IL.VK_State strand =>
			    [T.Stmt.assignState(lookup(env, x), trExp(env, exp))]
			| IL.VK_Local => [T.Stmt.assign(lookup(env, x), trExp(env, exp))]
		      (* end case *))
		  | IL.S_IfThen(cond, thenBlk) =>
		      [T.Stmt.ifthen(trExp(env, cond), trBlock(env, thenBlk))]
		  | IL.S_IfThenElse(cond, thenBlk, elseBlk) =>
		      [T.Stmt.ifthenelse(trExp(env, cond),
			trBlock(env, thenBlk),
			trBlock(env, elseBlk))]
		  | IL.S_For(x, e1, e2, blk) => let
		      val x' = T.Var.var(cvtTy(V.ty x), V.name x)
		      val env' = V.Map.insert(env, x, x')
		      in [
			T.Stmt.for(x', trExp(env, e1), trExp(env, e2), trBlock(env', blk))
		      ] end
		  | IL.S_Cons(lhs, args) =>
		      [T.Stmt.cons(lookup(env, lhs), trExps(env, args))]
		  | IL.S_LoadVoxels(lhs, dim, addr) =>
		      T.Stmt.getImgData(lookup(env, lhs), trExp(env, addr))
		  | IL.S_LoadImage(lhs, dim, name) =>
		      T.Stmt.loadImage (lookup(env, lhs), dim, trExp(env, name))
		  | IL.S_Input(lhs, name, optDflt) =>
		      T.Stmt.input(lookup(env, lhs), name, Option.map (fn e => trExp(env, e)) optDflt)
(* FIXME: what about the args? *)
		  | IL.S_Exit args => [T.Stmt.exit()]
		  | IL.S_Active args =>
		      ListPair.foldrEq
			(fn (x, e, stms) => T.Stmt.assignState(x, trExp(env, e))::stms)
			  [T.Stmt.active()]
			    (stateVars, args)
		  | IL.S_Stabilize args =>
		      ListPair.foldrEq
			(fn (x, e, stms) => T.Stmt.assignState(x, trExp(env, e))::stms)
			  [T.Stmt.stabilize()]
			    (stateVars, args)
		  | IL.S_Die => [T.Stmt.die()]
		(* end case *))
	  and trBlock (env, IL.Block{locals, body}) = let
		val env = List.foldl
		      (fn (x, env) => V.Map.insert(env, x, T.Var.var(cvtTy(V.ty x), V.name x)))
			env locals
		val stms = List.foldr (fn (stm, stms) => trStmt(env, stm)@stms) [] body
		val stms = List.foldr
		      (fn (x, stms) => T.Stmt.decl(lookup(env, x), NONE)::stms)
			stms locals
		in
		  T.Stmt.block stms
		end
	  in
	    trBlock (env, body)
	  end

    fun trMethod (strand, state, env) (IL.Method{name, body}) =
	  T.Strand.method (strand, Atom.toString name, trBody (state, env, body))

    fun trStrand (prog, env) (IL.Strand{name, params, state, stateInit, methods}) = let
	  val strand = T.Strand.define(prog, Atom.toString name)
	  val state' =
		List.map (fn x => T.Var.state(strand, cvtTy(V.ty x), V.name x)) state
	  val env = addBindings (env, state, state')
	(* define the parameters and add them to the environment *)
	  val params' = List.map (fn x => T.Var.param(cvtTy(V.ty x), V.name x)) params
	  val env = addBindings (env, params, params')
	  in
	    T.Strand.init (strand, params', trBody ([], env, stateInit));
	    List.app (trMethod (strand, state', env)) methods
	  end

    fun generate (fileStem, srcProg) = let
	  val treeProg as TreeIL.Program{globals, globalInit, strands} = LowToTree.translate srcProg
val _ = (
TextIO.output(Log.logFile(), "********** After translation to TreeIL **********\n");
TreeILPP.program (Log.logFile(), treeProg))
	  val prog = T.newProgram ()
	(* define the globals and initialize the environment *)
	  val env = let
		fun gvar (x, env) =
		      V.Map.insert(env, x, T.Var.global(prog, cvtTy(V.ty x), V.name x))
		in
		  List.foldl gvar V.Map.empty globals
		end
	  in
	  (* global initialization *)
	    T.globalInit (prog, trBody ([], env, globalInit));
	  (* translate strands *)
	    List.app (trStrand (prog, env)) strands;
	  (* output the program *)
	    T.generate (fileStem, prog)
	  end

  end