(* gen-output.sml
 *
 * COPYRIGHT (c) 2012 The Diderot Project (http://diderot-language.cs.uchicago.edu)
 * All rights reserved.
 *
 * Generate strand output functions.  There are four cases that we handle:
 *
 *	grid, fixed-size elements:
 *		nrrd has object dimensions followed by axes
 *
 *	collection, fixed-size elements
 *		nrrd has object dims followed by a single axis
 *
 *	grid, dynamic-size elements
 *		lengths nrrd has size 2 for objects (offset, length) followed by axes
 *		data nrrd has object dimensions followed by a single axis
 *
 *	collection, dynamic-size elements
 *		lengths nrrd has size 2 for objects (offset, length) followed by a single axis
 *		data nrrd has object dimensions followed by a single axis
 *)

structure GenOutput : sig

    val gen : TargetUtil.target_desc * int option -> (TreeIL.Ty.ty * string) -> CLang.decl

  end = struct

    structure IL = TreeIL
    structure V = IL.Var
    structure Ty = IL.Ty
    structure CL = CLang
    structure N = CNames

    fun mapi f l = let
	  fun mapf (i, [], l) = List.rev l
	    | mapf (i, x::xs, l) = mapf (i+1, xs, f(i, x)::l)
	  in
	    mapf (0, l, [])
	  end

    val nrrdPtrTy = CL.T_Ptr(CL.T_Named "Nrrd")
    val sizeTy = CL.T_Named "size_t"

    fun mkInt i = CL.mkInt(IntInf.fromInt i)

    fun infoOf (tgt : TargetUtil.target_desc, ty) = (case ty
           of Ty.IntTy => if #longint tgt
		then (CL.int64, "nrrdTypeLLong", [1])
		else (CL.int32, "nrrdTypeInt", [1])
            | Ty.TensorTy [] => if #double tgt
		then (CL.double, "nrrdTypeDouble", [1])
		else (CL.float, "nrrdTypeFloat", [1])
            | Ty.TensorTy dims => if #double tgt
		then (CL.double, "nrrdTypeDouble", dims)
		else (CL.float, "nrrdTypeFloat", dims)
            | Ty.SeqTy(ty, n) => let
		val (elemTy, nrrdTy, dims) = infoOf (tgt, ty)
		in
		  (elemTy, nrrdTy, n::dims)
		end
            | _ => raise Fail(concat["GetOutput.infoOf(", Ty.toString ty, ")"])
	  (* end case *))

  (* variables in the generated code *)
    val wrldV = CL.mkVar "wrld"
    val sizesV = CL.mkVar "sizes"
    val iV = CL.mkVar "i"
    val nV = CL.mkVar "n"
    val pV = CL.mkVar "p"
    val offsetV = CL.mkVar "offset"
    val dataV = CL.mkVar "data"
    val lengthsV = CL.mkVar "lengths"
    val numStableV = CL.mkVar "numStable"
    val numElemsV = CL.mkVar "numElems"
    val DIDEROT_STABLE = CL.mkVar "DIDEROT_STABLE"
    val NRRD = CL.mkVar "NRRD"

  (* dymanic sequence operations *)
    fun seqLength arg = CL.mkApply("Diderot_DynSeqLength", [arg])
    fun seqCopy (elemSz, dst, seq) = CL.mkApply("Diderot_DynSeqCopy", [elemSz, dst, seq])

  (* utility functions for initializing the sizes array *)
    fun sizes i = CL.mkSubscript(sizesV, mkInt i)
    fun setSizes (i, v) = CL.mkAssign(sizes i, v)

  (* code to access state variable
	wrld->outState[i].name
   *)
    fun stateVar name = CL.mkIndirect(CL.mkSubscript(CL.mkIndirect(wrldV, "outState"), iV), name)

  (* code fragment to loop over strands
	for (unsigned int i = 0;  i < wrld->numStrands;  i++) ...
   *)
    fun forStrands stm = CL.mkFor(
	  [(CL.uint32, "i", mkInt 0)],
	  CL.mkBinOp(iV, CL.#<, CL.mkIndirect(wrldV, "numStrands")),
	  [CL.mkPostOp(numStableV, CL.^++)],
	  stm)

  (* code fragment to test for stable strands in a loop
	if (wrld->status[i] == DIDEROT_STABLE)
	    ...
   *)
    fun ifStable stm = CL.mkIfThen(
	  CL.mkBinOp(CL.mkSubscript(CL.mkIndirect(wrldV, "status"), iV), CL.#==, DIDEROT_STABLE),
	  stm)

  (* code fragment to allocate nrrd data and check for errors
	if (nrrdMaybeAlloc_nva(<nrrdVar>, <nrrdType>, <nDims>, sizes) != 0) {
	    wrld->err = biffGetDone(NRRD);
	    return true;
	}
   *)
    fun maybeAlloc (nrrdVar, nrrdType, nDims) = 
	  CL.mkIfThen(
	    CL.mkBinOp(
	      CL.mkApply("nrrdMaybeAlloc_nva", [
		  nrrdVar, CL.mkVar nrrdType, mkInt nDims, sizesV
		]),
	      CL.#!=,
	      CL.mkInt 0),
	  (* then *)
	    CL.mkBlock[
		CL.mkAssign(CL.mkIndirect(wrldV, "err"),
		  CL.mkApply("biffGetDone", [NRRD])),
		CL.mkReturn(SOME(CL.mkVar "true"))
	      ]
	  (* endif*))

  (* create the body of an output function for dynamic-size outputs.  The structure of the
   * function body is:
   *
   *	declarations
   *    compute sizes array for lengths
   *	allocate nrrd for lengths
   *    compute sizes array for data
   *	allocate nrrd for data
   *	copy data from strands to nrrd
   *)
    fun genDynOutput (tgt, nAxes, ty, name) = let
	  val (elemCTy, nrrdType, dims) = infoOf (tgt, ty)
	  val nDims = List.length dims
	  val nElems = List.foldl Int.* 1 dims
	  val (isArray, nAxes) = (case nAxes of NONE => (false, 1) | SOME n => (true, n))
	(* declarations *)
	  val decls = [
		  CL.mkDecl(CL.T_Array(sizeTy, SOME(nDims+nAxes)), "sizes", NONE),
		  CL.mkDecl(CL.charPtr, "p", NONE)
		]
	(* count number of elements (and stable strands) *)
	  val countElems = let
		val nElemsInit = CL.mkAssign(numElemsV, CL.mkInt 0)
		val cntElems = CL.S_Exp(CL.mkAssignOp(numElemsV, CL.+=, seqLength(stateVar name)))
		in
		  if isArray
		    then [
			CL.mkComment["count number of elements"],
			nElemsInit, forStrands cntElems
		      ]
		    else [
			CL.mkComment["count number of output elements and stable strands"],
			CL.mkAssign(numStableV, CL.mkInt 0), nElemsInit,
			forStrands(ifStable(CL.mkBlock[cntElems, CL.S_Exp(CL.mkPostOp(numStableV, CL.^++))]))
		      ]
		end
	(* generate code to allocate the lengths nrrd *)
	  val lengthsNrrd = let
		val dimSizes = setSizes(0, CL.mkInt 2)  (* lengths is 2-element vector *)
		in
		  CL.mkComment["allocate lengths nrrd"] ::
		  (if isArray
		    then dimSizes ::
		      List.tabulate (nAxes, fn i =>
			setSizes(i+1, CL.mkSubscript(CL.mkIndirect(wrldV, "sizes"), mkInt(nAxes-i-1)))) @
		      [maybeAlloc (lengthsV, "nrrdTypeInt", nAxes+1)]
		    else [
			dimSizes, setSizes(nDims, numStableV),
			maybeAlloc (lengthsV, "nrrdTypeInt", 2)
		      ])
		end
	(* generate code to allocate the data nrrd *)
	  val dataNrrd = let
		val dimSizes = mapi (fn (i, d) => setSizes(i, mkInt d)) dims
		val sizes = setSizes(nDims, numElemsV)
		in
		  dimSizes @ [
		      sizes,
		      maybeAlloc (dataV, nrrdType, nDims+1)
		    ]
		end
	(* generate the lengths copy code *)
	  val copyLengths = let
		val pInit = CL.mkAssign(pV,
		      CL.mkCast(CL.T_Ptr(CL.uint32), CL.mkIndirect(lengthsV, "data")))
		val offsetDecl = CL.mkDecl(CL.uint32, "offset", SOME(CL.I_Exp(CL.mkInt 0)))
		val copyBlk = CL.mkBlock[
			CL.mkDecl(CL.uint32, "n", SOME(CL.I_Exp(seqLength(stateVar name)))),
			CL.mkAssign(CL.mkPostOp(pV, CL.^++), offsetV),
			CL.mkAssign(CL.mkPostOp(pV, CL.^++), nV),
			CL.S_Exp(CL.mkAssignOp(offsetV, CL.+=, nV))
		      ]
		val copyStm = if isArray
		      then copyBlk
		      else ifStable copyBlk
		in
		  [pInit, offsetDecl, copyStm]
		end
	(* generate the data copy code *)
	  val copyData = let
		val pInit = CL.mkAssign(pV, CL.mkCast(CL.charPtr, CL.mkIndirect(dataV, "data")))
		val copyStm = CL.mkAssign(pV, seqCopy(
		      CL.mkBinOp(mkInt nElems, CL.#*, CL.mkSizeof(elemCTy)), pV, stateVar name))
		val copyStm = if isArray
		      then copyStm
		      else ifStable copyStm
		in
		  [pInit, forStrands copyStm]
		end
	(* the function body *)
	  val stms =
		decls @
		countElems @
		lengthsNrrd @
		dataNrrd @
		copyLengths @
		copyData @
		[CL.mkReturn(SOME(CL.mkVar "false"))]	  
	  in
	    ([CL.PARAM([], nrrdPtrTy, "lengths"), CL.PARAM([], nrrdPtrTy, "data")], CL.mkBlock stms)
	  end

  (* create the body of an output function for fixed-size outputs.  The structure of the
   * function body is:
   *
   *	declare and compute sizes array
   *	allocate nrrd data
   *	copy data from strands to nrrd
   *)
    fun genFixedOutput (tgt, nAxes, ty, name) = let
	  val (elemCTy, nrrdType, dims) = infoOf (tgt, ty)
	  val nDims = List.length dims
	  val nElems = List.foldl Int.* 1 dims
	  val (isArray, nAxes) = (case nAxes of NONE => (false, 1) | SOME n => (true, n))
	(* generate the sizes initialization code *)
	  val initSizes = let
		val sizesDcl = CL.mkDecl(CL.T_Array(sizeTy, SOME(nDims+nAxes)), "sizes", NONE)
		val dimSizes = sizesDcl :: mapi (fn (i, d) => setSizes(i, mkInt d)) dims
		in
		  if isArray
		    then dimSizes @
		      List.tabulate (nAxes, fn i =>
			setSizes(nDims+i, CL.mkSubscript(CL.mkIndirect(wrldV, "sizes"), mkInt(nAxes-i-1))))
		    else
		      CL.mkDecl(sizeTy, "numStable", SOME(CL.I_Exp(mkInt 0))) ::
		      forStrands (ifStable(CL.S_Exp(CL.mkPostOp(numStableV, CL.^++)))) ::
		      dimSizes @ [setSizes(nDims, numStableV)]
		end
	(* generate the copy code *)
	  val copyCode = let
		val pDecl = CL.mkDecl(CL.charPtr, "p", SOME(CL.I_Exp(
			CL.mkCast(CL.charPtr, CL.mkIndirect(dataV, "data")))))
		val copyBlk = CL.mkBlock[
			CL.mkCall("memcpy", [
			    pV,
			    CL.mkUnOp(CL.%&, stateVar name),
			    CL.mkBinOp(mkInt nDims, CL.#*, CL.mkSizeof elemCTy)
			  ]),
			CL.mkExpStm(CL.mkAssignOp(pV, CL.+=, CL.mkSizeof elemCTy))
		      ]
		val copyStm = if isArray
		      then copyBlk
		      else ifStable copyBlk
		in
		  [pDecl, forStrands copyStm]
		end
	(* the function body *)
	  val stms =
		CL.mkComment["Compute sizes of nrrd file"] ::
		CL.mkDecl(sizeTy, "sizes", NONE) ::
		initSizes @
		CL.mkComment["Allocate nrrd data"] ::
		maybeAlloc (dataV, nrrdType, nDims+nAxes) ::
		CL.mkComment["copy data to output nrrd"] ::
		copyCode @
		[CL.mkReturn(SOME(CL.mkVar "false"))]
	  in
	    ([CL.PARAM([], nrrdPtrTy, "data")], CL.mkBlock stms)
	  end

    fun gen (tgt : TargetUtil.target_desc, nAxes) (ty, name) = let
	  val funcName = String.concat[#namespace tgt, "Get_", name]
	  val wrldPtr = CL.T_Ptr(CL.T_Named(#namespace tgt ^ "World_t"))
	  fun mkFunc (params, body) =
		CL.D_Func([], CL.boolTy, funcName, CL.PARAM([], wrldPtr, "wrld")::params, body)
	  in
	    case ty
	     of Ty.DynSeqTy ty' => mkFunc (genDynOutput(tgt, nAxes, ty', name))
	      | _ => mkFunc (genFixedOutput(tgt, nAxes, ty, name))
	    (* end case *)
	  end

  end