(* high-to-mid.sml
 *
 * Translation from HighIR to MidIR representations.
 *
 * This code is part of the Diderot Project (http://diderot-language.cs.uchicago.edu)
 *
 * COPYRIGHT (c) 2015 The University of Chicago
 * All rights reserved.
 *)

structure HighToMid : sig

    val translate : HighIR.program -> MidIR.program

  end = struct

    structure SrcIR = HighIR
    structure SrcTy = HighIRTypes
    structure SrcOp = HighOps
    structure SrcSV = SrcIR.StateVar
    structure VTbl = SrcIR.Var.Tbl
    structure DstIR = MidIR
    structure DstTy = MidIRTypes
    structure DstOp = MidOps
    structure InP = Inputs
    structure BCtl = BorderCtl

    fun getIncUse (DstIR.V{useCnt, ...})= !useCnt
    fun incUseD (DstIR.V{useCnt, ...}) = (useCnt := !useCnt + 1)
    fun useCount (SrcIR.V{useCnt, ...}) = !useCnt 
    fun useD x = (incUseD x; x)
    fun iTos e=Int.toString e
    val cnt = ref 0
    fun nameCnt e= String.concat[SrcIR.Var.toString e, "(",Int.toString(useCount e),")"]
    fun genName prefix = let
          val n = !cnt
          in
            cnt := n+1;
            String.concat[prefix, "_", Int.toString n]
          end

    fun getRHS x  = (case SrcIR.Var.binding x
           of SrcIR.VB_RHS(SrcIR.OP(rator, args)) => (rator, args)
            | SrcIR.VB_RHS(SrcIR.VAR x') => getRHS x'
            | SrcIR.VB_RHS(SrcIR.GLOBAL x') => getRHS(SrcIR.GlobalVar.binding x')
            | vb => raise Fail(concat[
                  "expected rhs operator for ", SrcIR.Var.toString x,
                  " but found ", SrcIR.vbToString vb
                ])
          (* end case *))

  (* get the image referenced on a RHS and its border control (if any) *)
    fun getRHSImage x = let
          fun get x = (case getRHS x
                 of (SrcOp.LoadImage(SrcTy.ImageTy v, _), _) => v
                  | (SrcOp.Input(InP.INP{init=SOME(InP.Proxy(_, v)), ...}), _) => v
                  | (SrcOp.Input(InP.INP{init=SOME(InP.Image v), ...}), _) => v
                  | _ => raise Fail "bogus image variable"
                (* end case *))
          in
            case getRHS x
             of (SrcOp.BorderCtlDefault _, [img, v]) => (get img, BCtl.Default v)
              | (SrcOp.BorderCtlClamp _, [img]) => (get img, BCtl.clamp)
              | (SrcOp.BorderCtlMirror _, [img]) => (get img, BCtl.mirror)
              | (SrcOp.BorderCtlWrap _, [img]) => (get img, BCtl.wrap)
              | (SrcOp.LoadImage(SrcTy.ImageTy v, _), _) => (v, BCtl.None)
              | (SrcOp.Input(InP.INP{init=SOME(InP.Proxy(_, v)), ...}), _) => (v, BCtl.None)
              | (SrcOp.Input(InP.INP{init=SOME(InP.Image v), ...}), _) => (v, BCtl.None)
              | _ => raise Fail "bogus image variable"
            (* end case *)
          end

    fun cvtTy SrcTy.BoolTy = DstTy.BoolTy
      | cvtTy SrcTy.StringTy = DstTy.StringTy
      | cvtTy SrcTy.IntTy = DstTy.intTy
      | cvtTy (SrcTy.TensorTy dd) = DstTy.tensorTy dd
      | cvtTy (SrcTy.TupleTy tys) = DstTy.TupleTy(List.map cvtTy tys)
      | cvtTy (SrcTy.SeqTy(ty, n)) = DstTy.SeqTy(cvtTy ty, n)
    (* we replace Kernel and Field operations by 0, so the types are mapped to int *)
      | cvtTy SrcTy.KernelTy = DstTy.KernelTy
      | cvtTy SrcTy.FieldTy = DstTy.intTy
      | cvtTy (SrcTy.ImageTy info) = DstTy.ImageTy info
      | cvtTy ty = raise Fail("unexpected type " ^ SrcTy.toString ty)

  (* instantiate the translation environment *)
    structure Env = TranslateEnvFn (
      struct
        structure SrcIR = SrcIR
        structure DstIR = DstIR
        val cvtTy = cvtTy
      end)

  (* expand raising a real to an integer power.  When we know the exponent, we can inline
   * multiplications.
   *)
    fun expandPower (env, y, [x, n]) = let
          fun getConst x = (case SrcIR.Var.binding x
                 of SrcIR.VB_RHS(SrcIR.VAR x') => getConst x'
                  | SrcIR.VB_RHS(SrcIR.LIT(Literal.Int n)) => SOME n
                  | vb => NONE
                (* end case *))
          val x = Env.rename(env, x)
          fun pow () = let
                val t = DstIR.Var.new("n", DstTy.realTy)
                in [
                  (t, DstIR.OP(DstOp.IntToReal, [Env.rename(env, n)])),
                  (y, DstIR.APPLY(MathFuns.pow, [x, t]))
                ] end
          in
            case getConst n
             of SOME 0 => [(y, DstIR.LIT(Literal.Float(FloatLit.one)))]
              | SOME 1 => [(y, DstIR.VAR x)]
              | SOME ~1 => let
                  val t = DstIR.Var.new("one", DstTy.realTy)
                  in [
                    (t, DstIR.LIT(Literal.Float(FloatLit.one))),
(* FIXME: shouldn't be IDiv; should be real division! *)
                    (y, DstIR.OP(DstOp.IDiv , [t, x]))
                  ] end
(* FIXME: shouldn't be IMul; should be real multiplication! *)
              | SOME 2 => [(y, DstIR.OP(DstOp.IMul , [x, x]))]
(* FIXME: expand into multiplications
              | SOME n =>
*) | SOME _ => pow()
              | NONE => pow()
            (* end case *)
          end

  (* expand the field Inside operator into a image-space test *)
    fun expandInside (env, result, pos, fld) = (case getRHS fld
           of SrcIR.EINAPP(_, [img, h]) => (case (getRHSImage img, getRHS h)
                 of (v, SrcIR.OP(SrcOp.Kernel(h, _), _)) => let
                      val pos = Env.rename (env, pos)
                      val img = Env.rename (env, img)    
                      val s = Kernel.support h
                      val dim = ImageInfo.dim  v
                      val (_, x, code) = TransformEin.WorldToImagespace(dim, v, pos, img)
                      in
                        code @ [(result, DstIR.OP(DstOp.Inside(v, s), [x, img]))]
                      end
                  | _ => raise Fail "bogus kernel binding"
                (* end case *))
            | _ => raise Fail "bogus field binding"
          (* end case *))

    fun arity (SrcTy.TensorTy[]) = 1
      | arity (SrcTy.TensorTy[d]) = d
      | arity _ = raise Fail "arity"

    fun expandOp (env, y, rator, args) = let
          fun assign rator' =
                [(y, DstIR.OP(rator', Env.renameList(env, args)))]
          fun cvtToInt rator' = let
                val t = DstIR.Var.new ("t", DstTy.realTy)
                in [
                  (t, DstIR.OP(rator', Env.renameList(env, args))),
                  (y, DstIR.OP(DstOp.RealToInt 1, [t]))
                ] end
          fun dummy () = [(y, DstIR.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.IMod  => assign (DstOp.IMod)
              | SrcOp.INeg  => assign (DstOp.INeg)
              | SrcOp.Abs ty => assign (DstOp.Abs(cvtTy ty))
              | SrcOp.LT ty => assign (DstOp.LT(cvtTy ty))
              | SrcOp.LTE ty => assign (DstOp.LTE(cvtTy ty))
              | SrcOp.EQ ty => assign (DstOp.EQ(cvtTy ty))
              | SrcOp.NEQ ty => assign (DstOp.NEQ(cvtTy ty))
              | SrcOp.GT ty => assign (DstOp.GT(cvtTy ty))
              | SrcOp.GTE ty => assign (DstOp.GTE(cvtTy ty))
              | SrcOp.Power => expandPower(env, y, args)
              | SrcOp.Not => assign DstOp.Not
              | SrcOp.Max => assign DstOp.Max
              | SrcOp.Min => assign DstOp.Min
              | SrcOp.Clamp ty => assign (DstOp.Clamp(cvtTy ty))
              | SrcOp.Lerp ty => assign (DstOp.Lerp(cvtTy ty))
              | SrcOp.Sqrt => assign DstOp.Sqrt
              | SrcOp.Norm(SrcTy.TensorTy []) => assign (DstOp.Abs DstTy.realTy)
              | SrcOp.Norm(SrcTy.TensorTy alpha) => let
		(* Note Norm is implemented with EINAPP as a summation over modulate then sqrt *)
                  val t = DstIR.Var.new (genName "t", DstTy.realTy)
                  val t = useD t
                  val a = Env.renameList(env, args)
                  in [
		    (t, DstIR.EINAPP(MkOperators.magnitudeTT alpha, a@a)),
		    (y, DstIR.OP(DstOp.Sqrt, [t]))
		  ] end
              | SrcOp.Normalize ty => assign (DstOp.Normalize(arity ty))
              | SrcOp.PrincipleEvec ty => assign (DstOp.PrincipleEvec(cvtTy ty))
              | SrcOp.Zero ty => assign (DstOp.Zero(cvtTy ty))
              | SrcOp.Slice(SrcTy.TensorTy argTy, mask) => let
		  val args' = let
			fun f e = let
			      val SrcIR.LIT(Literal.Int i) = getRHS e
			      in
				IntInf.toInt i
			      end
			in
			  List.map f (List.drop(args, 1))
			end
                  val DstTy.TensorTy rstTy = DstIR.Var.ty y
                  val rator = MkOperators.slice (argTy, mask, args', rstTy)
                  val a = Env.renameList(env, args)
                  in
                    [(y, DstIR.EINAPP(rator, a))]
                  end
              | SrcOp.Slice(ty, mask) => raise Fail "FIXME: Slice"
              | SrcOp.TensorSub(ty as SrcTy.TensorTy _) => assign (DstOp.Subscript(cvtTy ty))
              | SrcOp.Select(ty as SrcTy.TupleTy _, i) => assign (DstOp.Select(cvtTy ty, i))
              | SrcOp.Select(ty as SrcTy.SeqTy _, i) => assign (DstOp.Index(cvtTy ty, i))
              | SrcOp.SeqSub(ty as SrcTy.SeqTy _) => assign (DstOp.Subscript(cvtTy ty))
              | SrcOp.IntToReal => assign DstOp.IntToReal
              | SrcOp.TruncToInt => cvtToInt (DstOp.Trunc 1)
              | SrcOp.RoundToInt => cvtToInt (DstOp.Round 1)
              | SrcOp.CeilToInt => cvtToInt (DstOp.Ceiling 1)
              | SrcOp.FloorToInt => cvtToInt (DstOp.Floor 1)
              | SrcOp.Kernel h => assign (DstOp.Kernel h)
              | SrcOp.Inside _ => (case args
                   of [pos, fld] => expandInside(env, y, pos, fld)
                  (* end case *))
              | SrcOp.LoadImage(ty, nrrd, info) =>
                  assign (DstOp.LoadImage(DstTy.ImageTy info, nrrd, info))
              | SrcOp.Input inp => (case Inputs.imageInfo inp
                   of SOME info => let
                        val Inputs.INP{name, desc, init, ...} = inp
                        in
                          assign (DstOp.Input(Inputs.INP{
                              ty = DstTy.ImageTy info,
                              name = name, desc = desc,
                              init = init
                            }))
                        end
                    | _ => assign (DstOp.Input(Inputs.map cvtTy inp))
                  (* end case *))
              | rator => raise Fail("bogus operator " ^ SrcOp.toString rator)
            (* end case *)
          end
handle ex => (print(concat["error converting ", SrcOp.toString rator, "\n"]); raise ex)

  (* expandEINAPP: env* midil.var*EIN*mid-ilvar->DstIR.ASSGN list
   * Field operators are changed to zero
   *)
    fun expandEINAPP (env, srcy, y, rator, args) = (case SrcIR.Var.ty srcy
           of SrcTy.FieldTy => [DstIR.ASSGN(y, DstIR.LIT(Literal.Int 0))]
            | _ => if (useCount srcy > 0)
		then HandleEin.expandEinOp (srcy, (y, DstIR.EINAPP(rator, Env.renameList(env, args))))
		else []
          (* end case *))
handle ex => (print(concat["error converting: ", MidIR.Var.toString y, " = ",
EinPP.toString rator, " ", "(", String.concatWithMap ", " HighIR.Var.toString args, ")\n"]);
raise ex)

  (* expand a SrcIR assignment to a list of DstIR assignments *)
    fun expand (env, (y, rhs)) = let
          fun assign rhs = [DstIR.ASSGN(Env.rename (env, y), rhs)]
          in
            case rhs
             of SrcIR.GLOBAL x => assign (DstIR.GLOBAL(Env.renameGV(env, x)))
              | SrcIR.STATE x => assign (DstIR.STATE(Env.renameSV(env, x)))
              | SrcIR.VAR x => assign (DstIR.VAR(Env.rename(env, x)))
              | SrcIR.LIT lit => assign (DstIR.LIT lit)
              | SrcIR.OP(rator, args) =>
                  List.map DstIR.ASSGN (expandOp (env, Env.rename (env, y), rator, args))
              | SrcIR.CONS(ty, args) => assign (DstIR.CONS(cvtTy ty, Env.renameList(env, args)))
              | SrcIR.EINAPP(rator, args) =>
                  expandEINAPP (env, y, Env.rename (env, y), rator, args)
            (* end case *)
          end

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

    structure Trans =  TranslateFn (
      struct
        open Env
        val expand = DstIR.CFG.mkBlock o expand
        val mexpand = mexpand
      end)

    fun translate prog = let
(* need a prepass to add Inside tests for border control *)
          val prog = Trans.translate prog
          in
            MidIRCensus.init prog;
            prog
          end

  end