(* expr-fn.sml
 *
 * This code is part of the Diderot Project (http://diderot-language.cs.uchicago.edu)
 *
 * COPYRIGHT (c) 2015 The University of Chicago
 * All rights reserved.
 *
 * This functor implements hash-consing for IR expressions.  It can be used to implement
 * optimizations such as CSE and PRE.
 *)

signature EXPR =
  sig

    structure IR : SSA
    structure Op : OPERATORS where type rator = IR.Op.rator

    datatype expr_nd
      = GLOBAL of IR.global_var
      | STATE of IR.state_var
      | VAR of IR.var
      | LIT of Literal.t
      | OP of Op.rator * expr list
      | MULTIOP of int * Op.rator * expr list     (* n'th result of operator in multi-assignment *)
      | CONS of expr list * IR.Ty.ty
      | SEQ of expr list * IR.Ty.ty
      | PHI of expr list
      | EINAPP of Ein.ein * expr list

    withtype expr = expr_nd HashCons.obj

    val same : expr * expr -> bool
    val toString : expr -> string

  (* hash-cons construction of expressions *)
    type tbl

    val new : unit -> tbl

    val mkGLOBAL  : tbl -> IR.global_var -> expr
    val mkSTATE   : tbl -> IR.state_var -> expr
    val mkVAR     : tbl -> IR.var -> expr
    val mkLIT     : tbl -> Literal.t -> expr
    val mkOP      : tbl -> Op.rator * expr list -> expr
    val mkMULTIOP : tbl -> int * Op.rator * expr list -> expr
    val mkCONS    : tbl -> expr list * IR.Ty.ty -> expr
    val mkSEQ     : tbl -> expr list * IR.Ty.ty -> expr
    val mkPHI     : tbl -> expr list -> expr
    val mkEINAPP  : tbl -> Ein.ein * expr list -> expr

  (* tables, sets, and maps *)
    structure Tbl : MONO_HASH_TABLE where type Key.hash_key = expr
    structure Set : ORD_SET where type Key.ord_key = expr
    structure Map : ORD_MAP where type Key.ord_key = expr

  end

functor ExprFn (IR : SSA) : EXPR =
  struct

    structure IR = IR
    structure Op = IR.Op
    structure HC = HashCons

    datatype expr_nd
      = GLOBAL of IR.global_var
      | STATE of IR.state_var
      | VAR of IR.var
      | LIT of Literal.t
      | OP of Op.rator * expr list
      | MULTIOP of int * Op.rator * expr list     (* n'th result of operator *)
      | CONS of expr list * IR.Ty.ty
      | SEQ of expr list * IR.Ty.ty
      | PHI of expr list
      | EINAPP of Ein.ein * expr list

    withtype expr = expr_nd HashCons.obj

    val same : expr * expr -> bool = HC.same
    val same' = ListPair.allEq same

    fun sameNd (GLOBAL x, GLOBAL y) = IR.GlobalVar.same(x, y)
      | sameNd (STATE x, STATE y) = IR.StateVar.same(x, y)
      | sameNd (VAR x, VAR y) = IR.Var.same(x, y)
      | sameNd (LIT a, LIT b) = Literal.same(a, b)
      | sameNd (OP(op1, args1), OP(op2, args2)) =
          Op.same(op1, op2) andalso Op.isPure op1 andalso same'(args1, args2)
      | sameNd (MULTIOP(i1, op1, args1), MULTIOP(i2, op2, args2)) =
          (i1 = i2) andalso Op.same(op1, op2) andalso same'(args1, args2)
      | sameNd (CONS(args1, _), CONS(args2, _)) = same'(args1, args2)
      | sameNd (SEQ(args1, _), SEQ(args2, _)) = same'(args1, args2)
      | sameNd (PHI args1, PHI args2) = same'(args1, args2)
      | sameNd (EINAPP(ein1, a1), EINAPP(ein2, a2)) =
          EinUtil.same(ein1, ein2) andalso same'(a1, a2)
      | sameNd _ = false

    fun toString exp = let
          fun toS (e : expr, l) = (case #nd e
                 of GLOBAL x => IR.GlobalVar.toString x :: l
                  | STATE x => IR.StateVar.toString x :: l
                  | VAR x => IR.Var.toString x :: l
                  | LIT lit => Literal.toString lit :: l
                  | OP(rator, args) => Op.toString rator :: "(" :: argsToS (args, ")" :: l)
                  | MULTIOP(i, rator, args) =>
                      "#" :: Int.toString i :: "(" :: Op.toString rator :: "("
                        :: argsToS (args, "))" :: l)
                  | CONS(args, _) => "[" :: argsToS (args, "]" :: l)
                  | SEQ(args, _) => "{" :: argsToS (args, "}" :: l)
                  | PHI args => "PHI(" :: argsToS (args, ")" :: l)
		  | EINAPP _=> "EINAPP" :: l (* FIXME *)
                (* end case *))
          and argsToS ([], l) = l
            | argsToS ([e], l) = toS(e, l)
            | argsToS (e::es, l) = toS(e, ","::argsToS(es, l))
          in
            String.concat (toS (exp, []))
          end

  (* hash-cons construction of expressions *)
    datatype tbl = Tbl of expr_nd HC.tbl

    fun new () = Tbl(HC.new{eq = sameNd})

    fun mkGLOBAL (Tbl tbl) x = HC.cons0 tbl (0w9941 + IR.GlobalVar.hash x, GLOBAL x)
    fun mkSTATE (Tbl tbl) x = HC.cons0 tbl (0w7477 + IR.StateVar.hash x, STATE x)
    fun mkVAR (Tbl tbl) x = HC.cons0 tbl (0w7919 + IR.Var.hash x, VAR x)
    fun mkLIT (Tbl tbl) a = HC.cons0 tbl (0w6997 + Literal.hash a, LIT a)
    fun mkOP (Tbl tbl) (rator, args) =
          HC.consList tbl (Op.hash rator, fn args => OP(rator, args)) args
    fun mkMULTIOP (Tbl tbl) (i, rator, args) =
          HC.consList tbl (Op.hash rator + Word.fromInt i, fn args => OP(rator, args)) args
    fun mkCONS (Tbl tbl) (args, ty) =
          HC.consList tbl (0w5987, fn args => CONS(args, ty)) args
    fun mkSEQ (Tbl tbl) (args, ty) =
          HC.consList tbl (0w6011, fn args => SEQ(args, ty)) args
    fun mkPHI (Tbl tbl) args = HC.consList tbl (0w6079, PHI) args
    fun mkEINAPP (Tbl tbl) (rator, args) = (
          HC.consList tbl (EinUtil.hash rator, fn args => EINAPP(rator, args)) args)

  (* hash tables *)
    structure Tbl = HashTableFn (
      struct
        type hash_key = expr
        fun hashVal (e : expr) = #tag e
        val sameKey = same
      end)

  (* sets and maps *)
    structure Ord =
      struct
        type ord_key = expr
        fun compare (e1 : expr, e2 : expr) = Word.compare(#tag e1, #tag e2)
      end
    structure Set = RedBlackSetFn (Ord)
    structure Map = RedBlackMapFn (Ord)

  end