(* ssa-fn.sml
 *
 * COPYRIGHT (c) 2010 The Diderot Project (http://diderot-language.cs.uchicago.edu)
 * All rights reserved.
 *)

signature SSA =
  sig

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

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

    datatype cfg = CFG of {
	entry : node,	(* the entry node of a graph; not necessarily an ENTRY node *)
	exit : node	(* the exit node of a graph; not necessarily a DIE, STABILIZE,
			 * or EXIT node.
			 *)
      }

    and 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 : phi list ref,
	    succ : node ref
	  }
      | COND of {
	    pred : node ref,
	    cond : var,
	    trueBranch : node ref,
	    falseBranch : node ref
	  }
      | COM of  {			(* comment *)
	    pred : node ref,
	    text : string list,
	    succ : node ref
	  }
      | ASSIGN of {			(* assignment *)
	    pred : node ref,
	    stm : assign,
	    succ : node ref
	  }
      | NEW of {			(* create new strand instance *)
	    pred : node ref,
	    strand : Atom.atom,
	    args : var list,
	    succ : node ref
	  }
      | DIE of {
	    pred : node ref
	  }
      | STABILIZE of {
	    pred : node ref
	  }
      | EXIT of {
	    pred : node ref
	  }

    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)
	 and phi = (var * var list)


  (***** Program representation *****)

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

    and strand = Strand of {
	name : Atom.atom,
	params : var list,
	state : var list,
	stateInit : cfg,
	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 : cfg		(* method body *)
      }


  (* operations on CFGs *)
    structure CFG : sig
      (* the empty CFG *)
	val empty : cfg
      (* is a CFG empty? *)
	val isEmpty : cfg -> bool
      (* create a basic block from a list of assignments *)
	val mkBlock : assign list -> cfg
      (* entry/exit nodes of a CFG *)
	val entry : cfg -> node
	val exit : cfg -> node
      (* DFS sorting of the graph rooted at the entry to a statement; the resulting list will
       * be in preorder with parents before children.
       *)
	val sort : cfg -> node list
      (* apply a function to all of the nodes in the graph rooted at the entry to the statement *)
	val apply : (node -> unit) -> cfg -> unit
      (* delete a simple node from a CFG *)
	val deleteNode : node -> unit
      (* replace a simple node in a cfg with a different simple node *)
	val replaceNode : (node * node) -> unit
      (* replace a simple node in a cfg with a subgraph *)
	val splice : (node * cfg) -> unit
      (* concatenate two CFGs *)
	val concat : cfg * cfg -> cfg
      (* append a node to a CFG *)
	val appendNode : cfg * node -> cfg
      end

  (* operations on CFG nodes *)
    structure Node : sig
	val id : node -> Stamp.stamp
	val kind : node -> node_kind
	val same : node * node -> bool
	val compare : node * node -> order
	val hash : node -> word
	val toString : node -> string
	val isNULL : node -> bool
      (* dummy node *)
	val dummy : node
      (* CFG edges *)
	val hasPred : node -> bool
	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 mkCOM : string list -> node
	val mkASSIGN : assign -> node
	val mkNEW : {strand : 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

  (* operations on variables *)
    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 useCount : var -> int
	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

  (* return a string representation of a PHI node *)
    val phiToString : phi -> string

  (* return a string representation of an assignment *)
    val assignToString : assign -> string

  end

functor SSAFn (

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

  ) : SSA = struct

    structure Ty = Ty
    structure Op = Op

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

    datatype cfg = CFG of {
	entry : node,	(* the entry node of a graph; not necessarily an ENTRY node *)
	exit : node	(* the exit node of a graph; not necessarily a DIE, STABILIZE,
			 * or EXIT node.
			 *)
      }

    and 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 : phi list ref,
	    succ : node ref
	  }
      | COND of {
	    pred : node ref,
	    cond : var,
	    trueBranch : node ref,
	    falseBranch : node ref
	  }
      | COM of  {			(* comment *)
	    pred : node ref,
	    text : string list,
	    succ : node ref
	  }
      | ASSIGN of {			(* assignment *)
	    pred : node ref,
	    stm : assign,
	    succ : node ref
	  }
      | NEW of {			(* create new strand instance *)
	    pred : node ref,
	    strand : Atom.atom,
	    args : var list,
	    succ : node ref
	  }
      | DIE of {
	    pred : node ref
	  }
      | STABILIZE of {
	    pred : node ref
	  }
      | EXIT of {
	    pred : node ref
	  }

    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)
	 and phi = (var * var list)


  (***** Program representation *****)

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

    and strand = Strand of {
	name : Atom.atom,
	params : var list,
	state : var list,
	stateInit : cfg,
	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 : cfg		(* method body *)
      }

    structure Node =
      struct
	fun id (ND{id, ...}) = id
	fun kind (ND{kind, ...}) = kind
	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"
		      | COM _ => "COM"
		      | ASSIGN _ => "ASSIGN"
		      | 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 mkCOM text = new (COM{pred = ref dummy, text = text, succ = ref dummy})
	fun mkASSIGN stm = new (ASSIGN{pred = ref dummy, stm = stm, succ = ref dummy})
	fun mkNEW {strand, args} = new (NEW{
		pred = ref dummy, strand = strand, 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})
	fun isNULL (ND{kind=NULL, ...}) = true
	  | isNULL _ = false
      (* editing node edges *)
	fun hasPred (ND{kind, ...}) = (case kind
	       of NULL => false
		| ENTRY _ => false
		| _ => true
	      (* end case *))
	fun setPred (nd0 as ND{kind, ...}, nd) = (case kind
	       of NULL => raise Fail("setPred on NULL node "^toString nd0)
		| ENTRY _ => raise Fail("setPred on ENTRY node "^toString nd0)
		| JOIN{preds, ...} => if List.exists (fn nd' => same(nd, nd')) (!preds)
		    then ()
		    else preds := nd :: !preds
		| COND{pred, ...} => pred := nd
		| COM{pred, ...} => pred := nd
		| ASSIGN{pred, ...} => pred := nd
		| NEW{pred, ...} => pred := nd
		| DIE{pred} => pred := nd
		| STABILIZE{pred} => pred := nd
		| EXIT{pred} => pred := nd
	      (* end case *))
	fun preds (nd as ND{kind, ...}) = (case kind
	       of NULL => raise Fail("preds on NULL node "^toString nd)
		| ENTRY _ => []
		| JOIN{preds, ...} => !preds
		| COND{pred, ...} => [!pred]
		| COM{pred, ...} => [!pred]
		| ASSIGN{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 _ => true
		| JOIN _ => true
		| COND _ => true
		| COM _ => true
		| ASSIGN _ => true
		| NEW _ => true
		| DIE _ => false
		| STABILIZE _ => false
		| EXIT _ => false
	      (* end case *))
	fun setSucc (nd0 as ND{kind, ...}, nd) = (case kind
	       of NULL => raise Fail("setSucc on NULL node "^toString nd0)
		| ENTRY{succ} => succ := nd
		| JOIN{succ, ...} => succ := nd
		| COND _ => raise Fail("setSucc on COND node "^toString nd0)
		| COM{succ, ...} => succ := nd
		| ASSIGN{succ, ...} => succ := nd
		| NEW{succ, ...} => succ := nd
		| DIE _ => raise Fail("setSucc on DIE node "^toString nd0)
		| STABILIZE _ => raise Fail("setSucc on STABILIZE node "^toString nd0)
		| EXIT _ => raise Fail("setSucc on EXIT node "^toString nd0)
	      (* end case *))
	fun succs (nd as ND{kind, ...}) = (case kind
	       of NULL => raise Fail("succs on NULL node "^toString nd)
		| ENTRY{succ} => [!succ]
		| JOIN{succ, ...} => [!succ]
		| COND{trueBranch, falseBranch, ...} => [!trueBranch, !falseBranch]
		| COM{succ, ...} => [!succ]
		| ASSIGN{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 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 useCount (V{useCnt, ...}) = !useCnt
	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

    structure CFG =
      struct
	val empty = CFG{entry = Node.dummy, exit = Node.dummy}

	fun isEmpty (CFG{entry, exit}) =
	      Node.same(entry, exit) andalso Node.isNULL entry

      (* create a basic block from a list of assignments *)
	fun mkBlock [] = empty
	  | mkBlock (stm::stms) = let
	      val entry = Node.mkASSIGN stm
	      fun f (stm, prev) = let
		    val nd = Node.mkASSIGN stm
		    in
		      Node.addEdge (prev, nd);
		      nd
		    end
	      val exit = List.foldl f entry stms
	      in
		CFG{entry = entry, exit = exit}
	      end

      (* entry/exit nodes of a CFG *)
	fun entry (CFG{entry = nd, ...}) = nd
	fun exit (CFG{exit = nd, ...}) = nd

      (* DFS sorting of the graph rooted at the entry to a statement; the resulting list will
       * be in preorder with parents before children.
       *)
	fun sort (CFG{entry, ...}) = 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 (entry, [])
	      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 apply (f : node -> unit) (CFG{entry, ...}) = 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 (entry, [])
	      in
		List.app (fn nd => setFn(nd, false)) nodes
	      end
      (* delete a simple node from a CFG *)
	fun deleteNode (ND{kind, ...}) = (case kind
	       of ASSIGN{pred, succ, ...} => (
		    Node.setPred (!succ, !pred);
		    Node.setSucc (!pred, !succ))
		| _ => raise Fail "unsupported deleteNode"
	      (* end case *))

      (* replace a simple node in a cfg with a subgraph *)
	fun replaceNode (ND{kind, ...}, node) = (case kind
	       of ASSIGN{pred, succ, ...} => (
		    Node.addEdge (!pred, node);
		    Node.addEdge (node, !succ))
		| _ => raise Fail "unsupported splice"
	      (* end case *))

      (* replace a simple node in a cfg with a subgraph *)
	fun splice (nd as ND{kind, ...}, cfg as CFG{entry, exit}) =
	      if isEmpty cfg
		then deleteNode nd
		else (case kind
		   of ASSIGN{pred, succ, ...} => (
			Node.addEdge (!pred, entry);
			Node.addEdge (exit, !succ))
		    | _ => raise Fail "unsupported replaceNode"
		  (* end case *))

      (* concatenate two CFGs *)
	fun concat (cfg1 as CFG{entry=e1, exit=x1}, cfg2 as CFG{entry=e2, exit=x2}) =
	      if isEmpty cfg1 then cfg2
	      else if isEmpty cfg2 then cfg1
	      else (
		Node.setSucc (x1, e2);
		Node.setPred (e2, x1);
		CFG{entry = e1, exit = x2})

      (* append a node to a CFG *)
	fun appendNode (cfg as CFG{entry, exit}, nd) =
	      if isEmpty cfg
		then CFG{entry=nd, exit=nd}
		else (
		  Node.setPred (nd, exit);
		  Node.setSucc (exit, nd);
		  CFG{entry=entry, exit=nd})

      end

  (* return a string representation of a PHI node *)
    fun phiToString (y, xs) = String.concat [
	    Ty.toString(Var.ty y), " ", Var.toString y, " = PHI(",
	    String.concatWith "," (List.map Var.toString xs), ")"
	  ]

  (* return a string representation of an assignment *)
    fun assignToString (y, rhs) = let
	  val rhs = (case rhs
		 of VAR x => Var.toString x
		  | LIT lit => Literal.toString lit
		  | OP(rator, xs) => String.concat [
			Op.toString rator,
			"(", String.concatWith "," (List.map Var.toString xs), ")"
		      ]
		  | CONS xs => String.concat [
			"[", String.concatWith "," (List.map Var.toString xs), "]"
		      ]
		(* end case *))
	  in
	    String.concat [Ty.toString(Var.ty y), " ", Var.toString y, " = ", rhs]
	  end

  end (* SSAFn *)