(* 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 Op : OPERATORS

    datatype var = V of {
	name : string,			(* name *)
	id : Stamp.stamp,		(* unique ID *)
	useCnt : int ref,		(* count of uses *)
	props : PropList.holder
      }

  (* a statement is a CFG fragment with a single entrypoint and
   * a single continuation.
   *)
    datatype stmt
      = BLOCK of stmt
      | SEQ of stmt list
      | IF of {
	  pre : stmt,
	  cond : var,
	  trueBranch : stmt
	  falseBranch : stmt
	}
      | WHILE of {
	  hdr : block,
	  cond : var,
	  body : stmt
	}

    and block = BLK of {
	  parent : stmt ref,			(* parent statement of this block *)
	  id : Stamp.stamp,			(* unique ID *)
	  preds : block list ref,		(* list of predecessor blocks in the CFG *)
	  phi : (var * var list) list ref,	(* phi statements *)
	  body : simple_stmt list ref,		
	  succs : block list ref		(* successor blocks in the CFG *)
	}

    and simple_stmt
      = ASSIGN of var * rhs
      | DIE
      | STABILIZE
      | RETURN

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

    val newVar : string -> var
    val newBlock : unit -> block

    val entryBlock : stmt -> block
    val nextBlock : stmt -> block

  end

functor SSAFn (Op : OPERATORS) : SSA =
  struct

    datatype var = V of {
	name : string,			(* name *)
	id : Stamp.stamp,		(* unique ID *)
	useCnt : int ref,		(* count of uses *)
	props : PropList.holder
      }

  datatype stmt
    = EXIT
    | SEQ of block * stmt
    | IF of {
	pre : stmt,
	cond : var,
	trueBranch : stmt,
	falseBranch : stmt,
	next : stmt		(* ?? *)
      }
    | WHILE of {
	hdr : stmt,
	cond : var,
	body : stmt,
	exit : stmt		(* ?? *)
      }

  and block = BLK of {
	parent : stmt ref,			(* parent statement of this block *)
	id : Stamp.stamp,			(* unique ID *)
	preds : block list ref,			(* list of predecessor blocks in the CFG *)
	phi : (var * var list) list ref,	(* phi statements *)
	body : simple_stmt list ref,		
	succs : block list ref			(* successor blocks in the CFG *)
      }

  and simple_stmt
    = ASSIGN of var * rhs
    | DIE
    | STABILIZE
    | RETURN

    and rhs
      = VAR of var
      | OP of Op.rator * var list

  (* block properties *)
    fun parentOf (BLK{parent, ...}) = !parent
    fun predsOf (BLK{preds, ...}) = !preds
    fun succsOf (BLK{succs, ...}) = !succs

  (* IL construction code *)
    fun newVar name = V{
	    name = name,
	    id = Stamp.new(),
	    useCnt = ref 0,
	    props = PropList.newHolder()
	  }

    fun newBlock () = BLK{
	    parent = ref(SEQ[]),
	    id = Stamp.new(),
	    preds = ref[],
	    phi = ref[],
	    body = ref[],
	    succs = ref[]
	  }

    local
      fun setParent (BKL{parent, ...}, s) = (parent := s)
    in
    fun mkBLOCK blk = let
	  val s = BLOCK blk
	  in
	    setParent (blk, s);
	    s
	  end

    fun mkIF (pre, cond, t, f) = let
	  val s = IF{pre=pre, cond=cond, trueBranch=t, falseBranch=f}
	  in
	    setParent (pre, s);
	    s
	  end

    fun mkWHILE (hdr, cond, body) = let
	  val s = WHILE{hdr=hdr, cond=cond, body=body}
	  in
	    setParent (hdr, s);
	    s
	  end

    end (* local *)

    fun entryBlock (BLOCK blk) = blk
      | entryBlock (SEQ(s1::_)) = entryBlock s1
      | entryBlock (IF{pre, ...}) = entryBlock pre
      | entryBlock (WHILE{hdr, ...}) = entryBlock hdr

    fun nextBlock (SEQ(, next)) = nextBlock next
      | nextBlock (SEQ stms) = nextBlock(List.last stms)
      | nextBlock (IF{pre, ...}) = entryBlock pre
      | nextBlock (WHILE{hdr, ...}) = entryBlock hdr

  end