(* 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 SrcSV = SrcIL.StateVar
    structure SrcTy = MidILTypes
    structure VTbl = SrcIL.Var.Tbl
    structure DstIL = LowIL
    structure DstTy = LowILTypes
    structure DstOp = LowOps
    structure E=Ein
    structure P=Printer

    structure genEin=genEin


 
  (* instantiate the translation environment *)
    local
      type var_env = DstIL.var VTbl.hash_table
      type state_var_env = DstIL.state_var SrcSV.Tbl.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 renameSV (env : state_var_env, x) = (case SrcSV.Tbl.find env x
               of SOME x' => x'
                | NONE => let
                    val x' = DstIL.StateVar.new (SrcSV.isOutput x, SrcSV.name x, SrcSV.ty x)
                    in
                      SrcSV.Tbl.insert env (x, x');
                      x'
                    end
              (* end case *))
    in
    structure Env = TranslateEnvFn (
      struct
        structure SrcIL = SrcIL
        structure DstIL = DstIL
      
        type var_env = var_env
        type state_var_env = state_var_env
        val rename = rename
        val renameSV = renameSV
      end)
    end (* local *)


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

   
              
(* 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 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)
	
	(*Add here is of address type, assume it is okay to keep IADD since not tensors*)
	  val code = (result, DstIL.OP(DstOp.IAdd , [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} *)
(* FIXME: sizes is [] when the image does not have a proxy *)
	  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
(* FIXME: sizes is [] when the image does not have a proxy *)
	  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.IAdd , [base, offp])) ::
		(base, DstIL.OP(DstOp.ImageAddress info, [img])) ::
		code
	  in
	    List.rev code
	  end



    fun expandOp (env, y, rator, args) = let
	  val args' = Env.renameList (env, args)
	  fun assign rator' = [(y, DstIL.OP(rator', args'))]
     fun dummy () = [(y, DstIL.LIT(Literal.Int 0))]
	  in
	    case rator
	     of SrcOp.IAdd  => assign (DstOp.IAdd )
	      | SrcOp.ISub  => assign (DstOp.ISub )
	      | SrcOp.IMul  => assign (DstOp.IMul )
	      | SrcOp.IDiv  => assign (DstOp.IDiv )
	      | SrcOp.INeg  => assign (DstOp.INeg )
	      | 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.Clamp ty => assign (DstOp.Clamp ty)
	      | SrcOp.Lerp ty => assign (DstOp.Lerp ty)
		  | SrcOp.Norm ty => assign (DstOp.Norm ty)
	      | SrcOp.Normalize d => assign (DstOp.Normalize d)
	 
	      | SrcOp.Zero ty => assign (DstOp.Zero ty)
	      | SrcOp.PrincipleEvec ty => assign (DstOp.PrincipleEvec ty)
          | SrcOp.EigenVals2x2 => assign (DstOp.EigenVals2x2)
          | SrcOp.EigenVals3x3 => assign (DstOp.EigenVals3x3)
          | SrcOp.Select(ty as SrcTy.TupleTy tys, i) => assign (DstOp.Select(ty, i))
          | SrcOp.Index(ty, i) => assign (DstOp.Index(ty, i))
          | 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.Kernel h =>  assign(DstOp.Kernel h)
          | SrcOp.LoadImage info => assign (DstOp.LoadImage info)
	      | SrcOp.Inside info => assign (DstOp.Inside info)
       
            (*Input problems *)
	      (*| SrcOp.Input(ty, s, desc) => assign (DstOp.Input(ty, s, desc))
	      | SrcOp.InputWithDefault(ty, s, desc) =>assign (DstOp.InputWithDefault(ty, s, desc))
       *)
          | SrcOp.Transform V=> assign (DstOp.Transform V)
          | SrcOp.Translate V=> assign(DstOp.Translate V)
	      | rator => raise Fail("bogus operator " ^ SrcOp.toString rator)
	    (* end case *)
	  end
       
       
       
       fun printgetRHS x  = (case DstIL.Var.binding x
       of DstIL.VB_RHS(DstIL.OP(rator, args)) => print "Got Dstop"
       | DstIL.VB_RHS(DstIL.VAR x') => printgetRHS x'
       | DstIL.VB_RHS(DstIL.EINAPP(rator, args)) => print "got ein"
       | DstIL.VB_RHS(DstIL.CONS (ty,args))=>print "cons"
       | DstIL.VB_NONE=>print "\nnone"
            | vb => (print (String.concat["\n -- ", DstIL.Var.toString x,"but found ", DstIL.vbToString vb,"\n"]))
            (* end case *))
       
       
       
     
       fun printX(DstIL.ASSGN (x, _))=printgetRHS x
            
            
	  fun expandEinOp (env, y, e, args) = let
            val einargs=Env.renameList(env, args)
            val f=print(String.concat(["\n\n new ein \n", DstIL.Var.toString(y),"=",P.printerE(e)]@(List.map (fn e=> (DstIL.Var.toString(e)^",")) einargs)))
            val _=print "\n ** pre gen**"
            val (_,code)=genEin.genfn(y,e,args,einargs)
            val _=print "**\n post gen"
             val _=print(Int.toString(length(code)))
            val DstIL.ASSGN (a1,DstIL.OP(opss1,args1))=List.hd(List.rev(code))
            val c=DstIL.ASSGN (y,DstIL.OP(opss1,args1))
       
            val m=print (String.concat(["\n NEW-- ", DstIL.Var.toString y,"===", DstOp.toString opss1,"-"]@(List.map (fn e=> (DstIL.Var.toString(e)^",")) args1)))

            val g=print "\n made it post code"
       
              in code@[c]

            end

       
	handle ex => (print(concat["error converting  \n"]); raise ex)
	              


  (* expand a SrcIL assignment to a DstIL CFG *)
    fun expand (env, (y, rhs)) = let
	  val y' = Env.rename (env, y)

	  fun assign rhs = [DstIL.ASSGN(y', rhs)]
	  in
	    case rhs
       of SrcIL.STATE x => (assign (DstIL.STATE(Env.renameSV(env, x))))
          | SrcIL.VAR x => assign (DstIL.VAR(Env.rename(env, x)))
          | SrcIL.LIT lit => (assign (DstIL.LIT lit))
          | SrcIL.OP(rator, args) => (List.map DstIL.ASSGN (expandOp (env, y', rator, args)))
	      | SrcIL.APPLY(f, args) => assign (DstIL.APPLY(f, Env.renameList(env, args)))
	      | SrcIL.CONS(ty, args) => assign (DstIL.CONS(ty, Env.renameList(env, args)))
          | SrcIL.EINAPP(rator, args) => (*(print "\n Got one";List.map DstIL.ASSGN (expandEinOp (env, Env.rename (env, y), rator, args)))*)
       (print "\n ------------------------last name";print(SrcIL.Var.toString(y));print "---";expandEinOp (env, Env.rename (env, y), rator, args))
		 (* end case *)
	end


  (* expand a SrcIL multi-assignment to a DstIL CFG *)
    fun mexpand (env, (ys, rator, xs)) = let
          val ys' = Env.renameList(env, ys)
          val rator' = (case rator
                 of SrcOp.EigenVecs2x2 => DstOp.EigenVecs2x2
                  | SrcOp.EigenVecs3x3 => DstOp.EigenVecs3x3
                  | SrcOp.Print tys => DstOp.Print tys
                  | _ => raise Fail("bogus operator " ^ SrcOp.toString rator)
                (* end case *))
          val xs' = Env.renameList(env, xs)
          val nd = DstIL.Node.mkMASSIGN(ys', rator', xs')
          in
            DstIL.CFG{entry=nd, exit=nd}
          end

    structure Trans =  TranslateFn (
      struct
        open Env
        val expand = DstIL.CFG.mkBlock o expand
        val mexpand = mexpand
      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>