(* ssa-fn.sml
 *
 * COPYRIGHT (c) 2010 The Diderot Project (http://diderot.cs.uchicago.edu)
 * All rights reserved.
 *
 * The IL is a combination of a block-structured tree and an SSA control-flow
 * graph of blocks.
 *)

signature SSA =
  sig

    structure Ty : SSA_TYPES
    structure Op : OPERATORS where type ty = Ty.ty

  (***** CFG *****)

    datatype node = ND of {
	id : Stamp.stamp,
	props : PropList.holder,
	kind : node_kind
      }

    and node_kind
      = NULL
      | ENTRY of {
	    succ : node ref
	  }
      | JOIN of {
	    preds : node list ref,
	    phis : (var * var list) list ref,	(* phi statements *)
	    succ : node ref
	  }
      | COND of {
	    pred : node ref,
	    cond : var,
	    trueBranch : node ref,
	    falseBranch : node ref
	  }
      | BLOCK of {
	    pred : node ref,
	    body : assign list ref,
	    succ : node ref
	  }
      | NEW of {
	    pred : node ref,
	    actor : Atom.atom,
	    args : var list,
	    succ : node ref
	  }
      | DIE of {
	    pred : node ref
	  }
      | STABILIZE of {
	    pred : node ref
	  }
      | EXIT of {
	    pred : node ref
	  }

  (***** Statements *****)

    and stmt = STM of {
	id : Stamp.stamp,
	props : PropList.holder,
	kind : stmt_kind,
	next : stmt option		(* next statement at this structural level *)
      }

    and stmt_kind
      = S_SIMPLE of node		(* ENTRY, JOIN, BLOCK, NEW, DIE, STABILIZE, or EXIT node *)
      | S_IF of {
	    cond : node,		(* COND node *)
	    thenBranch : stmt,
	    elseBranch : stmt
	  }
      | S_LOOP of {
	    hdr : stmt,
	    cond : node,		(* COND node *)
	    body : stmt
	  }

    and rhs
      = VAR of var
      | LIT of Literal.literal
      | OP of Op.rator * var list
      | CONS of var list		(* tensor-value construction *)

    and var = V of {
	name : string,			(* name *)
	id : Stamp.stamp,		(* unique ID *)
	ty : Ty.ty,			(* type *)
	bind : var_bind ref,		(* binding *)
	useCnt : int ref,		(* count of uses *)
	props : PropList.holder
      }

    and var_bind
      = VB_NONE
      | VB_RHS of rhs
      | VB_PHI of var list
      | VB_PARAM
      | VB_STATE_VAR

    withtype assign = (var * rhs)

    datatype program = Program of {
	globals : var list,
	globalInit : stmt,
	actors : actor list
	(* initialization *)
      }

    and actor = Actor of {
	name : Atom.atom,
	params : var list,
	state : var list,
	stateInit : stmt,
	methods : method list
      }

    and method = Method of {
	name : Atom.atom,
	stateIn : var list,	(* names of state variables on method entry *)
	stateOut : var list,	(* names of state variables on method exit *)
	body : stmt		(* method body *)
      }

    structure Node : sig
	val same : node * node -> bool
	val compare : node * node -> order
	val hash : node -> word
	val toString : node -> string
      (* dummy node *)
	val dummy : node
      (* CFG edges *)
	val preds : node -> node list
	val setPred : node * node -> unit
	val hasSucc : node -> bool
	val succs : node -> node list
	val setSucc : node * node -> unit
	val setTrueBranch : node * node -> unit	 (* set trueBranch successor for COND node *)
	val setFalseBranch : node * node -> unit (* set falseBranch successor for COND node *)
	val addEdge : node * node -> unit
      (* constructors *)
	val mkENTRY : unit -> node
	val mkJOIN : (var * var list) list -> node
	val mkCOND : {cond : var, trueBranch : node, falseBranch : node} -> node
	val mkBLOCK : assign list -> node
	val mkNEW : {actor : Atom.atom, args : var list} -> node
	val mkDIE : unit -> node
	val mkSTABILIZE : unit -> node
	val mkEXIT : unit -> node
      (* properties *)
	val newProp : (node -> 'a) -> {
		getFn : node -> 'a,
		peekFn : node -> 'a option,
		setFn : node * 'a -> unit,
		clrFn : node -> unit
	      }
	val newFlag : unit -> {
		getFn : node -> bool,
		setFn : node * bool -> unit
	      }
      end

    structure Stmt : sig
	val same : stmt * stmt -> bool
	val compare : stmt * stmt -> order
	val hash : stmt -> word
	val toString : stmt -> string
      (* return the entry node of the statement *)
	val entry : stmt -> node
      (* return the tail-end node of a statement (not applicable to S_IF or S_LOOP) *)
	val tail : stmt -> node
      (* statement constructor functions *)
	val new : (stmt_kind * stmt option) -> stmt
	val mkENTRY : stmt option -> stmt
	val mkJOIN : (var * var list) list * stmt option -> stmt
	val mkIF : var * stmt * stmt * stmt option -> stmt
	val mkBLOCK : assign list * stmt option -> stmt
	val mkNEW : Atom.atom * var list * stmt option -> stmt
	val mkDIE : unit -> stmt
	val mkSTABILIZE : unit -> stmt
	val mkEXIT : unit -> stmt
      (* properties *)
	val newProp : (stmt -> 'a) -> {
		getFn : stmt -> 'a,
		peekFn : stmt -> 'a option,
		setFn : stmt * 'a -> unit,
		clrFn : stmt -> unit
	      }
	val newFlag : unit -> {
		getFn : stmt -> bool,
		setFn : stmt * bool -> unit
	      }
      end

    structure Var : sig
	val new : string * Ty.ty -> var
	val name : var -> string
	val ty : var -> Ty.ty
	val binding : var -> var_bind
	val setBinding : var * var_bind -> unit
	val same : var * var -> bool
	val compare : var * var -> order
	val hash : var -> word
	val toString : var -> string
      (* properties *)
	val newProp : (var -> 'a) -> {
		getFn : var -> 'a,
		peekFn : var -> 'a option,
		setFn : var * 'a -> unit,
		clrFn : var -> unit
	      }
	val newFlag : unit -> {
		getFn : var -> bool,
		setFn : var * bool -> unit
	      }
      (* collections *)
	structure Map : ORD_MAP where type Key.ord_key = var
	structure Set : ORD_SET where type Key.ord_key = var
	structure Tbl : MONO_HASH_TABLE where type Key.hash_key = var
      end

  (* DFS sorting of the graph rooted at the entry to a statement; the resulting list will
   * be in preorder with parents before children.
   *)
    val sortNodes : stmt -> node list

  (* apply a function to all of the nodes in the graph rooted at the entry to the statement *)
    val applyToNodes : (node -> unit) -> stmt -> unit

  end

functor SSAFn (

    structure Ty : SSA_TYPES
    structure Op : OPERATORS where type ty = Ty.ty

  ) : SSA = struct

    structure Ty = Ty
    structure Op = Op

    datatype node = ND of {
	id : Stamp.stamp,
	props : PropList.holder,
	kind : node_kind
      }

    and node_kind
      = NULL
      | ENTRY of {
	    succ : node ref
	  }
      | JOIN of {
	    preds : node list ref,
	    phis : (var * var list) list ref,	(* phi statements *)
	    succ : node ref
	  }
      | COND of {
	    pred : node ref,
	    cond : var,
	    trueBranch : node ref,
	    falseBranch : node ref
	  }
      | BLOCK of {
	    pred : node ref,
	    body : assign list ref,
	    succ : node ref
	  }
      | NEW of {
	    pred : node ref,
	    actor : Atom.atom,
	    args : var list,
	    succ : node ref
	  }
      | DIE of {
	    pred : node ref
	  }
      | STABILIZE of {
	    pred : node ref
	  }
      | EXIT of {
	    pred : node ref
	  }

  (***** Statements *****)

    and stmt = STM of {
	id : Stamp.stamp,
	props : PropList.holder,
	kind : stmt_kind,
	next : stmt option		(* next statement at this structural level *)
      }

    and stmt_kind
      = S_SIMPLE of node		(* ENTRY, JOIN, BLOCK, NEW, DIE, STABILIZE, or EXIT node *)
      | S_IF of {
	    cond : node,		(* COND node *)
	    thenBranch : stmt,
	    elseBranch : stmt
	  }
      | S_LOOP of {
	    hdr : stmt,
	    cond : node,		(* COND node *)
	    body : stmt
	  }

    and rhs
      = VAR of var
      | LIT of Literal.literal
      | OP of Op.rator * var list
      | CONS of var list		(* tensor-value construction *)

    and var = V of {
	name : string,			(* name *)
	ty : Ty.ty,			(* type *)
	bind : var_bind ref,		(* binding *)
	id : Stamp.stamp,		(* unique ID *)
	useCnt : int ref,		(* count of uses *)
	props : PropList.holder
      }

    and var_bind
      = VB_NONE
      | VB_RHS of rhs
      | VB_PHI of var list
      | VB_PARAM
      | VB_STATE_VAR

    withtype assign = (var * rhs)

    datatype program = Program of {
	globals : var list,
	globalInit : stmt,
	actors : actor list
	(* initialization *)
      }

    and actor = Actor of {
	name : Atom.atom,
	params : var list,
	state : var list,
	stateInit : stmt,
	methods : method list
      }

    and method = Method of {
	name : Atom.atom,
	stateIn : var list,	(* names of state variables on method entry *)
	stateOut : var list,	(* names of state variables on method exit *)
	body : stmt		(* method body *)
      }

    structure Node =
      struct
	fun same (ND{id=a, ...}, ND{id=b, ...}) = Stamp.same(a, b)
	fun compare (ND{id=a, ...}, ND{id=b, ...}) = Stamp.compare(a, b)
	fun hash (ND{id, ...}) = Stamp.hash id
	fun toString (ND{id, kind, ...}) = let
	      val tag = (case kind
		     of NULL => "NULL"
		      | ENTRY _ => "ENTRY"
		      | JOIN _ => "JOIN"
		      | COND _ => "COND"
		      | BLOCK _ => "BLOCK"
		      | NEW _ => "NEW"
		      | DIE _ => "DIE"
		      | STABILIZE _ => "STABILIZE"
		      | EXIT _ => "EXIT"
		    (* end case *))
	      in
		tag ^ Stamp.toString id
	      end
	fun new kind = ND{id = Stamp.new(), props = PropList.newHolder(), kind = kind}
	val dummy = new NULL
	fun mkENTRY () = new (ENTRY{succ = ref dummy})
	fun mkJOIN phis = new (JOIN{preds = ref [], phis = ref phis, succ = ref dummy})
	fun mkCOND {cond, trueBranch, falseBranch} = new (COND{
		pred = ref dummy, cond = cond,
		trueBranch = ref trueBranch, falseBranch = ref falseBranch
	      })
	fun mkBLOCK body = new (BLOCK{pred = ref dummy, body = ref body, succ = ref dummy})
	fun mkNEW {actor, args} = new (NEW{
		pred = ref dummy, actor = actor, args = args, succ = ref dummy
	      })
	fun mkDIE () = new (DIE{pred = ref dummy})
	fun mkSTABILIZE () = new (STABILIZE{pred = ref dummy})
	fun mkEXIT () = new (EXIT{pred = ref dummy})
      (* editing node edges *)
	fun setPred (ND{kind, ...}, nd) = (case kind
	       of NULL => raise Fail "setPred on NULL node"
		| ENTRY _ => raise Fail "setPred on ENTRY node"
		| JOIN{preds, ...} => if List.exists (fn nd' => same(nd, nd')) (!preds)
		    then ()
		    else preds := nd :: !preds
		| COND{pred, ...} => pred := nd
		| BLOCK{pred, ...} => pred := nd
		| NEW{pred, ...} => pred := nd
		| DIE{pred} => pred := nd
		| STABILIZE{pred} => pred := nd
		| EXIT{pred} => pred := nd
	      (* end case *))
	fun preds (ND{kind, ...}) = (case kind
	       of NULL => raise Fail "preds on NULL node"
		| ENTRY _ => []
		| JOIN{preds, ...} => !preds
		| COND{pred, ...} => [!pred]
		| BLOCK{pred, ...} => [!pred]
		| NEW{pred, ...} => [!pred]
		| DIE{pred} => [!pred]
		| STABILIZE{pred} => [!pred]
		| EXIT{pred} => [!pred]
	      (* end case *))
	fun hasSucc (ND{kind, ...}) = (case kind
	       of NULL => false
		| ENTRY{succ} => true
		| JOIN{succ, ...} => true
		| COND{trueBranch, falseBranch, ...} => true
		| BLOCK{succ, ...} => true
		| NEW{succ, ...} => true
		| DIE _ => false
		| STABILIZE _ => false
		| EXIT _ => false
	      (* end case *))
	fun setSucc (ND{kind, ...}, nd) = (case kind
	       of NULL => raise Fail "setSucc on NULL node"
		| ENTRY{succ} => succ := nd
		| JOIN{succ, ...} => succ := nd
		| COND _ => raise Fail "setSucc on COND node"
		| BLOCK{succ, ...} => succ := nd
		| NEW{succ, ...} => succ := nd
		| DIE _ => raise Fail "setSucc on DIE node"
		| STABILIZE _ => raise Fail "setSucc on STABILIZE node"
		| EXIT _ => raise Fail "setSucc on EXIT node"
	      (* end case *))
	fun succs (ND{kind, ...}) = (case kind
	       of NULL => raise Fail "succs on NULL node"
		| ENTRY{succ} => [!succ]
		| JOIN{succ, ...} => [!succ]
		| COND{trueBranch, falseBranch, ...} => [!trueBranch, !falseBranch]
		| BLOCK{succ, ...} => [!succ]
		| NEW{succ, ...} => [!succ]
		| DIE _ => []
		| STABILIZE _ => []
		| EXIT _ => []
	      (* end case *))
	fun setTrueBranch (ND{kind=COND{trueBranch, ...}, ...}, nd) = trueBranch := nd
	  | setTrueBranch (nd, _) = raise Fail("setTrueBranch on " ^ toString nd)
	fun setFalseBranch (ND{kind=COND{falseBranch, ...}, ...}, nd) = falseBranch := nd
	  | setFalseBranch (nd, _) = raise Fail("setFalseBranch on " ^ toString nd)
	fun addEdge (nd1, nd2) = (
	      if hasSucc nd1
		then (
		  setSucc (nd1, nd2);
		  setPred (nd2, nd1))
		else ())
