(* gen-inputs.sml
 *
 * COPYRIGHT (c) 2012 The Diderot Project (http://diderot-language.cs.uchicago.edu)
 * All rights reserved.
 *)

structure GenInputs : sig

  (* input variable descriptor: type, name, description, and default *)
    type input_desc = (TreeIL.global_var * string * string option * TreeIL.exp option)

    val gatherInputs : TreeIL.block -> input_desc list

  (*** Support for standalone executables ***)

  (* generate the input initialization structure that we use to initialize input
   * globals from command-line arguments in stand-alone executables.
   *)
    val genInputsStruct : Properties.props * input_desc list -> CLang.decl list

  (* generate the common functions that handle inputs for standalone executables.  These are:
   *    InitDefaults    -- called to initialize the default input values
   *    RegisterInputs  -- called to register the command-line options for the input globals
   *    InitInptus      -- called to initialize the input globals from the values specified
   *                       on the command line.
   *)
    val genExecInputFuns : Properties.props * input_desc list -> CLang.decl list

  (*** Support for libraries ***)

  (* generate the typedef for the defined-input flag struct. *)
    val genDefinedInpStruct : Properties.props * input_desc list -> CLang.decl list

  (* generated the functions to initialize inputs for the library API.  This function also
   * generates the function to initialize the defined input flags struct.
   *)
    val genInputFuns : Properties.props * input_desc list -> CLang.decl list

  end = struct

    structure IL = TreeIL
    structure Ty = TreeIL.Ty
    structure GVar = IL.GlobalVar
    structure CL = CLang
    structure N = CNames
    structure ToC = TreeToC
    structure TrTy = CTyTranslate

    type input_desc = (TreeIL.global_var * string * string option * IL.exp option)

    val nrrdPtrTy = CL.T_Ptr(CL.T_Named "Nrrd")
    val wrldPrefixTy = CL.T_Ptr(CL.T_Named "WorldPrefix_t")

  (* world pointer cast to the world prefix type *)
    val wrldPrefix = CL.mkCast(wrldPrefixTy, CL.mkVar "wrld")

    type props = Properties.props

  (* translate a TreeIL type to the C types used to represent it in the external API *)
    val trType = CTyTranslate.toCType

  (* an l-value expression for accessing a defined-input flag *)
    fun defined gv = CL.mkSelect(CL.mkIndirect(CL.mkVar "wrld", "definedInp"), GVar.name gv)

  (* an l-value expression for accessing a global variable *)
    fun global gv = CL.mkIndirect(CL.mkIndirect(CL.mkVar "wrld", "globals"), GVar.name gv)

  (* gather the input globals from the input initialization block *)
    fun gatherInputs (IL.Block{body, ...}) = let
          fun inputInfo cvt (x, name, desc, NONE) = SOME(x, name, desc, NONE)
            | inputInfo cvt (x, name, desc, SOME v) = SOME(x, name, desc, SOME(cvt v))
          fun gather (IL.S_Input inp) = inputInfo (fn e => e) inp
            | gather (IL.S_InputNrrd inp) = inputInfo (fn s => TreeIL.E_Lit(Literal.String s)) inp
            | gather _ = NONE
          in
            List.mapPartial gather body
          end

  (* generate the input initialization structure that we use to initialize input
   * globals from command-line arguments in stand-alone executables.
   *)
    fun genInputsStruct (_, []) = []
      | genInputsStruct (tgt : props, inputs) = let
          fun mkField (gv, _, _, _) = (case GVar.ty gv
                 of Ty.DynSeqTy _ => (CL.charPtr, GVar.name gv)
                  | Ty.ImageTy _ => (CL.charPtr, GVar.name gv)
                  | ty => (trType ty, GVar.name gv)
                (* end case *))
          in
            [CL.D_StructDef(NONE, List.map mkField inputs, SOME(N.inputsTy tgt))]
          end

    local
      fun subscript (e, i) = CL.mkSubscript(e, CL.mkInt(IntInf.fromInt i))

    (* initialize an external C representation from constant expressions *)
      fun initC {ty, dst, src} = let
            fun cvtExp e = TreeToC.trExp (TreeIL.Var.Map.empty, e)
            fun simpleInit () = [CL.mkAssign(dst, cvtExp src)]
            in
              case (ty, src)
               of (Ty.BoolTy, _) => simpleInit ()
                | (Ty.StringTy, _) => simpleInit ()
                | (Ty.IntTy, _) => simpleInit ()
                | (Ty.TensorTy[], _) => simpleInit ()
                | (Ty.TensorTy[n], TreeIL.E_Cons(_, es)) => let
                    fun lp (_, []) = []
                      | lp (i, e::es) = CL.mkAssign(subscript(dst, i), cvtExp e) :: lp(i+1, es)
                    in
                      lp(0, es)
                    end
                | (Ty.DynSeqTy _, _) => simpleInit ()
                | (Ty.ImageTy _, _) => simpleInit ()
                | _ => let
                    val (tmp, stm) = TreeToC.expToVar(TreeIL.Var.Map.empty, CTyTranslate.toType ty, "tmp", src)
                    in
                      stm @ CTyTranslate.copyToC{ty = ty, dst=dst, src=tmp}
                    end
              (* end case *)
            end
    in
  (* generate code to initialize the default input values *)
    fun genInitDefaults (tgt, inputs) = let
        (* the inputs pointer type *)
          val inputPtrTy = CL.T_Ptr(CL.T_Named(N.inputsTy tgt))
        (* some common variables *)
          val inpV = CL.mkVar "inp"
        (* initialize a given input *)
          fun initInput ((gv, name, _, SOME dflt), stms) = let
                val lhs = CL.mkIndirect(inpV, GVar.name gv)
                in
                  initC {ty=GVar.ty gv, dst=lhs, src=dflt} @ stms
                end
            | initInput (_, stms) = stms
          in
            CL.D_Func(
              ["static"], CL.voidTy, N.initDefaults,
              [CL.PARAM([], inputPtrTy, "inp")],
              CL.mkBlock(List.foldr initInput [] inputs))
          end
    end (* local *)

  (* generate code to register command-line options for setting the input variables *)
    fun genRegisterInputs (tgt, inputs) = let
        (* the inputs pointer type *)
          val inputPtrTy = CL.T_Ptr(CL.T_Named(N.inputsTy tgt))
        (* some common variables *)
          val inpV = CL.mkVar "inp"
          val optsV = CL.mkVar "opts"
        (* register a given input *)
          fun registerInput (gv, name, desc, optDflt) = CL.mkCall(N.input(GVar.ty gv), [
                  optsV, CL.mkStr name, CL.mkStr(Option.getOpt(desc, "")),
                  if TrTy.isCArrayTy(GVar.ty gv)
                    then CL.mkIndirect(inpV, GVar.name gv)
                    else CL.mkUnOp(CL.%&, CL.mkIndirect(inpV, GVar.name gv)),
                  CL.mkBool(Option.isSome optDflt)
                ])
          in
            CL.D_Func(
              ["static"], CL.voidTy, N.registerOpts,
              [CL.PARAM([], inputPtrTy, "inp"), CL.PARAM([], CL.T_Ptr(CL.T_Named N.optionsTy), "opts")],
              CL.mkBlock(List.map registerInput inputs))
          end

  (* generate code to initialize the global input variables from the command-line inputs *)
    fun genInitInputs (tgt, inputs) = let
        (* the world pointer type *)
          val worldPtrTy = CL.T_Ptr(CL.T_Named(N.worldTy tgt))
        (* the global state pointer type *)
          val globPtrTy = CL.T_Ptr(CL.T_Named(N.globalsTy tgt))
        (* the inputs pointer type *)
          val inputPtrTy = CL.T_Ptr(CL.T_Named(N.inputsTy tgt))
        (* some common variables *)
          val inpV = CL.mkVar "inp"
          val optsV = CL.mkVar "opts"
        (* initialize a given input global; for sequences and images, this requires
         * loading the value from the specified nrrd file, while for other types
         * we just copy the values.
         *)
          fun initInput ((gv, name, desc, optDflt), stms) = (case GVar.ty gv
                 of Ty.DynSeqTy elemTy => let
                      val (loadFn, nDims, dims) = (case elemTy
                             of Ty.BoolTy => ("Diderot_DynSeqLoadBoolFromFile", CL.mkInt 0, CL.mkInt 0)
                              | Ty.IntTy => ("Diderot_DynSeqLoadIntFromFile", CL.mkInt 0, CL.mkInt 0)
                              | Ty.TensorTy[] => ("Diderot_DynSeqLoadRealFromFile", CL.mkInt 0, CL.mkInt 0)
                              | Ty.TensorTy _ => raise Fail "TODO: sequences of tensors"
                              | Ty.SeqTy elemTy => raise Fail "TODO: sequences of sequences"
                              | _ => raise Fail "unsupported dynamic sequence type"
                            (* end case *))
                      in
                        CL.mkAssign(global gv,
                          CL.mkApply(loadFn, [wrldPrefix, CL.mkIndirect(inpV, GVar.name gv), nDims, dims])) ::
                        CL.mkIfThen(CL.mkBinOp(global gv, CL.#==, CL.mkInt 0),
                          CL.mkReturn(SOME(CL.mkVar "true"))) :: stms
                      end
                  | Ty.ImageTy info => let
                      val loadFn = (case ImageInfo.dim info
                             of 1 => "Diderot_LoadImage1D"
                              | 2 => "Diderot_LoadImage2D"
                              | 3 => "Diderot_LoadImage3D"
                              | _ => raise Fail "image with dimension > 3"
                            (* end case *))
                      in
                        CL.mkIfThen(
                          CL.mkApply(loadFn, [wrldPrefix, CL.mkIndirect(inpV, GVar.name gv), CL.mkAddrOf(global gv)]),
                          CL.mkReturn(SOME(CL.mkVar "true"))) :: stms
                      end
                  | ty => TrTy.copyFromC{ty=ty, dst=global gv, src=CL.mkIndirect(inpV, GVar.name gv)} @ stms
                (* end case *))
          in
            CL.D_Func(
              ["static"], CL.boolTy, N.initInputs,
              [CL.PARAM([], worldPtrTy, "wrld"), CL.PARAM([], inputPtrTy, "inp")],
              CL.mkBlock(
                CL.mkDeclInit(globPtrTy, "glob", CL.mkIndirect(CL.mkVar "wrld", "globals")) ::
                List.foldr initInput [CL.mkReturn(SOME(CL.mkVar "false"))] inputs))
          end

  (* generate the functions that handle inputs for standalone executables.  These are:
   *    InitDefaults    -- called to initialize the default input values
   *    RegisterInputs  -- called to register the command-line options for the input globals
   *    InitInptus      -- called to initialize the input globals from the values specified
   *                       on the command line.
   *)
    fun genExecInputFuns (_, []) = []
      | genExecInputFuns arg = [
            genInitDefaults arg,
            genRegisterInputs arg,
            genInitInputs arg
          ]

  (* generate the typedef for the defined-input flag struct. *)
    fun genDefinedInpStruct (_, []) = []
      | genDefinedInpStruct (tgt : props, inputs) = let
          fun mkField (gv, _, _, _) = (CL.boolTy, GVar.name gv)
          in
            [CL.D_StructDef(NONE, List.map mkField inputs, SOME(N.definedInpTy tgt))]
          end

    fun genDefineInp (tgt, inputs) = let
        (* the world pointer type *)
          val worldPtrTy = CL.T_Ptr(CL.T_Named(N.worldTy tgt))
          val wrldParam = CL.PARAM([], worldPtrTy, "wrld")
          fun initFlag (gv, _, _, _) = CL.mkAssign(defined gv, CL.mkBool false)
          in
            CL.D_Func(
              ["static"], CL.voidTy, N.initDefined tgt,
              [wrldParam],
              CL.mkBlock(List.map initFlag inputs))
          end

    fun genCheckInputs (tgt, inputs) = let
        (* the world pointer type *)
          val worldPtrTy = CL.T_Ptr(CL.T_Named(N.worldTy tgt))
          val wrldParam = CL.PARAM([], worldPtrTy, "wrld")
        (* the inputs pointer type *)
          fun check (gv, name, _, optDflt) = let
                val dfltStm = (case optDflt
                       of SOME v => let
                          (* get default file name for images and sequences *)
                            fun getName () = (case v
                                   of IL.E_Lit(Literal.String s) => CL.mkStr s
                                    | _ => raise Fail "expected string for default nrrd"
                                  (* end case *))
                            in
                              case (GVar.ty gv)
                               of Ty.DynSeqTy elemTy =>
                                    GenLoadNrrd.loadSeqFromFile (global gv, elemTy, getName())
                                | Ty.ImageTy info =>
                                    GenLoadNrrd.loadImage (global gv, info, getName())
                                | _ => CL.mkBlock(ToC.trAssign(ToC.empty, global gv, v))
                              (* end case *)
                            end
                        | NONE => CL.mkBlock[
                              World.errorMsgAdd(CL.mkStr(concat["undefined input \"", name, "\"\n"])),
                              CL.mkReturn(SOME(CL.mkBool true))
                            ]
                      (* end case *))
                in
                  CL.mkIfThen(CL.mkUnOp(CL.%!, defined gv), dfltStm)
                end
          in
            CL.D_Func(
              ["static"], CL.boolTy, N.checkDefined tgt,
              [wrldParam],
              CL.mkBlock(List.map check inputs @ [CL.mkReturn(SOME(CL.mkBool false))]))
          end

  (* for each input variable we generate two or three top-level declarations in the
   * exported API.
   *)
    fun genInputFuns (tgt, []) = []
      | genInputFuns (tgt, inputs) = let
        (* the world pointer type *)
          val worldPtrTy = CL.T_Ptr(CL.T_Named(N.worldTy tgt))
          val wrldParam = CL.PARAM([], worldPtrTy, "wrld")
        (* create decls for an input variable *)
          fun mkInputDecls (gv, name, desc, optDflt) = let
                val ty = GVar.ty gv
              (* create a description declaration for the input variable *)
                val descDcl = (case desc
                       of SOME desc => [
                              CL.D_Var([], CL.T_Ptr(CL.T_Named "const char"),
                                N.inputDesc(tgt, name),
                                SOME(CL.I_Exp(CL.mkStr desc)))
                            ]
                        | NONE => []
                      (* end case *))
                val getDcl = if (Option.isSome optDflt)
                        then let
                          val getName = N.inputGet(tgt, name)
                        (* generate code for a function that returns the current value of
                         * an input global.
                         *)
                          fun mkFunc () = let
                                val isArray = TrTy.isCArrayTy ty
                                val (outTy, dst) = if isArray
                                      then (trType ty, CL.mkVar "v")
                                      else (CL.T_Ptr(trType ty), CL.mkUnOp(CL.%*, CL.mkVar "v"))
                                in
                                  CL.D_Func([], CL.voidTy, getName,
                                    [wrldParam, CL.PARAM([], outTy, "v")],
                                    CL.mkBlock(TrTy.copyToC{ty=ty, dst=dst, src=global gv}))
                                end
(* FIXME: for images and sequences, it is not clear what the get function should return.
 * For now, we just return 0
 *)
                        (* for images and sequences, we currently return 0 *)
                          fun nrrdFunc () = CL.D_Func([], CL.voidTy, getName,
                                [wrldParam, CL.PARAM([], CL.T_Ptr CL.voidPtr, "v")],
                                CL.mkAssign(CL.mkUnOp(CL.%*, CL.mkVar "v"), CL.mkInt 0))
                          val func = (case ty
                                 of Ty.DynSeqTy _ => nrrdFunc ()
                                  | Ty.ImageTy _ => nrrdFunc ()
                                  | _ => mkFunc ()
                                (* end case *))
                          in
                            [func]
                          end
                        else []
                val setDcl = (case ty
                       of Ty.ImageTy info => [
                              CL.D_Func(
                                [], CL.boolTy, N.inputSetByName(tgt, name),
                                [wrldParam, CL.PARAM(["const"], CL.charPtr, "s")],
                                CL.appendStm(
                                  GenLoadNrrd.loadImage (global gv, info, CL.mkVar "s"),
                                  CL.mkReturn(SOME(CL.mkBool false)))),
                              CL.D_Func(
                                [], CL.boolTy, N.inputSet(tgt, name),
                                [wrldParam, CL.PARAM([], nrrdPtrTy, "nin")],
                                CL.appendStm(
                                  GenLoadNrrd.setImage (global gv, info, CL.mkVar "nin"),
                                  CL.mkReturn(SOME(CL.mkBool false))))
                            ]
                        | Ty.DynSeqTy elemTy => [
                              CL.D_Func(
                                [], CL.boolTy, N.inputSetByName(tgt, name),
                                [wrldParam, CL.PARAM(["const"], CL.charPtr, "s")],
                                CL.appendStm(
                                  GenLoadNrrd.loadSeqFromFile (global gv, elemTy, CL.mkVar "s"),
                                  CL.mkReturn(SOME(CL.mkBool false)))),
                              CL.D_Func(
                                [], CL.boolTy, N.inputSet(tgt, name),
                                [wrldParam, CL.PARAM([], nrrdPtrTy, "nin")],
                                CL.appendStm(
                                  GenLoadNrrd.loadSeq (global gv, elemTy, CL.mkVar "nin"),
                                  CL.mkReturn(SOME(CL.mkBool false))))
                            ]
                        | _ => [
                              CL.D_Func(
                                [], CL.boolTy, N.inputSet(tgt, name),
                                [wrldParam, CL.PARAM([], trType ty, "v")],
                                CL.mkBlock(
                                  CL.mkAssign(defined gv, CL.mkBool true) ::
                                  TrTy.copyFromC{ty=ty, dst=global gv, src=CL.mkVar "v"} @
                                    [CL.mkReturn(SOME(CL.mkVar "false"))]))
                            ]
                      (* end case *))
                in
                  (descDcl @ getDcl @ setDcl)
                end
          val extras = [
                  genCheckInputs (tgt, inputs),
                  genDefineInp (tgt, inputs)
                ]
          in
            List.foldr (fn (input, dcls) => mkInputDecls input @ dcls) extras inputs
          end

  end