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) -> False

;============== 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"
      ADD-UU-OP : "add-uu"
      ADD-US-OP : "add-us"
      ADD-SU-OP : "add-su"
      ADD-SS-OP : "add-ss"
      SUB-OP : "sub"
      SUB-UU-OP : "sub-uu"
      SUB-US-OP : "sub-us"
      SUB-SU-OP : "sub-su"
      SUB-SS-OP : "sub-ss"
      MUL-OP : "mul"
      MUL-UU-OP : "mul-uu"
      MUL-US-OP : "mul-us"
      MUL-SU-OP : "mul-su"
      MUL-SS-OP : "mul-ss"
      DIV-OP : "div"
      DIV-UU-OP : "div-uu"
      DIV-US-OP : "div-us"
      DIV-SU-OP : "div-su"
      DIV-SS-OP : "div-ss"
      MOD-OP : "mod"
      MOD-UU-OP : "mod-uu"
      MOD-US-OP : "mod-us"
      MOD-SU-OP : "mod-su"
      MOD-SS-OP : "mod-ss"
      QUO-OP : "quo"
      QUO-UU-OP : "quo-uu"
      QUO-US-OP : "quo-us"
      QUO-SU-OP : "quo-su"
      QUO-SS-OP : "quo-ss"
      REM-OP : "rem"
      REM-UU-OP : "rem-uu"
      REM-US-OP : "rem-us"
      REM-SU-OP : "rem-su"
      REM-SS-OP : "rem-ss"
      ADD-WRAP-OP : "add-wrap"
      ADD-WRAP-UU-OP : "add-wrap-uu"
      ADD-WRAP-US-OP : "add-wrap-us"
      ADD-WRAP-SU-OP : "add-wrap-su"
      ADD-WRAP-SS-OP : "add-wrap-ss"
      SUB-WRAP-OP : "sub-wrap"
      SUB-WRAP-UU-OP : "sub-wrap-uu"
      SUB-WRAP-US-OP : "sub-wrap-us"
      SUB-WRAP-SU-OP : "sub-wrap-su"
      SUB-WRAP-SS-OP : "sub-wrap-ss"
      LESS-OP : "lt"
      LESS-UU-OP : "lt-uu"
      LESS-US-OP : "lt-us"
      LESS-SU-OP : "lt-su"
      LESS-SS-OP : "lt-ss"
      LESS-EQ-OP : "leq"
      LESS-EQ-UU-OP : "leq-uu"
      LESS-EQ-US-OP : "leq-us"
      LESS-EQ-SU-OP : "leq-su"
      LESS-EQ-SS-OP : "leq-ss"
      GREATER-OP : "gt"
      GREATER-UU-OP : "gt-uu"
      GREATER-US-OP : "gt-us"
      GREATER-SU-OP : "gt-su"
      GREATER-SS-OP : "gt-ss"
      GREATER-EQ-OP : "geq"
      GREATER-EQ-UU-OP : "geq-uu"
      GREATER-EQ-US-OP : "geq-us"
      GREATER-EQ-SU-OP : "geq-su"
      GREATER-EQ-SS-OP : "geq-ss"
      EQUAL-OP : "eq"
      EQUAL-UU-OP : "eq-uu"
      EQUAL-SS-OP : "eq-ss"
      NEQUAL-OP : "neq"
      NEQUAL-UU-OP : "neq-uu"
      NEQUAL-SS-OP : "neq-ss"
      MUX-OP : "mux"
      MUX-UU-OP : "mux-uu"
      MUX-SS-OP : "mux-ss"
      PAD-OP : "pad"
      PAD-U-OP : "pad-u"
      PAD-S-OP : "pad-s"
      AS-UINT-OP : "as-UInt"
      AS-UINT-U-OP : "as-UInt-u"
      AS-UINT-S-OP : "as-UInt-s"
      AS-SINT-OP : "as-SInt"
      AS-SINT-U-OP : "as-SInt-u"
      AS-SINT-S-OP : "as-SInt-s"
      SHIFT-LEFT-OP : "shl"
      SHIFT-LEFT-U-OP : "shl-u"
      SHIFT-LEFT-S-OP : "shl-s"
      SHIFT-RIGHT-OP : "shr"
      SHIFT-RIGHT-U-OP : "shr-u"
      SHIFT-RIGHT-S-OP : "shr-s"
      CONVERT-OP : "convert"
      CONVERT-U-OP : "convert-u"
      CONVERT-S-OP : "convert-s"
      NEG-OP : "neg"
      NEG-U-OP : "neg-u"
      NEG-S-OP : "neg-s"
      BIT-NOT-OP : "bit-not"
      BIT-AND-OP : "bit-and"
      BIT-OR-OP : "bit-or"
      BIT-XOR-OP : "bit-xor"
      CONCAT-OP : "cat"
      BIT-SELECT-OP : "bit"
      BITS-SELECT-OP : "bits"
      BIT-AND-REDUCE-OP : "bit-and-reduce"
      BIT-OR-REDUCE-OP : "bit-or-reduce"
      BIT-XOR-REDUCE-OP : "bit-xor-reduce"
   
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) :
         print-all(o,["mem " 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) :
         println-all(o, ["when " pred(c) " :"])
         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:EmptyStmt) :
         print(o, "skip")
   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:Module) :
   println-all(o, ["module " name(m) " :"])
   val io = IndentedStream(o, 3)
   for p in ports(m) do :
      println(io,p)
   print(io,body(m))

defmethod print (o:OutputStream, c:Circuit) :
   println-all(o, ["circuit " main(c) " :"])
   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), enable(e))
      (e:WritePort) : WritePort(f(mem(e)), f(index(e)), type(e), 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(name(c), f(source(c)), f(index(c)))
      (c:DefNode) : DefNode(name(c), f(value(c)))
      (c:DefInstance) : DefInstance(name(c), f(module(c)))
      (c:Conditionally) : Conditionally(f(pred(c)), conseq(c), alt(c))
      (c:Connect) : Connect(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(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(name(c),f(type(c)))
      (c:DefRegister) : DefRegister(name(c),f(type(c)))
      (c:DefMemory) : DefRegister(name(c),f(type(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)