(* types.sml
 *
 * COPYRIGHT (c) 2010 The Diderot Project (http://diderot.cs.uchicago.edu)
 * All rights reserved.
 *
 * This file defines the internal representation of Diderot types.  Most values
 * have the tensor type, which includes scalars, vectors, etc.  The internal
 * type system supports polymorphism, which is used to give types to higher-order
 * operators (e.g., derivatives).
 *)

structure Types =
  struct

  (* kinds for type variables *)
    datatype kind
      = TK_DIM		(* ranges over dimensions (1, 2, ...) *)
      | TK_INT		(* ranges over integer types *)
      | TK_FLT		(* ranges over floating-point types *)
      | TK_RAW		(* ranges over raw (scalar) types *)
      | TK_TYPE		(* ranges over types *)
(* Question: do we want kinds for tensors of different orders (e.g., TK_TENSOR of word?);
 * Then TK_RAW would be TK_TENSOR of 0w0.
 *)

  (* raw numeric types as supported by NRRD *)
    datatype raw_ty
      = RT_Int8 | RT_UInt8
      | RT_Int16 | RT_UInt16
      | RT_Int32 | RT_UInt32
      | RT_Int64 | RT_UInt64
      | RT_Float | RT_Double
      | RT_Var of var

    and ty
      = T_Var of var
      | T_Bool
    (* scalars, vectors, matrices, etc. *)
(* Question: perhaps we want {order : word, dim : dim, ty : raw_ty} instead *)
      | T_Tensor of {order : dim list, ty : raw_ty}
(* Question: change "data set" to "image"? *)
    (* data sets from NRRD *)
      | T_Data of {
	  dim : dim,		(* 2D or 3D data set *)
	  shape : dim list,	(* tensor shape; order is length of list *)
	  repTy : raw_ty	(* representation type of elements (raw kind) *)
	}
    (* continuous field reconstructed from a data set *)
      | T_Field of {
	  diff : dim,		(* number of levels of differentiation supported *)
	  dim : dim,		(* 2D or 3D field *)
	  shape : dim list,	(* tensor shape; order is length of list *)
	  repTy : raw_ty,	(* representation type of elements (raw kind) *)
	  ty : raw_ty		(* type of samples (float kind) *)
	}

    and dim = DIM of word | DIMVAR of var

  (* type variables; the kind field restricts their range *)
    and var = TV of {
	  kind : kind,
	  stamp : Stamp.stamp
	}

  (* given a dataset or field type, return the element type, which will be
   * a tensor type.
   *)
    fun elemTypeOf (T_Data{order, repTy, ...}) = T_Tensor{order=order, ty=repTy}
      | elemTypeOf (T_Field{order, ty, ...}) = T_Tensor{order=order, ty=ty}
      | elemTypeOf _ = raise Fail "not a dataset/field"

  (* is a type well-formed? *)
    fun wellFormed ty = let
	(* is a dimension well-formed? *)
	  fun okDim (DIM w) = (w > 0w0)
	    | okDim (DIMVAR(TV{kind=TK_DIM, ...})) = true
	    | okDim _ = false
	(* is a raw type well formed *)
	  fun okRawTy (RT_Var(TV{kind=TK_INT, ...})) = true
	    | okRawTy (RT_Var(TV{kind=TK_FLT, ...})) = true
	    | okRawTy (RT_Var(TV{kind=TK_RAW, ...})) = true
	    | okRawTy (RT_Var _) = false
	    | okRawTy _ = true
	(* is a raw type a floating-point type? *)
	  fun okFloatTy (RT_Var(TV{kind=TK_FLT, ...})) = true
	    | okFloatTy RT_Float = true
	    | okFloatTy RT_Double = true
	    | okFloatTy _ = false
	  fun wellFormed' (T_Var(TV{kind=TK_TYPE, ...})) = true
	    | wellFormed' (T_Var _) = false
	    | wellFormed' (T_Tensor{order, ty}) =
		List.all okDim order
		andalso okRawTy ty
	    | wellFormed' (T_Data{dim, order, repTy}) =
		okDim dim
		andalso List.all okDim order
		andalso okRawTy repTy
	    | wellFormed' (T_Field{dim, order, repTy, ty}) =
		okDim dim
		andalso List.all okDim order
		andalso okRawTy repTy
		andalso okFloatTy ty
	    | wellFormed' _ = false
	  in
	    wellFormed' ty
	  end

  (* some common types *)
    local
      fun scalarTy rty = T_Tensor{order=[], ty=rty}
      fun vecTy (d, rty) = T_Tensor{order=[DIM d], ty=rty}
    in
    val intTy = scalarTy RT_Int32
    val int2Ty = vecTy (0w2, RT_Int32)
    val intT3y = vecTy (0w3, RT_Int32)
    val intT4y = vecTy (0w4, RT_Int32)
    val uintTy = scalarTy RT_UInt32
    val uint2Ty = vecTy (0w2, RT_UInt32)
    val uintT3y = vecTy (0w3, RT_UInt32)
    val uintT4y = vecTy (0w4, RT_UInt32)
    val floatTy = scalarTy RT_Float
    val float2Ty = vecTy (0w2, RT_Float)
    val floatT3y = vecTy (0w3, RT_Float)
    val floatT4y = vecTy (0w4, RT_Float)
    val doubleTy = scalarTy RT_Double
    val double2Ty = vecTy (0w2, RT_Double)
    val double3Ty = vecTy (0w3, RT_Double)
    val double4Ty = vecTy (0w4, RT_Double)
    end (* local *)

  (* new type variables *)
    fun newVar k = TV{kind=k, stamp=Stamp.new()}

  end