defpackage firrtl/ir-utils :
   import core
   import verse
   import firrtl/ir2

;============== DEBUG STUFF =============================

public defmulti print-debug (o:OutputStream, e:Expression|Stmt|Type|Port|Field|Module|Circuit) -> False

;============== GENSYM STUFF ======================


public defn firrtl-gensym (s:Symbol) -> Symbol : firrtl-gensym(s,HashTable<Symbol,Int>(symbol-hash))

public defn firrtl-gensym (s:Symbol,sym-hash:HashTable<Symbol,Int>) -> Symbol :
   defn get-new (s:Symbol, i:Int) -> Symbol :
      val s* = symbol-join([s i])
      if contains?(keys(sym-hash),s*) :
         get-new(s,i + 1)
      else :
         sym-hash[s] = i
         sym-hash[s*] = 0
         s*
   get-new(s,0)
   
public defn firrtl-gensym (sym-hash:HashTable<Symbol,Int>) -> Symbol :
   firrtl-gensym(`gen,sym-hash)

public defn get-sym-hash (m:InModule) -> HashTable<Symbol,Int> :
   val sym-hash = HashTable<Symbol,Int>(symbol-hash)
   defn add-name (s:Symbol) -> False :
      sym-hash[s] = 0
   defn to-port (p:Port) -> False : add-name(name(p))
   defn to-stmt (s:Stmt) -> Stmt :
     match(s) :
        (s:DefWire) : add-name(name(s))
        (s:DefRegister) : add-name(name(s))
        (s:DefInstance) : add-name(name(s))
        (s:DefMemory) : add-name(name(s))
        (s:DefNode) : add-name(name(s))
        (s:DefAccessor) : add-name(name(s))
        (s) : false
     map(to-stmt,s)

   to-stmt(body(m))
   map(to-port,ports(m))
   sym-hash
   
;============== Exceptions =====================

public definterface PassException <: Exception
public defn PassException (s:String) :
   new PassException :
      defmethod print (o:OutputStream, this) :
         print(o, s)

public defn PassExceptions (xs:Streamable<PassException>) :
   PassException(string-join(xs, "\n"))

;============== Pass/Compiler Structs ============

public definterface Compiler 
public defmulti passes (c:Compiler) -> List<Pass>
public defmulti file (c:Compiler) -> String

public definterface Pass
public defmulti pass (p:Pass) -> (Circuit -> Circuit)
public defmulti name (p:Pass) -> String
public defmulti short-name (p:Pass) -> String
public defmethod print (o:OutputStream, p:Pass) :
   print(o,name(p))
   

;============== PRINTERS ===================================

defmethod print (o:OutputStream, d:Flip) :
   print{o, _} $
   switch {d == _} :
      DEFAULT : ""
      REVERSE: "flip"

defmethod print (o:OutputStream, d:Direction) :
   print{o, _} $
   switch {d == _} :
      INPUT : "input"
      OUTPUT: "output"

defmethod print (o:OutputStream, w:Width) :
   print{o, _} $
   match(w) :
      (w:UnknownWidth) : "?"
      (w:IntWidth) : width(w)

defmethod print (o:OutputStream, op:PrimOp) :
   print{o, _} $
   switch {op == _} :
      ADD-OP : "add"
      SUB-OP : "sub"
      MUL-OP : "mul"
      DIV-OP : "div"
      MOD-OP : "mod"
      QUO-OP : "quo"
      REM-OP : "rem"
      ADD-WRAP-OP : "add-wrap"
      SUB-WRAP-OP : "sub-wrap"
      LESS-OP : "lt"
      LESS-EQ-OP : "leq"
      GREATER-OP : "gt"
      GREATER-EQ-OP : "geq"
      EQUAL-OP : "eq"
      NEQUAL-OP : "neq"
      MUX-OP : "mux"
      PAD-OP : "pad"
      AS-UINT-OP : "as-UInt"
      AS-SINT-OP : "as-SInt"
      DYN-SHIFT-LEFT-OP : "dshl"
      DYN-SHIFT-RIGHT-OP : "dshr"
      SHIFT-LEFT-OP : "shl"
      SHIFT-RIGHT-OP : "shr"
      CONVERT-OP : "convert"
      NEG-OP : "neg"
      BIT-NOT-OP : "bit-not"
      BIT-AND-OP : "bit-and"
      BIT-OR-OP : "bit-or"
      BIT-XOR-OP : "bit-xor"
      BIT-AND-REDUCE-OP : "bit-and-reduce"
      BIT-OR-REDUCE-OP : "bit-or-reduce"
      BIT-XOR-REDUCE-OP : "bit-xor-reduce"
      CONCAT-OP : "cat"
      BIT-SELECT-OP : "bit"
      BITS-SELECT-OP : "bits"
   
defmethod print (o:OutputStream, e:Expression) :
   match(e) :
      (e:Ref) : print(o, name(e))
      (e:Subfield) : print-all(o, [exp(e) "." name(e)])
      (e:Index) : print-all(o, [exp(e) "[" value(e) "]"])
      (e:UIntValue) : print-all(o, ["UInt(" value(e) ")"])
      (e:SIntValue) : print-all(o, ["SInt(" value(e) ")"])
      (e:DoPrim) :
         print-all(o, [op(e) "("])
         print-all(o, join(concat(args(e), consts(e)), ", "))
         print(o, ")")
      (e:ReadPort) : print-all(o, ["ReadPort(" mem(e) ", " index(e) ", " enable(e) ")"])
      (e:WritePort) : print-all(o, ["WritePort(" mem(e) ", " index(e) ", " enable(e) ")"])
      (e:Register) : print-all(o, ["Register(" value(e) ", " enable(e) ")"])
   print-debug(o,e)

defmethod print (o:OutputStream, c:Stmt) :
   val io = IndentedStream(o, 3)
   match(c) :
      (c:DefWire) :
         print-all(o,["wire " name(c) " : " type(c)])
      (c:DefRegister) : 
         print-all(o,["reg " name(c) " : " type(c)])
      (c:DefMemory) :
         if seq?(c) : print-all(o,["smem " name(c) " : " type(c)])
         else       : print-all(o,["cmem " name(c) " : " type(c)])
      (c:DefInstance) :
         print-all(o,["inst " name(c) " of " module(c)])
      (c:DefNode) :
         print-all(o,["node " name(c) " = " value(c)])
      (c:DefAccessor) :
         print-all(o,["accessor " name(c) " = " source(c) "[" index(c) "]"])      
      (c:Conditionally) :
         print-all(o, ["when " pred(c) " :"])
         print-debug(o,c)
         print(o,"\n")
         print(io,conseq(c))
         if alt(c) not-typeof EmptyStmt :
            print(o, "\nelse :")
            print(io, "\n")
            print(io,alt(c))
      (c:Begin) :
         do(print{o,_}, join(body(c), "\n"))
      (c:Connect) :
         print-all(o, [loc(c) " := " exp(c)])
      (c:BulkConnect) :
         print-all(o, [loc(c) " <> " exp(c)])
      (c:OnReset) :
         print-all(o, ["on-reset " loc(c) " := " exp(c)])
      (c:EmptyStmt) :
         print(o, "skip")
   if not c typeof Conditionally|Begin|EmptyStmt : print-debug(o,c)

defmethod print (o:OutputStream, t:Type) :
   match(t) :
      (t:UnknownType) :
         print(o, "?")
      (t:UIntType) :
         match(width(t)) :
            (w:UnknownWidth) : print-all(o, ["UInt"])
            (w) : print-all(o, ["UInt<" width(t) ">"])
      (t:SIntType) :
         match(width(t)) :
            (w:UnknownWidth) : print-all(o, ["SInt"])
            (w) : print-all(o, ["SInt<" width(t) ">"])
      (t:BundleType) :
         print(o, "{")
         print-all(o, join(fields(t), ", "))
         print(o, "}")
      (t:VectorType) :
         print-all(o, [type(t) "[" size(t) "]"])
   print-debug(o,t)

defmethod print (o:OutputStream, f:Field) :
   print-all(o, [flip(f) " " name(f) " : " type(f)])
   print-debug(o,f)

defmethod print (o:OutputStream, p:Port) :
   print-all(o, [direction(p) " " name(p) " : " type(p)])
   print-debug(o,p)

defmethod print (o:OutputStream, m:InModule) :
   print-all(o, ["module " name(m) " :"])
   print-debug(o,m)
   print(o,"\n")
   val io = IndentedStream(o, 3)
   for p in ports(m) do :
      println(io,p)
   print(io,body(m))

defmethod print (o:OutputStream, m:ExModule) :
   print-all(o, ["exmodule " name(m) " :"])
   print-debug(o,m)
   print(o,"\n")
   val io = IndentedStream(o, 3)
   for p in ports(m) do :
      println(io,p)

defmethod print (o:OutputStream, c:Circuit) :
   print-all(o, ["circuit " main(c) " :"])
   print-debug(o,c)
   print(o,"\n")
   val io = IndentedStream(o, 3)
   for m in modules(c) do :
      println(io, m)

;=================== MAPPERS ===============================
public defn map<?T> (f: Type -> Type, t:?T&Type) -> T :
   val type = 
      match(t) :
         (t:T&BundleType) :
            BundleType $
            for p in fields(t) map :
               Field(name(p), flip(p), f(type(p)))
         (t:T&VectorType) :
            VectorType(f(type(t)), size(t))
         (t) :
            t
   type as T&Type

public defmulti map<?T> (f: Expression -> Expression, e:?T&Expression) -> T 
defmethod map (f: Expression -> Expression, e:Expression) -> Expression :
   match(e) :
      (e:Subfield) : Subfield(f(exp(e)), name(e), type(e)) 
      (e:Index) : Index(f(exp(e)), value(e), type(e))
      (e:DoPrim) : DoPrim(op(e), map(f, args(e)), consts(e), type(e))
      (e:ReadPort) : ReadPort(f(mem(e)), f(index(e)), type(e), f(enable(e)))
      (e:WritePort) : WritePort(f(mem(e)), f(index(e)), type(e), f(enable(e)))
      (e:Register) : Register(type(e),f(value(e)),f(enable(e)))
      (e) : e

public defmulti map<?T> (f: Expression -> Expression, c:?T&Stmt) -> T
defmethod map (f: Expression -> Expression, c:Stmt) -> Stmt :
   match(c) :
      (c:DefAccessor) : DefAccessor(info(c),name(c), f(source(c)), f(index(c)))
      (c:DefNode) : DefNode(info(c),name(c), f(value(c)))
      (c:DefInstance) : DefInstance(info(c),name(c), f(module(c)))
      (c:Conditionally) : Conditionally(info(c),f(pred(c)), conseq(c), alt(c))
      (c:Connect) : Connect(info(c),f(loc(c)), f(exp(c)))
      (c:BulkConnect) : BulkConnect(info(c),f(loc(c)), f(exp(c)))
      (c:OnReset) : OnReset(info(c),f(loc(c)),f(exp(c)))
      (c) : c

public defmulti map<?T> (f: Stmt -> Stmt, c:?T&Stmt) -> T
defmethod map (f: Stmt -> Stmt, c:Stmt) -> Stmt :
   match(c) :
      (c:Conditionally) : Conditionally(info(c),pred(c), f(conseq(c)), f(alt(c)))
      (c:Begin) : Begin(map(f, body(c)))
      (c) : c

public defmulti map<?T> (f: Width -> Width, c:?T&Expression) -> T
defmethod map (f: Width -> Width, c:Expression) -> Expression :
   match(c) :
      (c:UIntValue) : UIntValue(value(c),f(width(c)))
      (c:SIntValue) : SIntValue(value(c),f(width(c)))
      (c) : c

public defmulti map<?T> (f: Width -> Width, c:?T&Type) -> T
defmethod map (f: Width -> Width, c:Type) -> Type :
   match(c) :
      (c:UIntType) : UIntType(f(width(c)))
      (c:SIntType) : SIntType(f(width(c)))
      (c) : c

public defmulti map<?T> (f: Type -> Type, c:?T&Expression) -> T
defmethod map (f: Type -> Type, c:Expression) -> Expression :
   match(c) :
      (c:Ref) : Ref(name(c),f(type(c)))
      (c:Subfield) : Subfield(exp(c),name(c),f(type(c)))
      (c:Index) : Index(exp(c),value(c),f(type(c)))
      (c:DoPrim) : DoPrim(op(c),args(c),consts(c),f(type(c)))
      (c:ReadPort) : ReadPort(mem(c),index(c),f(type(c)),enable(c))
      (c:WritePort) : WritePort(mem(c),index(c),f(type(c)),enable(c))
      (c:Register) : Register(f(type(c)),value(c),enable(c))
      (c) : c

public defmulti map<?T> (f: Type -> Type, c:?T&Stmt) -> T
defmethod map (f: Type -> Type, c:Stmt) -> Stmt :
   match(c) :
      (c:DefWire) : DefWire(info(c),name(c),f(type(c)))
      (c:DefRegister) : DefRegister(info(c),name(c),f(type(c)))
      (c:DefMemory) : DefMemory(info(c),name(c),f(type(c)) as VectorType,seq?(c))
      (c) : c

public defmulti mapr<?T> (f: Width -> Width, t:?T&Type) -> T
defmethod mapr (f: Width -> Width, t:Type) -> Type :
   defn apply-t (t:Type) -> Type :
      map{f,_} $ map(apply-t,t)
   apply-t(t)

public defmulti mapr<?T> (f: Width -> Width, s:?T&Stmt) -> T
defmethod mapr (f: Width -> Width, s:Stmt) -> Stmt :
   defn apply-t (t:Type) -> Type : mapr(f,t)
   defn apply-e (e:Expression) -> Expression :
      map{f,_} $ map{apply-t,_} $ map(apply-e,e)
   defn apply-s (s:Stmt) -> Stmt :
      map{apply-t,_} $ map{apply-e,_} $ map(apply-s,s)
   apply-s(s)


;================= HELPER FUNCTIONS USING MAP ===================
public defmulti do (f:Expression -> ?, e:Expression) -> False
defmethod do (f:Expression -> ?, e:Expression) -> False :
   defn f* (x:Expression) :
      f(x)
      x
   map(f*,e)
   false

public defmulti do (f:Expression -> ?, s:Stmt) -> False
defmethod do (f:Expression -> ?, s:Stmt) -> False :
   defn f* (x:Expression) : 
      f(x)
      x
   map(f*,s)
   false

public defmulti do (f:Stmt -> ?, s:Stmt) -> False
defmethod do (f:Stmt -> ?, s:Stmt) -> False :
   defn f* (x:Stmt) :
      f(x)
      x
   map(f*,s)
   false

; Not well defined - usually use dor on fields of a recursive type
;public defmulti dor (f:Expression -> ?, e:Expression) -> False
;defmethod dor (f:Expression -> ?, e:Expression) -> False :
;   f(e)
;   for x in e map : 
;      dor(f,x)
;      x
;   false
;
;public defmulti dor (f:Expression -> ?, s:Stmt) -> False
;defmethod dor (f:Expression -> ?, s:Stmt) -> False :
;   defn f* (x:Expression) :
;      dor(f,x)
;      x
;   map(f*,s)
;   false
;
;public defmulti dor (f:Stmt -> ?, s:Stmt) -> False
;defmethod dor (f:Stmt -> ?, s:Stmt) -> False :
;   f(s)
;   defn f* (x:Stmt) :
;      dor(f,x)
;      x
;   map(f*,s)
;   false
;
;public defmulti sub-exps (s:Expression|Stmt) -> List<Expression>
;defmethod sub-exps (e:Expression) -> List<Expression> :
;   val l = Vector<Expression>()
;   defn f (x:Expression) : add(l,x)
;   do(f,e)
;   to-list(l)
;defmethod sub-exps (e:Stmt) -> List<Expression> :
;   val l = Vector<Expression>()
;   defn f (x:Expression) : add(l,x)
;   do(f,e)
;   to-list(l)
;
;public defmulti sub-stmts (s:Stmt) -> List<Stmt>
;defmethod sub-stmts (s:Stmt) :
;   val l = Vector<Stmt>()
;   defn f (x:Stmt) : add(l,x)
;   do(f,s)
;   to-list(l)

;=================== ADAM OPS ===============================
public defn split (s:String,c:Char) -> List<String> :
   if not contains(to-list(s),c) : list(s)
   else :
      val index = label<Int> ret :
         var i = 0
         for c* in to-list(s) do :
            if c* == c : ret(i)
            else : 
               i = i + 1
         ret(0)
      val h = substring(s,0,index)
      val t = substring(s,index + 1,length(s))
      List(h,split(t,c))
     
public defn contains (l:List<Char>, c:Char) :
   label<True|False> myret :
      for x in l do :
         if x == c : myret(true)
      false

public defn merge!<?K,?V> (a:HashTable<?K,?V>, b:HashTable<K,V>) :
   for e in b do :
      a[key(e)] = value(e)

public defn pow (x:Int,y:Int) -> Int :
   var x* = 1
   var y* = y
   while y* != 0 :
      x* = x* * x
      y* = y* - 1
   x*