(* mid-to-low.sml
 *
 * COPYRIGHT (c) 2010 The Diderot Project (http://diderot-language.cs.uchicago.edu)
 * All rights reserved.
 *
 * Translation from MidIL to LowIL representations.
 *)

structure MidToLow : sig

    val translate : MidIL.program -> LowIL.program

  end = struct

    structure SrcIL = MidIL
    structure SrcOp = MidOps
    structure VTbl = SrcIL.Var.Tbl
    structure DstIL = LowIL
    structure DstTy = LowILTypes
    structure DstOp = LowOps

    type var_env = DstIL.var VTbl.hash_table

    fun rename (env : var_env, x) = (case VTbl.find env x
	   of SOME x' => x'
	    | NONE => let
		val x' = DstIL.Var.new (SrcIL.Var.name x, SrcIL.Var.ty x)
		in
		  VTbl.insert env (x, x');
		  x'
		end
	  (* end case *))
    fun renameList (env, xs) = List.map (fn x => rename(env, x)) xs

  (* convert a rational to a FloatLit.float value.  We do this by long division
   * with a cutoff when we get to 12 digits.
   *)
    fun ratToFloat r = (case Rational.explode r
	   of {sign=0, ...} => FloatLit.zero false
	    | {sign, num, denom=1} => FloatLit.fromInt(sign * IntInf.toInt num)
	    | {sign, num, denom} => let
	      (* normalize so that num <= denom *)
		val (denom, exp) = let
		      fun lp (n, denom) = if (denom < num)
			    then lp(n+1, denom*10)
			    else (denom, n)
		      in
			lp (1, denom)
		      end
	      (* normalize so that num <= denom < 10*num *)
		val (num, exp) = let
		      fun lp (n, num) = if (10*num < denom)
			    then lp(n-1, 10*num)
			    else (num, n)
		      in
			lp (exp, num)
		      end
	      (* divide num/denom, computing the resulting digits *)
		fun divLp (n, a) = let
		      val (q, r) = IntInf.divMod(a, denom)
		      in
			if (r = 0) then (q, [])
			else if (n < 12) then let
			  val (d, dd) = divLp(n+1, 10*r)
			  in
			    if (d < 10)
			      then (q, (IntInf.toInt d)::dd)
			      else (q+1, 0::dd)
			  end
			else if (IntInf.div(10*r, denom) < 5)
			  then (q, [])
			  else (q+1, []) (* round up *)
		      end
		val digits = let
		      val (d, dd) = divLp (0, num)
		      in
			(IntInf.toInt d)::dd
		      end
		in
		  FloatLit.fromDigits{isNeg=(sign < 0), digits=digits, exp=exp}
		end
	  (* end case *))

    fun imul (r : DstIL.var, a, b) = (r, DstIL.OP(DstOp.Mul DstTy.intTy, [a, b]))
    fun iadd (r : DstIL.var, a, b) = (r, DstIL.OP(DstOp.Add DstTy.intTy, [a, b]))
    fun ilit (r : DstIL.var, n) = (r, DstIL.LIT(Literal.Int(IntInf.fromInt n)))

  (* expand the EvalKernel operations into vector operations.  The parameters
   * are
   *	result	-- the lhs variable to store the result
   *	d	-- the vector width of the operation, which should be equal
   *		   to twice the support of the kernel
   *	h	-- the kernel
   *	k	-- the derivative of the kernel to evaluate
   *
   * The generated code is computing
   *
   *	result = a_0 + x*(a_1 + x*(a_2 + ... x*a_n) ... )
   *
   * as a d-wide vector operation, where n is the degree of the kth derivative
   * of h and the a_i are coefficient vectors that have an element for each
   * piece of h.  The computation is implemented as follows
   *
   *	m_n	= x * a_n
   *	s_{n-1}	= a_{n-1} + m_n
   *	m_{n-1}	= x * s_{n-1}
   *	s_{n-2}	= a_{n-2} + m_{n-1}
   *	m_{n-2}	= x * s_{n-2}
   *	...
   *	s_1	= a_1 + m_2
   *	m_1	= x * s_1
   *	result	= a_0 + m_1
   *
   * Note that the coeffient vectors are flipped (cf high-to-low/probe.sml).
   *)
    fun expandEvalKernel (result, d, h, k, [x]) = let
	  val {isCont, segs} = Kernel.curve (h, k)
	(* degree of polynomial *)
	  val deg = List.length(hd segs) - 1
	(* convert to a vector of vectors to give fast access *)
	  val segs = Vector.fromList (List.rev (List.map Vector.fromList segs))
	(* get the kernel coefficient value for the d'th term of the i'th
	 * segment.
	 *)
	  fun coefficient d i =
		Literal.Float(ratToFloat (Vector.sub (Vector.sub(segs, i), d)))
	  val ty = DstTy.vecTy d
	  val coeffs = List.tabulate (deg+1,
		fn i => DstIL.Var.new("a"^Int.toString i, ty))
	(* code to define the coefficient vectors *)
	  val coeffVecs = let
		fun mk (x, (i, code)) = let
		      val lits = List.tabulate(d, coefficient i)
		      val vars = List.tabulate(d, fn _ => DstIL.Var.new("_f", DstTy.realTy))
		      val code =
			    ListPair.map (fn (x, lit) => (x, DstIL.LIT lit)) (vars, lits) @
			      (x, DstIL.CONS(DstIL.Var.ty x, vars)) :: code
		      in
			(i-1, code)
		      end
		in
		  #2 (List.foldr mk (deg, []) coeffs)
		end
	(* build the evaluation of the polynomials in reverse order *)
	  fun pTmp i = DstIL.Var.new("prod" ^ Int.toString i, ty)
	  fun sTmp i = DstIL.Var.new("sum" ^ Int.toString i, ty)
	  fun eval (i, [coeff]) = let
		val m = pTmp i
		in
		  (m, [(m, DstIL.OP(DstOp.Mul ty, [x, coeff]))])
		end
	    | eval (i, coeff::r) = let
		val (m, stms) = eval(i+1, r)
		val s = sTmp i
		val m' = pTmp i
		val stms =
		      (m', DstIL.OP(DstOp.Mul ty, [x, s])) ::
		      (s, DstIL.OP(DstOp.Add ty, [coeff, m])) ::
		      stms
		in
		  (m', stms)
		end
	  val evalCode = (case coeffs
		 of [a0] => (* constant function *)
		      [(result, DstIL.VAR a0)]
		  | a0::r => let
		      val (m, stms) = eval (1, r)
		      in
			List.rev ((result, DstIL.OP(DstOp.Add ty, [a0, m]))::stms)
		      end
		(* end case *))
	  in
	    coeffVecs @ evalCode
	  end

(* FIXME: we will get better down-stream CSE if we structure the address computation
 * as
 *	(base + stride * (...)) + offset
 * since the lhs argument will be the same for each sample.
 *)
  (* add code to handle the offset and stride when addressing non-scalar image data *)
    fun adjustForStrideAndOffset (1, _, ix, code) = (ix, code)
      | adjustForStrideAndOffset (stride, 0, ix, code) = let
	  val offp = DstIL.Var.new ("offp", DstTy.intTy)
	  val stride' = DstIL.Var.new ("stride", DstTy.intTy)
	  in
	    (offp, imul(offp, stride', ix) :: ilit(stride', stride) :: code)
	  end
      | adjustForStrideAndOffset (stride, offset, ix, code) = let
	  val offp = DstIL.Var.new ("offp", DstTy.intTy)
	  val stride' = DstIL.Var.new ("stride", DstTy.intTy)
	  val offset' = DstIL.Var.new ("offset", DstTy.intTy)
	  val t = DstIL.Var.new ("t", DstTy.intTy)
	  val code =
		iadd(offp, offset', t) ::
		ilit (offset', offset) ::
		imul(t, stride', ix) ::
		ilit (stride', stride) ::
		code
	  in
	    (offp, code)
	  end

  (* compute the load address for a given set of voxels indices.  For the
   * operation
   *
   *	VoxelAddress<info,offset>(i_1, ..., i_d)
   *
   * the address is given by
   *
   *	base + offset + stride * (i_1 + N_1 * (i_2 + N_2 * (... + N_{d-1} * i_d) ...))
   *
   * where
   *	base	-- base address of the image data
   *	stride	-- number of samples per voxel
   *	offset  -- offset of sample being addressed
   *	N_i	-- size of ith axis in elements
   *
   * Note that we are following the Nrrd convention that the axes are ordered
   * in fastest to slowest order.  We are also assuming the C semantics of address
   * arithmetic, where the offset will be automatically scaled by the size of the
   * elements.
   *) 
    fun expandVoxelAddress (result, info, offset, [img, ix]) = let
	  val dim = ImageInfo.dim info
	  val sizes = ImageInfo.sizes info
	  val stride = ImageInfo.stride info
	  val shape = ImageInfo.voxelShape info
	  val (offp, code) = adjustForStrideAndOffset (stride, offset, ix, [])
	  val addrTy = DstTy.AddrTy info
	  val base = DstIL.Var.new ("imgBaseAddr", addrTy)
	  val code = (result, DstIL.OP(DstOp.Add addrTy, [base, offp])) ::
		(base, DstIL.OP(DstOp.ImageAddress info, [img])) ::
		code
	  in
	    List.rev code
	  end
      | expandVoxelAddress (result, info, offset, img::ix1::indices) = let
	  val dim = ImageInfo.dim info
	  val sizes = ImageInfo.sizes info
	  val stride = ImageInfo.stride info
	  val shape = ImageInfo.voxelShape info
	(* get N_1 ... N_{d-1} *)
	  val sizes = List.take (sizes, List.length sizes - 1)
	(* generate the address computation code in reverse order *)
	  fun gen (d, [n], [ix]) = let
		val n' = DstIL.Var.new ("n" ^ Int.toString d, DstTy.intTy)
		val t = DstIL.Var.new ("t", DstTy.intTy)
		val code = [
			imul(t, n', ix),
			ilit(n', n)
		      ]
		in
		  (t, code)
		end
	    | gen (d, n::ns, ix::ixs) = let
		val n' = DstIL.Var.new ("n" ^ Int.toString d, DstTy.intTy)
		val t1 = DstIL.Var.new ("t1", DstTy.intTy)
		val t2 = DstIL.Var.new ("t2", DstTy.intTy)
		val (t, code) = gen (d+1, ns, ixs)
		val code =
		      imul(t2, n', t1) ::
		      ilit(n', n) ::
		      iadd(t1, ix, t) :: code
		in
		  (t2, code)
		end
	  val (tmp, code) = gen (0, sizes, indices)
	  val t = DstIL.Var.new ("index", DstTy.intTy)
	  val code = iadd(t, ix1, tmp) :: code
	  val (offp, code) = adjustForStrideAndOffset (stride, offset, t, code)
	  val addrTy = DstTy.AddrTy info
	  val base = DstIL.Var.new ("imgBaseAddr", addrTy)
	  val code = (result, DstIL.OP(DstOp.Add addrTy, [base, offp])) ::
		(base, DstIL.OP(DstOp.ImageAddress info, [img])) ::
		code
	  in
	    List.rev code
	  end

    fun expandOp (env, y, rator, args) = let
	  val args' = renameList(env, args)
	  fun assign rator' = [(y, DstIL.OP(rator', args'))]
	  in
	    case rator
	     of SrcOp.Add ty => assign (DstOp.Add ty)
	      | SrcOp.Sub ty => assign (DstOp.Sub ty)
	      | SrcOp.Mul ty => assign (DstOp.Mul ty)
	      | SrcOp.Div ty => assign (DstOp.Div ty)
	      | SrcOp.Neg ty => assign (DstOp.Neg ty)
	      | SrcOp.Abs ty => assign (DstOp.Abs ty)
	      | SrcOp.LT ty => assign (DstOp.LT ty)
	      | SrcOp.LTE ty => assign (DstOp.LTE ty)
	      | SrcOp.EQ ty => assign (DstOp.EQ ty)
	      | SrcOp.NEQ ty => assign (DstOp.NEQ ty)
	      | SrcOp.GT ty => assign (DstOp.GT ty)
	      | SrcOp.GTE ty => assign (DstOp.GTE ty)
	      | SrcOp.Not => assign (DstOp.Not)
	      | SrcOp.Max => assign (DstOp.Max)
	      | SrcOp.Min => assign (DstOp.Min)
	      | SrcOp.Lerp ty => assign (DstOp.Lerp ty)
	      | SrcOp.Dot d => assign (DstOp.Dot d)
	      | SrcOp.MulVecMat(d1, d2) => assign (DstOp.MulVecMat(d1, d2))
	      | SrcOp.MulMatVec(d1, d2) => assign (DstOp.MulMatVec(d1, d2))
	      | SrcOp.MulMatMat(d1, d2, d3) => assign (DstOp.MulMatMat(d1, d2, d3))
	      | SrcOp.Cross => assign (DstOp.Cross)
	      | SrcOp.Select(ty, i) => assign (DstOp.Select(ty, i))
	      | SrcOp.Norm ty => assign (DstOp.Norm ty)
	      | SrcOp.Normalize d => assign (DstOp.Normalize d)
	      | SrcOp.Scale ty => assign (DstOp.Scale ty)
	      | SrcOp.Zero ty => assign (DstOp.Zero ty)
	      | SrcOp.CL => assign (DstOp.CL)
	      | SrcOp.PrincipleEvec ty => assign (DstOp.PrincipleEvec ty)
	      | SrcOp.Identity n => assign (DstOp.Identity n)
	      | SrcOp.Trace d => assign (DstOp.Trace d)
	      | SrcOp.Subscript ty => assign (DstOp.Subscript ty)
	      | SrcOp.Ceiling d => assign (DstOp.Ceiling d)
	      | SrcOp.Floor d => assign (DstOp.Floor d)
	      | SrcOp.Round d => assign (DstOp.Round d)
	      | SrcOp.Trunc d => assign (DstOp.Trunc d)
	      | SrcOp.IntToReal => assign (DstOp.IntToReal)
	      | SrcOp.RealToInt d => assign (DstOp.RealToInt d)
	      | SrcOp.VoxelAddress(info, offset) => expandVoxelAddress (y, info, offset, args')
	      | SrcOp.LoadVoxels(rty, d) => assign (DstOp.LoadVoxels(rty, d))
	      | SrcOp.PosToImgSpace info => assign (DstOp.PosToImgSpace info)
	      | SrcOp.TensorToWorldSpace(info, ty) => assign (DstOp.TensorToWorldSpace(info, ty))
	      | SrcOp.EvalKernel(d, h, k) => expandEvalKernel(y, d, h, k, args')
	      | SrcOp.LoadImage info => assign (DstOp.LoadImage info)
	      | SrcOp.Inside info => assign (DstOp.Inside info)
	      | SrcOp.Input(ty, name) => assign (DstOp.Input(ty, name))
	      | SrcOp.InputWithDefault(ty, name) => assign (DstOp.InputWithDefault(ty, name))
	    (* end case *)
	  end

  (* expand a SrcIL assignment to a DstIL CFG *)
    fun expand (env, (y, rhs)) = let
	  val y' = rename (env, y)
	  fun assign rhs = [(y', rhs)]
	  in
	    case rhs
	     of SrcIL.VAR x => assign (DstIL.VAR(rename(env, x)))
	      | SrcIL.LIT lit => assign (DstIL.LIT lit)
	      | SrcIL.OP(rator, args) => expandOp (env, y', rator, args)
	      | SrcIL.APPLY(f, args) => assign(DstIL.APPLY(f, renameList(env, args)))
	      | SrcIL.CONS(ty, args) => assign (DstIL.CONS(ty, renameList(env, args)))
	    (* end case *)
	  end

    structure Trans =  TranslateFn (
      struct
	structure SrcIL = SrcIL
	structure DstIL = DstIL

	type var_env = var_env

	val rename = rename
	val renameList = renameList
	val expand = DstIL.CFG.mkBlock o expand
      end)

    fun translate prog = let
	  val prog = Trans.translate prog
	  in
	    LowILCensus.init prog;
	    prog
	  end

  end

<script type="text/javascript">//<![CDATA[
		function toggle_ffErrors() {
			var errorsblock = document.getElementById("ffErrorsBlock");
			var errorsgroup = document.getElementById("ffErrors");
			if (errorsblock.style.display == "none") {
				errorsblock.style.display = "block";
				errorsgroup.style.right = "10px";
			} else {
				errorsblock.style.display = "none";
				errorsgroup.style.right = "300px";
			}
		}
//]]></script>
<div id="ffErrors">
<a href="javascript:toggle_ffErrors();">Click to toggle</a>
<div id="ffErrorsBlock">
	<div class="error">does not end with &lt;/html&gt; tag</div>
	<div class="error">does not end with &lt;/body&gt; tag</div>
	<div class="info">The output has ended thus: g = let
	  val prog = Trans.translate prog
	  in
	    LowILCensus.init prog;
	    prog
	  end

  end</div>
</div></div>
</body></html>