(*DEBUG*)handle ex => (
print(concat["error in addEdge(", toString nd1, ",", toString nd2, ")\n"]);
raise ex)
      (* properties *)
	fun newProp initFn =
	      PropList.newProp (fn (ND{props, ...}) => props, initFn) 
	fun newFlag () =
	      PropList.newFlag (fn (ND{props, ...}) => props) 
      end

    structure Stmt =
      struct
	fun same (STM{id=a, ...}, STM{id=b, ...}) = Stamp.same(a, b)
	fun compare (STM{id=a, ...}, STM{id=b, ...}) = Stamp.compare(a, b)
	fun hash (STM{id, ...}) = Stamp.hash id
	fun toString (STM{id, kind, ...}) = let
	      val tag = (case kind
		     of S_SIMPLE(ND{kind, ...}) => (case kind
			   of NULL => "NULL"
			    | ENTRY _ => "ENTRY"
			    | JOIN _ => "JOIN"
			    | COND _ => raise Fail "unexpected S_SIMPLE with COND node"
			    | BLOCK _ => "BLOCK"
			    | NEW _ => "NEW"
			    | DIE _ => "DIE"
			    | STABILIZE _ => "STABILIZE"
			    | EXIT _ => "EXIT"
			  (* end case *))
		      | S_IF _ => "IF"
		      | S_LOOP _ => "LOOP"
		    (* end case *))
	      in
		tag ^ Stamp.toString id
	      end
      (* return the entry node of the statement *)
	fun entry (STM{kind, ...}) = (case kind
	       of S_SIMPLE nd => nd
		| S_IF{cond, ...} => cond
		| S_LOOP{hdr, ...} => entry hdr
	      (* end case *))
      (* return the tail-end node of a statement (not applicable to S_IF or S_LOOP) *)
	fun tail (STM{kind, ...}) = (case kind
	       of S_SIMPLE nd => nd
		| S_IF{cond, ...} => raise Fail "tail of IF"
		| S_LOOP{hdr, ...} => raise Fail "tail of LOOP"
	      (* end case *))
      (* statement constructor functions *)
	fun new (kind, next) = STM{
	    id = Stamp.new(),
	    props = PropList.newHolder(),
	    kind = kind,
	    next = next
	  }
	val dummy = new (S_SIMPLE(Node.dummy), NONE)
	fun mkENTRY next = new (S_SIMPLE(Node.mkENTRY ()), next)
	fun mkJOIN (phis, next) = new (S_SIMPLE(Node.mkJOIN phis), next)
	fun mkIF (cond, thenBranch, elseBranch, next) = let
	      val cond = Node.mkCOND {
		    cond = cond,
		    trueBranch = entry thenBranch,
		    falseBranch = entry elseBranch
		  }
	      in
		Node.setPred (entry thenBranch, cond);
		Node.setPred (entry elseBranch, cond);
		new (S_IF{cond = cond, thenBranch = thenBranch, elseBranch = elseBranch}, next)
	      end
	fun mkBLOCK (body, next) = new (S_SIMPLE(Node.mkBLOCK body), next)
	fun mkNEW (actor, args, next) = new (S_SIMPLE(Node.mkNEW{actor=actor, args=args}), next)
	fun mkDIE () = new (S_SIMPLE(Node.mkDIE ()), NONE)
	fun mkSTABILIZE () = new (S_SIMPLE(Node.mkSTABILIZE ()), NONE)
	fun mkEXIT () = new (S_SIMPLE(Node.mkEXIT ()), NONE)
      (* properties *)
	fun newProp initFn =
	      PropList.newProp (fn (STM{props, ...}) => props, initFn) 
	fun newFlag () =
	      PropList.newFlag (fn (STM{props, ...}) => props) 
      end

    structure Var =
      struct
	fun new (name, ty) = V{
		name = name,
		id = Stamp.new(),
		ty = ty,
		bind = ref VB_NONE,
		useCnt = ref 0,
		props = PropList.newHolder()
	      }
	fun name (V{name, ...}) = name
	fun ty (V{ty, ...}) = ty
	fun binding (V{bind, ...}) = !bind
	fun setBinding (V{bind, ...}, vb) = bind := vb
	fun same (V{id=a, ...}, V{id=b, ...}) = Stamp.same(a, b)
	fun compare (V{id=a, ...}, V{id=b, ...}) = Stamp.compare(a, b)
	fun hash (V{id, ...}) = Stamp.hash id
	fun toString (V{name, id, ...}) = name ^ Stamp.toString id
      (* properties *)
	fun newProp initFn =
	      PropList.newProp (fn (V{props, ...}) => props, initFn) 
	fun newFlag () =
	      PropList.newFlag (fn (V{props, ...}) => props) 
	local
	  structure V =
	    struct
	      type ord_key = var
	      val compare = compare
	    end
	in
	structure Map = RedBlackMapFn (V)
	structure Set = RedBlackSetFn (V)
	end
	structure Tbl = HashTableFn (
	  struct
	    type hash_key = var
	    val hashVal = hash
	    val sameKey = same
	  end)
      end

  (* DFS sorting of the graph rooted at the entry to a statement; the resulting list will
   * be in preorder with parents before children.
   *)
    fun sortNodes stmt = let
	  val {getFn, setFn} = PropList.newFlag (fn (ND{props, ...}) => props)
	  fun dfs (nd, l) =
		if getFn nd
		  then l
		  else (
		    setFn (nd, true);
		    nd :: List.foldl dfs l (Node.succs nd))
	  val nodes = dfs (Stmt.entry stmt, [])
	  in
	    List.app (fn nd => setFn(nd, false)) nodes;
	    nodes
	  end

  (* apply a function to all of the nodes in the graph rooted at the entry to the statement *)
    fun applyToNodes (f : node -> unit) stmt = let
	  val {getFn, setFn} = PropList.newFlag (fn (ND{props, ...}) => props)
	  fun dfs (nd, l) =
		if getFn nd
		  then l
		  else (
		    f nd; (* visit *)
		    setFn (nd, true);
		    nd :: List.foldl dfs l (Node.succs nd))
	  val nodes = dfs (Stmt.entry stmt, [])
	  in
	    List.app (fn nd => setFn(nd, false)) nodes
	  end

  end