(* tree-il.sml
 *
 * COPYRIGHT (c) 2011 The Diderot Project (http://diderot-language.cs.uchicago.edu)
 * All rights reserved.
 *
 * This representation restores the block structure and nested expression syntax
 * of the source language.
 *)

structure TreeIL =
  struct

    structure Op = LowOps
    structure Ty = LowILTypes

    datatype program = Program of {
	globals : var list,
	globalInit : block,
	strands : strand list
	(* initialization *)
      }

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

    and method = Method of {
	name : Atom.atom,
	body : block		(* method body *)
      }

    and block = Block of {
	locals : var list,
	body : stm list
      }

    and stm
      = S_Comment of string list
      | S_Assign of var * exp
      | S_IfThen of exp * block
      | S_IfThenElse of exp * block * block
      | S_For of var * exp * exp * block	(* limited for loop *)
    (* special Diderot forms *)
      | S_Cons of var * exp list	(* tensor-value construction *)
      | S_LoadVoxels of var * int * exp
      | S_LoadImage of var * int * exp	(* load image data *)
      | S_Input of var * string * exp option (* get input *)
      | S_Exit of exp list
    (* return functions for methods *)
      | S_Active of exp list
      | S_Stabilize of exp list
      | S_Die

    and exp
      = E_Var of var
      | E_Lit of Literal.literal
      | E_Op of Op.rator * exp list

    and var = V of {
	name : string,			(* name (should be unique) *)
	id : Stamp.stamp,		(* unique ID *)
	kind : var_kind,
	ty : Ty.ty			(* type *)
      }

    and var_kind
      = VK_Global			(* global variable *)
      | VK_State of Atom.atom		(* strand state variable *)
      | VK_Local			(* includes strand parameters *)

    structure Var : sig

	val kind : var -> var_kind
	val name : var -> string
	val ty : var -> Ty.ty

	structure Map : ORD_MAP where type Key.ord_key = var

      end = struct
	fun kind (V{kind, ...}) = kind
	fun name (V{name, ...}) = name
	fun ty (V{ty, ...}) = ty
	local
	  structure VarOrd =
	    struct
	      type ord_key = var
	      fun compare (V{id=a, ...}, V{id=b, ...}) = Stamp.compare(a, b)
	    end
	in
	structure Map = RedBlackMapFn (VarOrd)
	end (* local *)
      end

  end