(* c-target.sml
 *
 * COPYRIGHT (c) 2011 The Diderot Project (http://diderot-language.cs.uchicago.edu)
 * All rights reserved.
 *
 * Generate C code with SSE 4.2 intrinsics.
 *)

structure CTarget : TARGET =
  struct

    structure CL = CLang

    datatype ty
      = T_Bool
      | T_String
      | T_Int			(* default float type *)
      | T_Real			(* default real type *)
      | T_Vec of int
      | T_IVec of int
      | T_Image of int * RawTypes.ty	(* n-dimensional image data *)
      | T_Ptr of RawTypes.ty		(* pointer type *)

  (* string representation of types (for debugging *)
    fun tyToString ty = (case ty
	   of T_Bool => "T_Bool"
	    | T_String => "T_String"
	    | T_Int => "T_Int"
	    | T_Real => "T_Real"
	    | T_Vec n => concat["T_Vec(", Int.toString n, ")"]
	    | T_IVec n => concat["T_IVec(", Int.toString n, ")"]
	    | T_Image(n, ty) => concat[
		  "T_Image", Int.toString n, "D(", RawTypes.toString ty, ")"
		]
	    | T_Ptr ty => concat["T_Ptr(", RawTypes.toString ty, ")"]
	  (* end case *))

    datatype strand = Strand of {
	name : string,
	tyName : string,
	state : (ty * string) list ref,
	code : CL.decl list ref
      }

    type var = (ty * string) (* FIXME *)

    type exp = CLang.exp * ty

    type stm = CL.stm

    type method = unit (* FIXME *)

    datatype program = Prog of {
	globals : CL.decl list ref,
	topDecls : CL.decl list ref,
	strands : strand list ref
      }

  (* globals that specify the target characteristics.  These should be initialized
   * when the program object is created.
   *)
    val gVectorWid = ref 4
    val gIntTy = ref CL.int32
    val gRealTy = ref CL.float
    val gRealSuffix = ref "f"

  (* initialize globals based on target precision *)
    fun initTargetSpec () = if !Controls.doublePrecision
	  then (
	    gVectorWid := 2;
	    gIntTy := CL.int64;
	    gRealTy := CL.double;
	    gRealSuffix := "d")
	  else (
	    gVectorWid := 4;
	    gIntTy := CL.int32;
	    gRealTy := CL.float;
	    gRealSuffix := "f")

  (* for SSE, we have 128-bit vectors *)
    fun vectorWidth () = !gVectorWid

    fun vectorSuffix n = Int.toString n ^ !gRealSuffix

  (* target types *)
    val boolTy = T_Bool
    val intTy = T_Int
    val realTy = T_Real
    fun vecTy 1 = T_Real
      | vecTy n = if (n < 1) orelse (!gVectorWid < n)
	  then raise Size
	  else T_Vec n
    fun ivecTy 1 = T_Int
      | ivecTy n = if (n < 1) orelse (!gVectorWid < n)
	  then raise Size
	  else T_IVec n
    fun imageTy (ImageInfo.ImgInfo{ty=([], rTy), dim, ...}) = T_Image(dim, rTy)
    fun imageDataTy (ImageInfo.ImgInfo{ty=([], rTy), ...}) = T_Ptr rTy
    val stringTy = T_String

    val statusTy = CL.T_Named "Status_t"

  (* convert target types to CLang types *)
    fun cvtTy T_Bool = CLang.T_Named "bool"
      | cvtTy T_String = CL.charPtr
      | cvtTy T_Int = !gIntTy
      | cvtTy T_Real = !gRealTy
      | cvtTy (T_Vec n) = CLang.T_Named(concat["Diderot_vec", Int.toString n, "D_t"])
      | cvtTy (T_IVec n) = raise Fail "FIXME: T_IVec"
      | cvtTy (T_Image(n, _)) = CLang.T_Named(concat["Diderot_image", Int.toString n, "D_t"])
      | cvtTy (T_Ptr ty) = CL.T_Ptr(CL.T_Num ty)

  (* report invalid arguments *)
    fun invalid (name, []) = raise Fail("invaild "^name)
      | invalid (name, args) = let
	  fun arg2s (e, ty) = concat["(", CL.expToString e, " : ", tyToString ty, ")"]
	  val args = String.concatWith ", " (List.map arg2s args)
	  in
	    raise Fail(concat["invalid arguments to ", name, ": ", args])
	  end

  (* helper functions for checking the types of arguments *)
    fun scalarTy T_Int = true
      | scalarTy T_Real = true
      | scalarTy _ = false
    fun numTy T_Int = true
      | numTy T_Real = true
      | numTy (T_Vec _) = true
      | numTy (T_IVec _) = true
      | numTy _ = false

    fun newProgram () = (
	  initTargetSpec();
	  Prog{
	      globals = ref [],
	      topDecls = ref [],
	      strands = ref []
	    })

    fun globalInit (Prog{topDecls, ...}, init) = let
	  val initFn = CL.D_Func([], CL.voidTy, "Diderot_InitGlobals", [], init)
	  in
	    topDecls := initFn :: !topDecls
	  end

    structure Var =
      struct
	fun global (Prog{globals, ...}, ty, name) = (
	      globals := CL.D_Var([], cvtTy ty, name) :: !globals;
	      (ty, name))
	fun param (ty, name) = (ty, name)
	fun state (Strand{state, ...}, ty, name) = (
	      state := (ty, name) :: !state;
	      (ty, name))
	fun var (ty, name) = (ty, name)
	fun tmp ty = raise Fail "FIXME: Var.tmp"
      end

  (* expression construction *)
    structure Expr =
      struct
      (* return true if the given expression from is allowed as a subexpression *)
	fun allowedInline _ = true (* FIXME *)

      (* variable references *)
	fun global (ty, x) = (CL.mkVar x, ty)
	fun getState (ty, x) = (CL.mkIndirect(CL.mkVar "selfIn", x), ty)
	fun param (ty, x) = (CL.mkVar x, ty)
	fun var (ty, x) = (CL.mkVar x, ty)

      (* literals *)
	fun intLit n = (CL.mkInt(n, !gIntTy), intTy)
	fun floatLit f = (CL.mkFlt(f, !gRealTy), realTy)
	fun stringLit s = (CL.mkStr s, stringTy)
	fun boolLit b = (CL.mkBool b, boolTy)

      (* select from a vector *)
	fun select (i, (e, T_Vec n)) =
	      if (i < 0) orelse (n <= i)
		then raise Subscript
		else (CL.mkSubscript(e, CL.mkInt(IntInf.fromInt i, CL.int32)), T_Real)
	  | select (i, (e, T_IVec n)) =
	      if (i < 0) orelse (n <= i)
		then raise Subscript
		else (CL.mkSubscript(e, CL.mkInt(IntInf.fromInt i, CL.int32)), T_Int)
	  | select (_, x) = invalid("select", [x])

      (* vector (and scalar) arithmetic *)
	local
	  fun checkTys (ty1, ty2) = (ty1 = ty2) andalso numTy ty1
	  fun binop rator ((e1, ty1), (e2, ty2)) =
		if checkTys (ty1, ty2)
		  then (CL.mkBinOp(e1, rator, e2), ty1)
		  else invalid (
		    concat["binary operator \"", CL.binopToString rator, "\""],
		    [(e1, ty1), (e2, ty2)])
	in
	fun add ((e1, ty as T_Ptr _), (e2, T_Int)) = (CL.mkBinOp(e1, CL.#+, e2), ty)
	  | add args = binop CL.#+ args
	fun sub ((e1, ty as T_Ptr _), (e2, T_Int)) = (CL.mkBinOp(e1, CL.#-, e2), ty)
	  | sub args = binop CL.#- args
      (* NOTE: multiplication and division are also used for scaling *)
	fun mul ((e1, T_Real), (e2, T_Vec n)) =
	      (CL.E_Apply("Diderot_scale"^vectorSuffix n, [e1, e2]), T_Vec n)
	  | mul args = binop CL.#* args
	fun divide ((e1, T_Vec n), (e2, T_Real)) =
	      (CL.E_Apply("Diderot_scale"^vectorSuffix n,
		[CL.mkBinOp(CL.mkFlt(FloatLit.one, !gRealTy), CL.#/, e2), e1]), T_Vec n)
	  | divide args = binop CL.#/ args
	end (* local *)
	fun neg (e, T_Bool) = raise Fail "invalid argument to neg"
	  | neg (e, ty) = (CL.mkUnOp(CL.%-, e), ty)

	fun abs (e, T_Int) = (CL.mkApply("abs", [e]), T_Int)	(* FIXME: not the right type for 64-bit ints *)
	  | abs (e, T_Real) = (CL.mkApply("fabs" ^ !gRealSuffix, [e]), T_Real)
	  | abs (e, T_Vec n) = raise Fail "FIXME: Expr.abs"
	  | abs (e, T_IVec n) = raise Fail "FIXME: Expr.abs"
	  | abs _ = raise Fail "invalid argument to abs"

	fun dot ((e1, T_Vec n1), (e2, T_Vec n2)) =
	      (CL.E_Apply("Diderot_dot"^vectorSuffix n1, [e1, e2]), T_Real)
	  | dot _ = raise Fail "invalid argument to dot"

	fun cross ((e1, T_Vec 3), (e2, T_Vec 3)) = raise Fail "FIXME: Expr.cross"
	  | cross _ = raise Fail "invalid argument to cross"

	fun length (e, T_Vec n) =
	      (CL.E_Apply("Diderot_length"^vectorSuffix n, [e]), T_Real)
	  | length _ = raise Fail "invalid argument to length"

	fun normalize (e, T_Vec n) =
	      (CL.E_Apply("Diderot_normalize"^vectorSuffix n, [e]), T_Vec n)
	  | normalize _ = raise Fail "invalid argument to length"

      (* comparisons *)
	local
	  fun checkTys (ty1, ty2) =
		(ty1 = ty2) andalso scalarTy ty1
	  fun cmpop rator ((e1, ty1), (e2, ty2)) =
		if checkTys (ty1, ty2)
		  then (CL.mkBinOp(e1, rator, e2), T_Bool)
		  else invalid (
		    concat["compare operator \"", CL.binopToString rator, "\""],
		    [(e1, ty1), (e2, ty2)])
	in
	val lt = cmpop CL.#<
	val lte = cmpop CL.#<=
	val equ = cmpop CL.#==
	val neq = cmpop CL.#!=
	val gte = cmpop CL.#>=
	val gt = cmpop CL.#>
	end (* local *)

      (* logical connectives *)
	fun not (e, T_Bool) = (CL.mkUnOp(CL.%!, e), T_Bool)
	  | not _ = raise Fail "invalid argument to not"
	fun && ((e1, T_Bool), (e2, T_Bool)) = (CL.mkBinOp(e1, CL.#&&, e2), T_Bool)
	  | && _ = raise Fail "invalid arguments to &&"
	fun || ((e1, T_Bool), (e2, T_Bool)) = (CL.mkBinOp(e1, CL.#||, e2), T_Bool)
	  | || _ = raise Fail "invalid arguments to ||"

	local
	  fun checkTys (ty1, ty2) = (ty1 = ty2) andalso scalarTy ty1
	  fun binFn f ((e1, ty1), (e2, ty2)) =
		if checkTys (ty1, ty2)
		  then (CL.mkApply(f, [e1, e2]), ty1)
		  else raise Fail "invalid arguments to binary function"
	in
      (* misc functions *)
	val min = binFn "Diderot_min"
	val max = binFn "Diderot_max"
	end (* local *)

      (* math functions *)
	fun pow ((e1, T_Real), (e2, T_Real)) =
	      if !Controls.doublePrecision
		then (CL.mkApply("pow", [e1, e2]), T_Real)
		else (CL.mkApply("powf", [e1, e2]), T_Real)
	  | pow _ = raise Fail "invalid arguments to pow"

	local
	  fun r2r (ff, fd) (e, T_Real) = if !Controls.doublePrecision
		then (CL.mkApply(fd, [e]), T_Real)
		else (CL.mkApply(ff, [e]), T_Real)
	    | r2r (_, fd) _ = raise Fail("invalid argument for "^fd)
	in
	val sin = r2r ("sinf", "sin")
	val cos = r2r ("cosf", "cos")
	val sqrt = r2r ("sqrtf", "sqrt")
      (* rounding *)
	val trunc = r2r ("truncf", "trunc")
	val round = r2r ("roundf", "round")
	val floor = r2r ("floorf", "floor")
	val ceil  = r2r ("ceilf", "ceil")
	end (* local *)

      (* conversions *)
	fun toReal (e, T_Int) = (CL.mkCast(!gRealTy, e), T_Real)
	  | toReal e = invalid ("toReal", [e])

	fun truncToInt (e as (_, T_Real)) = (CL.mkCast(!gIntTy, #1(trunc e)), T_Int)
	  | truncToInt e = invalid ("truncToInt", [e])
	fun roundToInt (e as (_, T_Real)) = (CL.mkCast(!gIntTy, #1(round e)), T_Int)
	  | roundToInt e = invalid ("roundToInt", [e])
	fun ceilToInt (e as (_, T_Real)) = (CL.mkCast(!gIntTy, #1(floor e)), T_Int)
	  | ceilToInt e = invalid ("ceilToInt", [e])
	fun floorToInt (e as (_, T_Real)) = (CL.mkCast(!gIntTy, #1(ceil e)), T_Int)
	  | floorToInt e = invalid ("floorToInt", [e])

      (* runtime system hooks *)
	fun imageAddr (e, T_Image(_, rTy)) = let
	      val cTy = CL.T_Ptr(!gRealTy)
	      in
		(CL.mkCast(cTy, CL.mkIndirect(e, "data")), T_Ptr rTy)
	      end
	  | imageAddr a = invalid("imageAddr", [a])
	fun getImgData (e, T_Ptr rTy) = let
	      val realTy as CL.T_Num rTy' = !gRealTy
	      val e = CL.E_UnOp(CL.%*, e)
	      in
		if (rTy' = rTy)
		  then (e, T_Real)
		  else (CL.E_Cast(realTy, e), T_Real)
	      end
	  | getImgData a = invalid("getImgData", [a])
	fun posToImgSpace ((img, T_Image(d, _)), (pos, T_Vec n)) = let
	      val fnName = concat["Diderot_ToImageSpace", Int.toString d, "D"]
	      val e = CL.mkApply(fnName, [img, pos])
	      in
		(e, T_Vec n)
	      end
	  | posToImgSpace (a, b) = invalid("posToImgSpace", [a, b])
	fun inside ((pos, T_Vec n), (img, T_Image(d, _)), s) = let
	      val fnName = concat["Diderot_Inside", Int.toString d, "D"]
	      val e = CL.mkApply(fnName,
		    [pos, img, CL.mkInt(IntInf.fromInt n, CL.int32)])
	      in
		(e, T_Bool)
	      end
	  | inside (a, b, _) = invalid("inside", [a, b])

      end (* Expr *)

  (* statement construction *)
    structure Stmt =
      struct
	val comment = CL.S_Comment
	fun assignState ((_, x), (e, _)) =
	      CL.mkAssign(CL.mkIndirect(CL.mkVar "selfOut", x), e)
	fun assign ((_, x), (e, _)) = CL.mkAssign(CL.mkVar x, e)
	fun decl ((ty, x), SOME(e, _)) = CL.mkDecl(cvtTy ty, x, SOME e)
	  | decl ((ty, x), NONE) = CL.mkDecl(cvtTy ty, x, NONE)
	val block = CL.mkBlock
	fun ifthen ((e, T_Bool), s1) = CL.mkIfThen(e, s1)
	fun ifthenelse ((e, T_Bool), s1, s2) = CL.mkIfThenElse(e, s1, s2)
      (* special Diderot forms *)
	fun cons (lhs, args) = comment ["**** cons ****"] (* FIXME *)
	fun getImgData (lhs, n, e) = comment ["**** getImgData ****"] (* FIXME *)
	fun loadImage (lhs : var, dim, name : exp) = let
	      val sts = "sts"
	      val imgTy = CL.T_Named(concat["Diderot_image", Int.toString dim, "D_t"])
	      val loadFn = concat["Diderot_LoadImage", Int.toString dim, "D"]
	      in [
		CL.S_Decl(
		  statusTy, sts,
		  SOME(CL.E_Apply(loadFn, [#1 name, CL.mkUnOp(CL.%&, CL.E_Var(#2 lhs))])))
	      ] end
	fun input (lhs : var, name, optDflt) = let
	      val sts = "sts"
	      val inputFn = (case #1 lhs
		     of T_String => "Diderot_InputString"
		      | T_Real => "Diderot_InputReal"
		      | T_Vec 3 => "Diderot_InputVec3"
		      | ty => raise Fail("unsupported input type " ^ tyToString ty)
		    (* end case *))
	      val lhs = CL.E_Var(#2 lhs)
	      val (initCode, hasDflt) = (case optDflt
		     of SOME(e, _) => ([CL.S_Assign(lhs, e)], true)
		      | NONE => ([], false)
		    (* end case *))
	      val code = [
		    CL.S_Decl(
		      statusTy, sts,
		      SOME(CL.E_Apply(inputFn, [
			  CL.E_Str name, CL.mkUnOp(CL.%&, lhs), CL.mkBool hasDflt
			])))
		    ]
	      in
		initCode @ code
	      end
	fun die () = comment ["**** die ****"] (* FIXME *)
	fun stabilize () = comment ["**** stabilize ****"] (* FIXME *)
      end

    structure Strand =
      struct
	fun define (Prog{strands, ...}, strandId) = let
	      val strand = Strand{
		      name = strandId,
		      tyName = strandId ^ "_t",
		      state = ref [],
		      code = ref []
		    }
	      in
		strands := strand :: !strands;
		strand
	      end

      (* register the strand-state initialization code.  The variables are the strand
       * parameters.
       *)
	fun init (Strand{name, tyName, code, ...}, params, init) = let
	      val fName = name ^ "_InitState"
	      val params =
		    CL.PARAM([], CL.T_Ptr(CL.T_Named tyName), "selfOut") ::
		      List.map (fn (ty, x) => CL.PARAM([], cvtTy ty, x)) params
	      val initFn = CL.D_Func([], CL.voidTy, fName, params, init)
	      in
		code := initFn :: !code
	      end

      (* register a strand method *)
	fun method (Strand{name, tyName, code, ...}, methName, body) = let
	      val fName = concat[name, "_", methName]
	      val params = [
		      CL.PARAM([], CL.T_Ptr(CL.T_Named tyName), "selfIn"),
		      CL.PARAM([], CL.T_Ptr(CL.T_Named tyName), "selfOut")
		    ]
	      val methFn = CL.D_Func([], CL.int32, fName, params, body)
	      in
		code := methFn :: !code
	      end
      end (* Strand *)

    fun genStrand (Strand{name, tyName, state, code}) = let
	  val selfTyDef = CL.D_StructDef(
		  List.rev (List.map (fn (ty, x) => (cvtTy ty, x)) (!state)),
		  tyName)
	  in
	    selfTyDef :: List.rev (!code)
	  end

    fun generate (baseName, Prog{globals, topDecls, strands}) = let
	  val fileName = OS.Path.joinBaseExt{base=baseName, ext=SOME "c"}
	  val outS = TextIO.openOut fileName
	  val ppStrm = PrintAsC.new outS
	  fun ppDecl dcl = PrintAsC.output(ppStrm, dcl)
	  in
	    List.app ppDecl (List.rev (!globals));
	    List.app ppDecl (List.rev (!topDecls));
(* what about the strands, etc? *)
	    List.app (fn strand => List.app ppDecl (genStrand strand)) (!strands);
	    PrintAsC.close ppStrm;
	    TextIO.closeOut outS
	  end

  end

structure CBackEnd = CodeGenFn(CTarget)