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

structure HighToMid : sig

    val translate : HighIL.program -> MidIL.program

  end = struct

    structure SrcIL = HighIL
    structure SrcOp = SrcIL.Op
    structure DstIL = MidIL
    structure DstOp = DstIL.Op
    structure VMap = SrcIL.Var.Map

  (* a tree representation of nested iterations over the image space, where the
   * height of the tree corresponds to the number of dimensions and at each node
   * we have as many children as there are iterations.
   *)
    structure IT =
      struct
	datatype ('nd, 'lf) iter_tree
	  = LF of 'lf
	  | ND of ('nd * ('nd, 'lf) iter_tree list)

	fun create (depth, width, ndAttr, f, lfAttr, init) = let
	      fun mk (d, i, arg) = if (d < depth)
		    then ND(ndAttr arg, List.tabulate(width, fn j => mk(d+1, j, f(j, arg))))
		    else LF(lfAttr arg)
	      in
		mk (0, 0, init)
	      end

	fun map (nd, lf) t = let
	      fun mapf (LF x) = LF(lf x)
		| mapf (ND(i, kids)) = ND(nd i, List.map mapf kids)
	      in
		mapf t
	      end

	fun foldr f init t = let
	      fun fold (LF x, acc) = f(x, acc)
		| fold (ND(_, kids), acc) = List.foldr fold acc kids
	      in
		fold t
	      end

      end

  (* expand the field Inside operator into a image-space test *)
    fun expandInside (env, result, pos, fld) = let
	  val pos' = lookupVar (env, pos)
	  val fld = (case valueOf fld
		 of SrcIL.OP(SrcOp.Field fld, []) => fld
		  | _ => raise Fail "bogus field binding"
		(* end case *))
	  fun expand (FieldDef.CONV(_, img, _)) => let
		val imgPos = newVar ??
		in [
		  (imgPos, DstIL.OP(DstOp.Transform img, [pos'])),
		  (result, DstIL.OP(DstOp.Inside img, [imgPos]))
		] end
	    | expand (FieldDef.NEG fld) => expand fld
	    | expand (FieldDef.SUM(fld1, dlf2)) => raise Fail "expandInside: SUM"
	  in
	    expand fld
	  end

  (* generate a new variable indexed by dimension *)
    local
      val dimNames = Vector.fromList[ "x", "y", "z" ];
    in
    fun newVar_dim (prefix, d) =
	  DstIL.Var.new (prefix ^ Vector.sub(dimNames, d))

    fun assign (x, rator, args) = (x, DstIL.OP(rator, args))
    fun cons (x, args) = (x, DstIL.CONS args)
    fun intLit (x, i) = (x, DstIL.LIT(Literal.Int(IntInf.fromInt i)))

    fun generate (n, f) = List.concat(List.tabulate(n, f))

  (* generate code for probing the field (D^k (v * h)) at pos *)
    fun probe (result, (k, v, h), pos) = let
	  val ImageInfo.ImgInfo{dim, ty=([], ty), ...} = v
	  val dimTy = DstOp.VecTy dim
	  val s = Kernel.support h
	  val sTy = DstOp.VecTy(2*s)
	(* generate the transform code *)
	  val x = DstIL.Var.new "x"	(* image-space position *)
	  val f = DstIL.Var.new "f"
	  val nd = DstIL.Var.new "nd"
	  val n = DstIL.Var.new "n"
	  val transformCode = [
		  assign(x, DstIL.Transform v, [pos]),
		  assign(nd, DstIL.Floor dim, [x]),
		  assign(f, DstIL.Sub dimTy, [x, nd]),
		  assign(n, DstOp.TruncToInt dim, [nd])
		]
	(* generate code to load the voxel data *)
	  val voxIter = let
		fun f (i, (offsets, id)) = (i - (s - 1) :: offsets, i::id)
		fun g (offsets, id) = {
			offsets = offsets,
			vox = DstIL.Var.new(String.concat("v" :: List.map Int.toString id))
		      }
		in
		  IT.create (depth, width, fn _ => (), f, g, ([], []))
		end
	  val loadCode = let
		fun genCode ({offsets, vox}, code) = let
		      fun computeIndices (_, []) = ([], [])
			| computeIndices (i, offset::offsets) = let
			    val index = newVar_dim("i", i)
			    val t1 = newVar "t1"
			    val t2 = newVar "t2"
			    val (indices, code) = computeIndices (i+1, offsets)
			    val code =
				  intLit(t1, offset) ::
				  assign(t2, DstOp.Select i, [n]) ::
				  assign(index, DstOp.Add(DstOp.IntTy), [t1, t2]) ::
				  code
			    val indices = index::indices
			    in
			      (indices, code)
			    end
		      val (indices, indicesCode) = computeIndices (0, ~(s-1) :: offsets)
		      val a = DstIL.Var.new "a"
		      in
			indicesCode :: [
			    assign(a, VoxelAddress v, indices),
			    assign(vox, LoadVoxels(ty, 2*s))
			  ] @ code
		      end
		in
		  IT.foldr genCode [] voxIter
		end
	  val voxVars = IT.foldr (fn ({vox, ...}, vs) => vox::vs) [] voxIter
	(* generate the code for computing the convolution coefficients *)
(* SOMETHING *)
	(* generate the reduction code *)
	  fun genReduce (d, IT.ND(kids), code) =
		if (d < dim)
		  then List.foldr (fn (nd, code) => genReduce(d+1, nd, code)) code kids
		  else let (* the kids will all be leaves *)
		    val vv = newVar "vv"
		    fun getVox (IT.LF{vox, offsets}) = vox
		    in
		      cons (vv, List.map getVox kids) ::
		      assign (t, DstIL.Dot, [h, vv]) :: code
		    end
	  val reduceCode = genReduce (1, voxIter, [])
	  in
	    transformCode @ loadCode @ reduceCode
	  end

    end

    fun expandProbe (env, result, fld, pos) = let
	  val pos' = lookupVar (env, pos)
	  val fld = (case valueOf fld
		 of SrcIL.OP(SrcOp.Field fld, []) => fld
		  | _ => raise Fail "bogus field binding"
		(* end case *))
	  fun expand (result, FieldDef.CONV(0, img, h)) => let
		val imgPos = newVar ??
		val xformStm = (imgPos, DstIL.OP(DstOp.Transform img, [pos']))
(* generate samples based on kernel support and dimensionality of image *)
		in
		  xformStm :: probeStms
		end
	    | expand (FieldDef.CONV(k, img, h)) => ??
	    | expand (FieldDef.NEG fld) => let
		val r = newVar ??
		val stms = expand (r, fld)
		in
		  (r, DstIL.OP(DstOp.Neg ty, [r])) :: stms
		end
	    | expand (FieldDef.SUM(fld1, dlf2)) => raise Fail "expandInside: SUM"
	  in
	    List.rev (expand (result, fld))
	  end

  end