(* simple.sml
 *
 * COPYRIGHT (c) 2013 The Diderot Project (http://diderot-language.cs.uchicago.edu)
 * All rights reserved.
 *
 * A simplified AST representation of a Diderot program.  This representation has the property
 * that the arguments to ifs, operators, etc. are variables and that the rhs of assignments
 * consist of a single operation.  It is not, however, a single-assignment representation.
 *)

structure Simple =
  struct

  (* types with the meta variables resolved *)
    type ty = SimpleTypes.ty

  (* resolved meta-variable arguments to basis functions *)
    datatype meta_arg = datatype SimpleTypes.meta_arg

    datatype var_kind = datatype AST.var_kind

    datatype var = V of {
        name : string,          (* print name of variable *)
        id : Stamp.stamp,       (* unique ID *)
        kind : var_kind,        (* variable kind *)
        ty : ty                 (* type *)
      }

    datatype program = Program of {
        props : StrandUtil.program_prop list,
        inputs : (var * SimpleTypes.ty Inputs.input) list,	(* input globals *)
        globals : var list,		(* non-input globals *)
        globalInit : block,
        funcs : func list,
        strands : strand list,
        init : init                     (* block evaluates any vars used in init *)
      }

    and func = Func of {
        f : var,
        params : var list,
        body : block
      }

    and init = Initially of {
        isArray : bool,
        rangeInit : block,
        iters : {param : var, lo : var, hi : var} list,
        create : create
      }

    and create = C_Create of {
        argInit : block,
        name : Atom.atom,
        args : var list
      }

(* do we use this?
    and iter = I_Range of {
        rangeInit : block,
        param : var,
        lo : var,
        hi : var
      }
*)

    and strand = Strand of {
        name : Atom.atom,
        params : var list,
        state : var list,
        stateInit : block,
        methods : method list
      }

    and method = Method of StrandUtil.method_name * block

    and block = Block of stmt list

    and stmt
      = S_Var of var                            (* introduce an uninitialized local variable. *)
                                                (* These stmts are needed for the results of *)
                                                (* conditional expressions *)
      | S_Assign of var * exp
      | S_IfThenElse of var * block * block
      | S_New of Atom.atom * var list
      | S_Die
      | S_Stabilize
      | S_Return of var
      | S_Print of var list

    and exp
      = E_Var of var
      | E_Lit of Literal.literal
      | E_Tuple of var list
      | E_Apply of var * var list * ty                          (* user-defined function *)
      | E_Prim of AST.var * meta_arg list * var list * ty       (* Diderot builtin *)
      | E_Cons of var list
      | E_Seq of var list * ty                                  (* sequence (ty is result type) *)
      | E_Slice of var * var option list * ty                   (* tensor slicing (ty is result type) *)
      | E_Coerce of {srcTy : ty, dstTy : ty, x : var}
      | E_LoadSeq of ty * string
      | E_LoadImage of ty * string * ImageInfo.info

    fun typeOf (E_Var(V{ty, ...})) = ty
      | typeOf (E_Lit lit) = (case lit
           of (Literal.Int _) => SimpleTypes.T_Int
            | (Literal.Float _) => SimpleTypes.T_Tensor[]
            | (Literal.String s) => SimpleTypes.T_String
            | (Literal.Bool _) => SimpleTypes.T_Bool
          (* end case *))
      | typeOf (E_Tuple _) = raise Fail "E_Tuple"
      | typeOf (E_Apply(_, _, ty)) = ty
      | typeOf (E_Prim(_, _, _, ty)) = ty
      | typeOf (E_Cons[]) = raise Fail "impossible empty E_Cons"
      | typeOf (E_Cons(V{ty, ...}::xs)) = let
          val d = List.length xs + 1
          in
            case ty
             of SimpleTypes.T_Tensor shape => SimpleTypes.T_Tensor(shape @ [d])
              | _ => raise Fail(concat[
                    "element of tensor construction is ", SimpleTypes.toString ty,
                    ", expected tensor"
                  ])
            (* end case *)
          end
      | typeOf (E_Seq(_, ty)) = ty
      | typeOf (E_Slice(_, _, ty)) = ty
      | typeOf (E_Coerce{dstTy, ...}) = dstTy
      | typeOf (E_LoadSeq(ty, _)) = ty
      | typeOf (E_LoadImage(ty, _, _)) = ty

  end