(* ssa-fn.sml
 *
 * COPYRIGHT (c) 2010 The Diderot Project (http://diderot-language.cs.uchicago.edu)
 * All rights reserved.
 *
 * The SSAFn functor is used to generate the High, Med, and Low ILs in the Diderot
 * compiler.  These ILs have the same program and control-flow structure, but differ
 * in their types and operators.
 *)

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 an 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
	  }
      | EXIT of {			(* includes die and stabilize *)
	    pred : node ref,
	    kind : ExitKind.kind,	(* kind of exit node *)
	    live : var list		(* live variables *)
	  }

    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			(* defined by an assignment (includes globals) *)
      | VB_PHI of var list		(* defined by a phi node *)
      | VB_PARAM			(* parameter to a strand *)
      | VB_STATE_VAR			(* use of a strand state variable *)

    withtype assign = (var * rhs)
	 and phi = (var * var list)


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

    datatype program = Program of {
	globals : var list,
	globalInit : cfg,
	initially : initially,
	strands : strand list
      }

    and initially = Initially of {
	isArray : bool,			(* true for initially array, false for collection *)
	rangeInit : cfg,		(* code to compute bounds of iteration *)
	iters : (var * var * var) list,
	create : (cfg * Atom.atom * var list)
      }

    and strand = Strand of {
	name : Atom.atom,
	params : var list,
	state : (bool * var) list,	(* output variables are marked with true *)
	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
      (* replace in-going and out-going edges *)
	val replaceInEdge : {src : node, oldDst : node, dst : node} -> unit
	val replaceOutEdge : {oldSrc : node, src : node, dst : 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 mkEXIT : ExitKind.kind * var list -> node
	val mkFRAGMENT : var list -> node
	val mkRETURN : var list -> node
	val mkACTIVE : var list -> node
	val mkSTABILIZE : var list -> node
	val mkDIE : 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 copy : var -> 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 rhs *)
    val rhsToString : rhs -> string

  (* return a string representation of a variable binding *)
    val vbToString : var_bind -> string

  (* 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 an 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
	  }
      | EXIT of {			(* includes die and stabilize *)
	    pred : node ref,
	    kind : ExitKind.kind,	(* kind of exit node *)
	    live : var list		(* live variables *)
	  }

    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,
	initially : initially,
	strands : strand list
      }

    and initially = Initially of {
	isArray : bool,			(* true for initially array, false for collection *)
	rangeInit : cfg,		(* code to compute bounds of iteration *)
	iters : (var * var * var) list,
	create : (cfg * Atom.atom * var list)
      }

    and strand = Strand of {
	name : Atom.atom,
	params : var list,
	state : (bool * var) list,	(* output variables are marked with true *)
	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 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 copy (V{name, ty, ...}) = new (name, ty)
	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 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"
		      | EXIT{kind, ...} => ExitKind.toString kind
		    (* 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 (lhs, rhs) = (
	      Var.setBinding (lhs, VB_RHS rhs);
	      new (ASSIGN{pred = ref dummy, stm = (lhs, rhs), succ = ref dummy}))
	fun mkNEW {strand, args} = new (NEW{
		pred = ref dummy, strand = strand, args = args, succ = ref dummy
	      })
	fun mkEXIT (kind, xs) = new (EXIT{kind = kind, live = xs, pred = ref dummy})
	fun mkFRAGMENT xs = mkEXIT (ExitKind.FRAGMENT, xs)
	fun mkRETURN xs = mkEXIT (ExitKind.RETURN, xs)
	fun mkACTIVE xs = mkEXIT (ExitKind.ACTIVE, xs)
	fun mkSTABILIZE xs = mkEXIT (ExitKind.STABILIZE, xs)
	fun mkDIE () = mkEXIT (ExitKind.DIE, [])
	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 := !preds @ [nd]  (* assume preds are added in order *)
		| COND{pred, ...} => pred := nd
		| COM{pred, ...} => pred := nd
		| ASSIGN{pred, ...} => pred := nd
		| NEW{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]
		| 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
		| 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
		| 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]
		| 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)
	fun replaceInEdge {src, oldDst, dst} = (
	    (* first set the successor of src *)
	      case kind src
	       of COND{trueBranch, falseBranch, ...} =>
		    if same(!trueBranch, oldDst)
		      then trueBranch := dst
		      else falseBranch := dst
		| _ => setSucc (src, dst)
	      (* end case *);
	    (* then set the predecessor of dst *)
	      setPred (dst, src))
(*DEBUG*)handle ex => (
print(concat["error in replaceInEdge(", toString src, ",", toString oldDst, ",", toString dst, ")\n"]);
raise ex)
	fun replaceOutEdge {oldSrc, src, dst} = (
	    (* first set the successor of src *)
	      setSucc (src, dst);
	    (* then set the predecessor of dst *)
	      case kind dst
	       of JOIN{preds, ...} => let
		    fun edit [] = raise Fail "replaceOutEdge: cannot find predecessor"
		      | edit (nd::nds) = if same(nd, oldSrc) then src::nds else nd::edit nds
		    in
		      preds := edit (!preds)
		    end
		| _ => setPred (dst, src)
	      (* end case *))
(*DEBUG*)handle ex => (
print(concat["error in replaceOutEdge(", toString oldSrc, ",", toString src, ",", toString dst, ")\n"]);
raise ex)
      (* properties *)
	fun newProp initFn =
	      PropList.newProp (fn (ND{props, ...}) => props, initFn) 
	fun newFlag () =
	      PropList.newFlag (fn (ND{props, ...}) => props) 
      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 as ND{kind, ...}) = (case kind
	       of ASSIGN{pred = ref pred, succ = ref succ, ...} => (
		    Node.setPred (succ, pred);
		    case Node.kind pred
		     of COND{trueBranch, falseBranch, ...} => (
			(* note that we treat each branch independently, so that we handle the
			 * situation where both branches are the same node.
			 *)
			  if Node.same(!trueBranch, nd)
			    then Node.setTrueBranch(pred, succ)
			    else ();
			  if Node.same(!falseBranch, nd)
			    then Node.setFalseBranch(pred, succ)
			    else ())
		      | _ => Node.setSucc (pred, succ)
		    (* end case *))
		| _ => raise Fail "unsupported deleteNode"
	      (* end case *))

      (* replace a simple node in a cfg with a subgraph *)
	fun replaceNode (oldNd as ND{kind, ...}, node) = (case kind
	       of ASSIGN{pred, succ, ...} => (
		    Node.replaceInEdge {src = !pred, oldDst = oldNd, dst = node};
		    Node.replaceOutEdge {oldSrc = oldNd, src = node, dst = !succ})
		| _ => raise Fail "unsupported replaceNode"
	      (* 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.replaceInEdge {src = !pred, oldDst = nd, dst = entry};
			Node.replaceOutEdge {oldSrc = nd, src = exit, dst = !succ})
		    | _ => raise Fail "unsupported splice"
		  (* 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 rhs *)
    fun rhsToString 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 *))

  (* return a string representation of a variable binding *)
    fun vbToString VB_NONE = "NONE"
      | vbToString (VB_RHS rhs) = concat["RHS(", rhsToString rhs, ")"]
      | vbToString (VB_PHI xs) = concat[
	    "PHI(", String.concatWith "," (List.map Var.toString xs), ")"
	  ]
      | vbToString VB_PARAM = "PARAM"
      | vbToString VB_STATE_VAR = "STATE_VAR"


  (* 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) =
	  String.concat [Ty.toString(Var.ty y), " ", Var.toString y, " = ", rhsToString rhs]

  end (* SSAFn *)