Added on-reset

12 files changed, 265 insertions, 300 deletions

diff --git a/TODO b/TODO
index 47588213..fd5685bd 100644
--- a/TODO
+++ b/TODO
@@ -3,7 +3,6 @@
 ================================================
 
 ======== Current Tasks ========
-on-reset
 Make instances always male, flip the bundles on declaration
 dlsh,drsh
 move Infer-Widths to before vec expansion?
diff --git a/src/main/stanza/ir-parser.stanza b/src/main/stanza/ir-parser.stanza
index 57a10f0a..6e6ba1a2 100644
--- a/src/main/stanza/ir-parser.stanza
+++ b/src/main/stanza/ir-parser.stanza
@@ -129,6 +129,7 @@ OPERATORS[`convert-s] = CONVERT-S-OP
 OPERATORS[`bit-and-reduce] = BIT-AND-REDUCE-OP
 OPERATORS[`bit-or-reduce] = BIT-OR-REDUCE-OP
 OPERATORS[`bit-xor-reduce] = BIT-XOR-REDUCE-OP
+OPERATORS[`bit-not] = BIT-NOT-OP
 OPERATORS[`bit-and] = BIT-AND-OP
 OPERATORS[`bit-or] = BIT-OR-OP
 OPERATORS[`bit-xor] = BIT-XOR-OP
diff --git a/src/main/stanza/ir-utils.stanza b/src/main/stanza/ir-utils.stanza
index 67dfc89d..1131c1b3 100644
--- a/src/main/stanza/ir-utils.stanza
+++ b/src/main/stanza/ir-utils.stanza
@@ -129,12 +129,12 @@ defmethod print (o:OutputStream, op:PrimOp) :
       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"
+      CONCAT-OP : "cat"
+      BIT-SELECT-OP : "bit"
+      BITS-SELECT-OP : "bits"
    
 defmethod print (o:OutputStream, e:Expression) :
    match(e) :
diff --git a/src/main/stanza/passes.stanza b/src/main/stanza/passes.stanza
index 88871bce..06ad6a97 100644
--- a/src/main/stanza/passes.stanza
+++ b/src/main/stanza/passes.stanza
@@ -132,7 +132,6 @@ defn times (f1:Flip,f2:Flip) -> Flip :
       REVERSE : swap(f1)
      
 defn to-field (p:Port) -> Field :
-   Field(name(p),REVERSE,type(p))
    if direction(p) == OUTPUT : Field(name(p),REVERSE,type(p))
    else if direction(p) == INPUT : Field(name(p),DEFAULT,type(p))
    else : error("Shouldn't be here")
@@ -929,70 +928,70 @@ defn expand-connect-indexed (c: Circuit) -> Circuit :
 ;   This ensures proper behavior if this pass is run multiple 
 ;   times.
 
-defn initialize-registers (c:Circuit) :
-   defn to-wire-name (y:Symbol) : symbol-join([ y "$init"])
-   defn add-when (s:Stmt,h:HashTable<Symbol,Type>) -> Stmt :
-      var inits = List<Stmt>()
-      for kv in h do :
-         val refreg = WRef(key(kv),value(kv),RegKind(),FEMALE)
-         val refwire = WRef(to-wire-name(key(kv)),value(kv),NodeKind(),MALE)
-         val connect = Connect(refreg,refwire)
-         inits = append(inits,list(connect))
-      if empty?(inits) : s
-      else :
-         val pred = WRef(`reset, UIntType(IntWidth(1)), PortKind(), MALE)
-         val when-reset = Conditionally(pred,Begin(inits),Begin(List<Stmt>()))
-         Begin(list(s,when-reset))
-
-   defn rename (s:Stmt,h:HashTable<Symbol,True|False>) -> [Stmt HashTable<Symbol,Type>] :
-      val t = HashTable<Symbol,Type>(symbol-hash)
-      defn rename-expr (e:Expression) -> Expression :
-         match(map(rename-expr,e)) :
-            (e:WRegInit) : 
-               val new-name = to-wire-name(name(reg(e) as WRef))
-               WRef(new-name,type(reg(e)),RegKind(),gender(e))
-            (e) : e
-      defn rename-stmt (s:Stmt) -> Stmt :
-         match(map(rename-stmt,s)) :
-            (s:DefRegister) : 
-               if h[name(s)] :
-                  t[name(s)] = type(s)
-                  Begin(list(s,DefWire(to-wire-name(name(s)),type(s))))
-               else : s
-            (s) : map(rename-expr,s)
-      [rename-stmt(s) t]
-
-   defn init? (y:Symbol,s:Stmt) -> True|False :
-      var used? = false
-      defn has? (e:Expression) -> Expression :
-         match(map(has?,e)) :
-            (e:WRegInit) : 
-               if name(reg(e) as WRef) == y : used? = true
-            (e) : map(has?,e)
-         e
-      map(has?,s)
-      used?
-
-   defn using-init (s:Stmt,h:HashTable<Symbol,True|False>) -> Stmt :
-      match(s) :
-         (s:DefRegister) : h[name(s)] = false
-         (s) : 
-            for x in h do :
-               h[key(x)] = value(x) or init?(key(x),s)
-      map(using-init{_,h},s)
-
-   defn explicit-init-scope (s:Stmt) -> Stmt :
-      val h = HashTable<Symbol,True|False>(symbol-hash)
-      using-init(s,h)
-      ;println-debug(h)
-      val [s* t] = rename(s,h)
-      add-when(s*,t)
-     
-   Circuit(modules*, main(c)) where :
-      val modules* = 
-         for m in modules(c) map :
-            Module(name(m), ports(m), body*) where :
-               val body* = explicit-init-scope(body(m))
+;defn initialize-registers (c:Circuit) :
+;   defn to-wire-name (y:Symbol) : symbol-join([ y "$init"])
+;   defn add-when (s:Stmt,h:HashTable<Symbol,Type>) -> Stmt :
+;      var inits = List<Stmt>()
+;      for kv in h do :
+;         val refreg = WRef(key(kv),value(kv),RegKind(),FEMALE)
+;         val refwire = WRef(to-wire-name(key(kv)),value(kv),NodeKind(),MALE)
+;         val connect = Connect(refreg,refwire)
+;         inits = append(inits,list(connect))
+;      if empty?(inits) : s
+;      else :
+;         val pred = WRef(`reset, UIntType(IntWidth(1)), PortKind(), MALE)
+;         val when-reset = Conditionally(pred,Begin(inits),Begin(List<Stmt>()))
+;         Begin(list(s,when-reset))
+;
+;   defn rename (s:Stmt,h:HashTable<Symbol,True|False>) -> [Stmt HashTable<Symbol,Type>] :
+;      val t = HashTable<Symbol,Type>(symbol-hash)
+;      defn rename-expr (e:Expression) -> Expression :
+;         match(map(rename-expr,e)) :
+;            (e:WRegInit) : 
+;               val new-name = to-wire-name(name(reg(e) as WRef))
+;               WRef(new-name,type(reg(e)),RegKind(),gender(e))
+;            (e) : e
+;      defn rename-stmt (s:Stmt) -> Stmt :
+;         match(map(rename-stmt,s)) :
+;            (s:DefRegister) : 
+;               if h[name(s)] :
+;                  t[name(s)] = type(s)
+;                  Begin(list(s,DefWire(to-wire-name(name(s)),type(s))))
+;               else : s
+;            (s) : map(rename-expr,s)
+;      [rename-stmt(s) t]
+;
+;   defn init? (y:Symbol,s:Stmt) -> True|False :
+;      var used? = false
+;      defn has? (e:Expression) -> Expression :
+;         match(map(has?,e)) :
+;            (e:WRegInit) : 
+;               if name(reg(e) as WRef) == y : used? = true
+;            (e) : map(has?,e)
+;         e
+;      map(has?,s)
+;      used?
+;
+;   defn using-init (s:Stmt,h:HashTable<Symbol,True|False>) -> Stmt :
+;      match(s) :
+;         (s:DefRegister) : h[name(s)] = false
+;         (s) : 
+;            for x in h do :
+;               h[key(x)] = value(x) or init?(key(x),s)
+;      map(using-init{_,h},s)
+;
+;   defn explicit-init-scope (s:Stmt) -> Stmt :
+;      val h = HashTable<Symbol,True|False>(symbol-hash)
+;      using-init(s,h)
+;      ;println-debug(h)
+;      val [s* t] = rename(s,h)
+;      add-when(s*,t)
+;     
+;   Circuit(modules*, main(c)) where :
+;      val modules* = 
+;         for m in modules(c) map :
+;            Module(name(m), ports(m), body*) where :
+;               val body* = explicit-init-scope(body(m))
 
 ;;================ EXPAND WHENS =============================
 ; This pass does three things: remove last connect semantics,
@@ -1187,108 +1186,69 @@ defn reduce-or (l:List<Expression>) -> Expression :
 ; enables:Calculated off of assigns.
 
 ; I think I'm going to restructure this so that not all information is held in the tables, but instead, we walk the graph again, and do stuff on declarations, and delete other stuff
-defn expand-whens (body:Stmt, table:HashTable<Symbol,SymbolicValue>) -> Stmt :
-           
-
-   match(map(expand-whens{_,table,resets},s)) : 
-      (s:DefWire) : Begin $ list{s,_} $
-         val ref = WRef(name(s),type(s),NodeKind,FEMALE)
-         if has-nul?(table[n]) : 
-            println("Uninitialized: ~" % [to-string(name(s))]);TODO actually collect error
-            EmptyStmt()
-         else : Connect(ref,to-exp(table[name(s)]))
-      (s:DefRegister) : Begin $ list{s,_} $
-         val ref = WRef(name(s),type(s),RegKind,FEMALE)
-         val e = to-exp(table[name(s)])
-         match(e) : 
-            (e:False) : EmptyStmt()
-            (e:Expression) : Register(type(s),e, to-exp $ get-write-enable(table[name(s)]))
-      (s:WDefAccessor) : Begin $ 
+defn expand-whens (s:Stmt, table:HashTable<Symbol,SymbolicValue>,decs:Vector<Stmt>,cons:Vector<Stmt>) -> Stmt :
+   match(map(expand-whens{_,table,decs,cons},s)) : 
+      (s:DefNode|DefMemory) : add(decs,s)
+      (s:DefWire) : 
+         add(decs,s)
+         add{cons,_} $ { 
+            val ref = WRef(name(s),type(s),NodeKind(),FEMALE)
+            if has-nul?(table[name(s)]) : 
+               println("Uninitialized: ~" % [to-string(name(s))]);TODO actually collect error
+               EmptyStmt()
+            else : Connect(ref,to-exp(table[name(s)]) as Expression)
+            }()
+      (s:DefRegister) : 
+         add(decs,DefWire(name(s),type(s)))
+         add{cons,_} $ { 
+            val ref = WRef(name(s),type(s),RegKind(),FEMALE)
+            val e = to-exp(table[name(s)])
+            match(e) : 
+               (e:False) : EmptyStmt()
+               (e:Expression) : Connect(ref,Register(type(s),e, to-exp(optimize $ get-write-enable(table[name(s)])) as Expression))
+         }()
+      (s:WDefAccessor) : 
          val t = type(type(source(s)) as VectorType)
          val n = name(s)
-         switch {_ == gender(s)} :
-            MALE : 
-               val ref = WRef(n,t,ReadAccessorKind(),FEMALE)
-               list(DefWire(n,t),
-                    Connect(ref,ReadPort(source(s),index(s),t,to-exp $ get-read-enable(table,n))))
-            FEMALE :
-               val ref = WRef(n,t,WriteAccessorKind(),FEMALE)
-               val e = to-exp(table[n])
-               val s* = match(e) :
-                  (e:False) : 
-                     println("Uninitialized: ~" % [to-string(n)]) ;TODO actually collect error
+         add(decs,DefWire(n,t))
+         add{cons,_} $ { 
+            switch {_ == gender(s)} :
+               MALE : 
+                  val ref = WRef(n,t,ReadAccessorKind(),FEMALE)
+                  Begin $ list $ Connect(ref,ReadPort(source(s),index(s),t,get-read-enable(n,table)))
+               FEMALE :
+                  val ref = WRef(n,t,WriteAccessorKind(),FEMALE)
+                  val e = to-exp(table[n])
+                  val s* = match(e) :
+                     (e:False) : 
+                        println("Uninitialized: ~" % [to-string(n)]) ;TODO actually collect error
+                        EmptyStmt()
+                     (e:Expression) : 
+                        Connect(ref,e)
+                  val enable = (to-exp $ optimize $ get-write-enable(table[n])) as Expression
+                  val wp = WritePort(source(s),index(s),t,enable as Expression)
+                  Begin $ list(Connect(wp,ref),s*)
+         }()
+      (s:DefInstance) : 
+         add(decs,s)
+         add{cons,_} $ Begin $ 
+            for f in fields(type(module(s)) as BundleType) map :
+               if flip(f) == DEFAULT : 
+                  val n = to-symbol("~.~" % [name(s),name(f)]) ; only on inputs
+                  val x = to-symbol(split(to-string(n),'.')[0])
+                  val f = to-symbol(split(to-string(n),'.')[1])
+                  val ref = WRef(x,type(module(s)),InstanceKind(),FEMALE)
+                  val sref = WSubfield(ref,f,bundle-field-type(type(module(s)),f),FEMALE)
+                  if has-nul?(table[n]) : 
+                     println("Uninitialized: ~" % [to-string(n)]);TODO actually collect error
                      EmptyStmt()
-                  (e:Expression) : 
-                     Connect(ref,e)
-               list(DefWire(name(s),t),
-                    Connect(WritePort(source(s),index(s),t,to-exp $ get-write-enable(table[n])),ref),
-                    s*)
-      (s:DefInstance) : Begin $
-         for f in fields(type(module(s)) as BundleType) do :
-            val n = to-symbol("~.~" % [name(s),name(f)]) ; only on inputs
-            val ref = WRef(n,type(s),PortKind,FEMALE)
-            if has-nul?(table[n]) : 
-               println("Uninitialized: ~" % [to-string(name(s))]);TODO actually collect error
-               EmptyStmt()
-            else : Connect(ref,to-exp(table[name(s)]))
-         val x = to-symbol(split(to-string(n),'.')[0])
-         val f = to-symbol(split(to-string(n),'.')[1])
-         val ref = WRef(x,type(module(s)),k,FEMALE)
-         val sref = WSubfield(ref,f,bundle-field-type(type(module(s)),f),FEMALE)
-         if has-nul?(assign[n]) : println("Uninitialized: ~" % [to-string(n)]);TODO actually collect error
-         else : add(decs,Connect(sref,to-exp(assign[n])))
-      (s:Connect) : EmptyStmt()
-      (s:Conditionally) : Begin(list(conseq(s),alt(s)))
-      (s:OnReset) : EmptyStmt()
-      (s:Begin) : s
-         
+                  else : Connect(sref,to-exp(table[n]) as Expression)
+               else : EmptyStmt()
+      (s:Connect|Conditionally|OnReset|Begin|EmptyStmt) : false
+   s
 
-   for x in assign do :
-      val [n sv] = [key(x) value(x)]
-      match(kinds[n]) : 
-         (k:WriteAccessorKind) :
-            ;First create WritePort and assign from accessor-turned-wire
-            val s = stmts[n] as WDefAccessor
-            val t = type(type(source(s)) as VectorType)
-            val ref = WRef(n,t,k,MALE)
-            val wp = WritePort(source(s),index(s),t,to-exp(enables[n]))
-            add(decs,Connect(wp,ref))
-            ;If initialized, assign input to accessor-turned-wire
-            val sv = remove-nul(assign[n])
-            if sv == SVNul : println("Uninitialized: ~" % [to-string(n)]) ;TODO actually collect error
-            else : add(decs,Connect(ref,to-exp(sv)))
-         (k:ReadAccessorKind) :
-            val s = stmts[n] as WDefAccessor
-            val t = type(type(source(s)) as VectorType)
-            val ref = WRef(n,t,k,FEMALE)
-            val rp = ReadPort(source(s),index(s),t,to-exp(enables[n]))
-            add(decs,Connect(ref,rp))
-         (k:RegKind) :
-            val s = stmts[n] as DefRegister
-            val ref = WRef(n,type(s),k,FEMALE)
-            val sv = remove-nul(assign[n])
-            val reg = 
-               if sv typeof SVNul : Register(type(s),UIntValue(0,width(type(s) as ?)),zero)
-               else : Register(type(s),to-exp(sv),to-exp(enables[n]))
-            add(decs,Connect(ref,reg))
-         (k:InstanceKind) :
-            val s = stmts[n] as DefInstance
-            val x = to-symbol(split(to-string(n),'.')[0])
-            val f = to-symbol(split(to-string(n),'.')[1])
-            val ref = WRef(x,type(module(s)),k,FEMALE)
-            val sref = WSubfield(ref,f,bundle-field-type(type(module(s)),f),FEMALE)
-            if has-nul?(assign[n]) : println("Uninitialized: ~" % [to-string(n)]);TODO actually collect error
-            else : add(decs,Connect(sref,to-exp(assign[n])))
-         (k) : 
-            val s = stmts[n] as DefWire
-            val ref = WRef(n,type(s),k,FEMALE)
-            if has-nul?(assign[n]) : println("Uninitialized: ~" % [to-string(n)]);TODO actually collect error
-            else : add(decs,Connect(ref,to-exp(assign[n])))
-   Begin(to-list(decs))
-
-defn get-enables (assign:HashTable<Symbol,SymbolicValue>,
-                  kinds:HashTable<Symbol,Kind>) -> HashTable<Symbol,SymbolicValue> :
-   defn get-read-enable (sym:Symbol,sv:SymbolicValue) -> Expression :
+defn get-read-enable (sym:Symbol,table:HashTable<Symbol,SymbolicValue>) -> Expression :
+   defn get-single-read-enable (sym:Symbol,sv:SymbolicValue) -> Expression :
       defn active (e:Expression) -> True|False : 
          match(e) :
             (e:WRef) : name(e) == sym
@@ -1300,30 +1260,30 @@ defn get-enables (assign:HashTable<Symbol,SymbolicValue>,
             if active(exp(sv)) : one 
             else : zero
          (sv: SVMux) : 
-            val e0 = get-read-enable(sym,SVExp(pred(sv)))
-            val e1 = get-read-enable(sym,conseq(sv))
-            val e2 = get-read-enable(sym,alt(sv))
+            val e0 = get-single-read-enable(sym,SVExp(pred(sv)))
+            val e1 = get-single-read-enable(sym,conseq(sv))
+            val e2 = get-single-read-enable(sym,alt(sv))
             if e1 == e2 : OR(e0,e1)
             else : OR(e0,OR(AND(pred(sv),e1),AND(NOT(pred(sv)),e2)))
+   DoPrim{BIT-OR-OP,_,list(),UIntType(IntWidth(1))} $ to-list $
+      for y in table stream : get-single-read-enable(sym,value(y))
+
+defn get-write-enable (sv:SymbolicValue) -> SymbolicValue :
+   match(map(get-write-enable,sv)) :
+      (sv: SVExp) : SVExp(one)
+      (sv: SVNul) : SVExp(zero)
+      (sv) : sv
 
-   defn get-write-enable (sv:SymbolicValue) -> SymbolicValue :
-      match(map(get-write-enable,sv)) :
-         (sv: SVExp) : SVExp(one)
-         (sv: SVNul) : SVExp(zero)
-         (sv) : sv
-
-   val enables = HashTable<Symbol,SymbolicValue>(symbol-hash)
-   for x in assign do :
-      val sym = key(x)
-      match(kinds[sym]) :
-         (k:ReadAccessorKind) : 
-            enables[sym] = SVExp{_} $ reduce-or{_} $ to-list{_} $
-               for y in assign stream :
-                  get-read-enable(sym,value(y))
-         (k:WriteAccessorKind) : enables[sym] = get-write-enable(value(x))
-         (k:RegKind) : enables[sym] = get-write-enable(value(x))
-         (k) : k
-   enables
+defn merge-resets (assign:HashTable<Symbol,SymbolicValue>, resets:HashTable<Symbol,SymbolicValue>) -> HashTable<Symbol,SymbolicValue> :
+   val table = HashTable<Symbol,SymbolicValue>(symbol-hash)
+   val reset = WRef(`reset, UnknownType(), PortKind(), MALE)
+   for i in get-unique-keys(list(assign,resets)) do :
+      table[i] = match(get?(assign,i,false),get?(resets,i,false)) : 
+            (a:SymbolicValue,r:SymbolicValue) : SVMux(reset,r,a)
+            (a:SymbolicValue,r:False) : a
+            (a:False,r:SymbolicValue) : SVMux(reset,r,SVNul())
+            (a:False,r:False) : error("Shouldn't be here")
+   table
 
 defn build-tables (s:Stmt, 
                    assign:HashTable<Symbol,SymbolicValue>,
@@ -1333,26 +1293,22 @@ defn build-tables (s:Stmt,
    match(s) :
       (s:DefWire) :
          assign[name(s)] = SVNul()
-         resets[name(s)] = SVNul()
          flattn[name(s)] = true
       (s:DefRegister|WDefAccessor) : 
          assign[name(s)] = SVNul()
-         resets[name(s)] = SVNul()
          flattn[name(s)] = false
-      (s:DefInstance) :
+      (s:DefInstance) : ;TODO only add instance input ports. This probably involves correcting instance genders
          for f in fields(type(module(s)) as BundleType) do :
-            if 
-            val n = to-symbol("~.~" % [name(s),name(f)]) ; only on inputs
-            assign[n] = SVNul()
-            resets[n] = SVNul()
-            flattn[name(s)] = true
-      (s:DefMemory|DefNode) : false
+            if flip(f) == DEFAULT : 
+               println-all-debug(["Instance: " s " has input " f])
+               val n = to-symbol("~.~" % [name(s),name(f)]) ; only on inputs
+               assign[n] = SVNul()
+               flattn[n] = true
       (s:Conditionally) :
          defn combine (flattn:HashTable<Symbol,True|False>,
                        table-c:HashTable<Symbol,SymbolicValue>,
                        table-a:HashTable<Symbol,SymbolicValue>,
-                       i:Symbol
-                       ) -> SymbolicValue :
+                       i:Symbol) -> SymbolicValue|False :
             match(get?(table-c,i,false),get?(table-a,i,false)) : 
                   (c:SymbolicValue,a:SymbolicValue) : 
                      if c == a : c
@@ -1363,7 +1319,7 @@ defn build-tables (s:Stmt,
                   (c:False,a:SymbolicValue) :
                      if flattn[i] : a
                      else : SVMux(pred(s),SVNul(),a)
-                  (c:False,a:False) : error("Shouldn't be here")
+                  (c:False,a:False) : false
 
          val assign-c = deepcopy(assign)
          val assign-a = deepcopy(assign)
@@ -1372,8 +1328,11 @@ defn build-tables (s:Stmt,
          build-tables(conseq(s),assign-c,resets-c,flattn)
          build-tables(alt(s),assign-a,resets-a,flattn)
          for i in get-unique-keys(list(assign-c,assign-a)) do :
-            assign[i] = combine(flattn,assign-c,assign-a,i)
-            resets[i] = combine(flattn,resets-c,resets-a,i)
+            assign[i] = combine(flattn,assign-c,assign-a,i) as SymbolicValue
+            val r = combine(flattn,resets-c,resets-a,i)
+            match(r) : 
+               (r:SymbolicValue) : resets[i] = r
+               (r) : false
          ;println-debug("TABLE-C")
          ;for x in assign-c do : println-debug(x)
          ;println-debug("TABLE-A")
@@ -1385,9 +1344,10 @@ defn build-tables (s:Stmt,
             (e:WRef) : name(e)
             (e:WSubfield) : symbol-join([name(exp(e) as ?) `. name(e)])
             (e) : error("Shouldn't be here with ~" % [e])
-         if c typeof Connect : assign[key*] = SVExp(exp(s))
-         if c typeof OnReset : resets[key*] = SVExp(exp(s))
-      (s:Begin) : for s* in body(s) do: build-tables(s*,assign,resets)
+         if s typeof Connect : assign[key*] = SVExp(exp(s))
+         if s typeof OnReset : resets[key*] = SVExp(exp(s))
+      (s:Begin) : for s* in body(s) do: build-tables(s*,assign,resets,flattn)
+      (s:DefMemory|DefNode|EmptyStmt) : false
          
 defn expand-whens (m:Module) -> Module :
    val assign = HashTable<Symbol,SymbolicValue>(symbol-hash)
@@ -1397,22 +1357,26 @@ defn expand-whens (m:Module) -> Module :
    for p in ports(m) do :
       if direction(p) == OUTPUT :
          assign[name(p)] = SVNul()
-         resets[name(p)] = SVNul()
          flattn[name(p)] = false
          
-   build-tables(body(m),assign,resets)
+   build-tables(body(m),assign,resets,flattn)
    for x in assign do : assign[key(x)] = optimize(value(x))
    for x in resets do : resets[key(x)] = optimize(value(x))
    ;val enables = get-enables(assign,kinds)
    ;for x in enables do : enables[key(x)] = optimize(value(x))
          
-   println-debug("Assigns")
+   println-debug("====== Assigns ======")
    for x in assign do : println-debug(x)
-   println-debug("Resets")
+   println-debug("====== Resets ======")
    for x in resets do : println-debug(x)
 
    val table = merge-resets(assign,resets)
-   Module(name(m),ports(m),expand-whens(body(m),table))
+   println-debug("====== Table ======")
+   for x in table do : println-debug(x)
+   val decs = Vector<Stmt>()
+   val cons = Vector<Stmt>()
+   expand-whens(body(m),table,decs,cons)
+   Module(name(m),ports(m),Begin(append(to-list(decs),to-list(cons))))
 
 defn expand-whens (c:Circuit) -> Circuit :
    Circuit(modules*, main(c)) where :
@@ -2103,7 +2067,6 @@ public defn run-passes (c: Circuit, p: List<Char>,file:String) :
    if contains(p,'X') or contains(p,'g') : do-stage("Expand Accessors", expand-accessors)
    if contains(p,'X') or contains(p,'h') : do-stage("Lower To Ground", lower-to-ground)
    if contains(p,'X') or contains(p,'i') : do-stage("Expand Indexed Connects", expand-connect-indexed)
-   if contains(p,'X') or contains(p,'j') : do-stage("Initialize Registers", initialize-registers)
    if contains(p,'X') or contains(p,'k') : do-stage("Expand Whens", expand-whens)
    if contains(p,'X') or contains(p,'l') : do-stage("Infer Widths", infer-widths)
    if contains(p,'X') or contains(p,'m') : do-stage("Inline Instances", inline-instances)
diff --git a/test/passes/expand-whens/bundle-init.fir b/test/passes/expand-whens/bundle-init.fir
new file mode 100644
index 00000000..48336c93
--- /dev/null
+++ b/test/passes/expand-whens/bundle-init.fir
@@ -0,0 +1,25 @@
+; RUN: firrtl -i %s -o %s.flo -x abcdefghijk -p cd | tee %s.out | FileCheck %s
+; CHECK: Expand Whens
+circuit top :
+   module A :
+      reg r : { x : UInt, flip y : UInt}
+      wire a : UInt
+      wire b : UInt
+      wire w : { x : UInt, flip y : UInt}
+      a := UInt(1)
+      b := UInt(2)
+
+      w.x := b
+      w.y := a
+      r.x := a
+      r.y := b
+      on-reset r := w
+
+; CHECK: r$x := Register(mux-uu(reset, w$x, a), UInt(1))
+; CHECK: r$y := Register(b, UInt(1))
+; CHECK: a := UInt(1)
+; CHECK: b := UInt(2)
+; CHECK: w$x := b
+; CHECK: w$y := mux-uu(reset, r$y, a)   
+      
+; CHECK: Finished Expand Whens
diff --git a/test/passes/initialize-regs/nested-whens.fir b/test/passes/expand-whens/nested-whens.fir
index 28cbc53d..8185dade 100644
--- a/test/passes/initialize-regs/nested-whens.fir
+++ b/test/passes/expand-whens/nested-whens.fir
@@ -1,5 +1,5 @@
-; RUN: firrtl -i %s -o %s.flo -x abcdefghij -p c | tee %s.out | FileCheck %s
-; CHECK: Done!
+; RUN: firrtl -i %s -o %s.flo -x abcdefghijk -p c | tee %s.out | FileCheck %s
+; CHECK: Expand Whens
 circuit top :
    module A :
       wire p : UInt
@@ -10,17 +10,15 @@ circuit top :
       wire x : UInt
       wire y : UInt
       wire z : UInt
+      wire w : UInt
 
       on-reset r := w
-      when p
+      when p :
         on-reset r := x
         r := a
-      when q
+      when q :
         on-reset r := y
         r := b
       r := z
-      
-; CHECK: r := z
-; CHECK: when reset
-; CHECK:    r := q?b:(p?a:w)
-
+; CHECK: r := Register(mux-uu(reset, mux-uu(q, y, mux-uu(p, x, w)), z), UInt(1))
+; CHECK: Finished Expand Whens
diff --git a/test/passes/infer-widths/gcd.fir b/test/passes/infer-widths/gcd.fir
index 0e9b5fed..3cd5c542 100644
--- a/test/passes/infer-widths/gcd.fir
+++ b/test/passes/infer-widths/gcd.fir
@@ -1,4 +1,4 @@
-; RUN: firrtl -i %s -o %s.flo -x abcdefghijkl -p cTd | tee %s.out | FileCheck %s
+; RUN: firrtl -i %s -o %s.flo -x abcdefghijkl -p cd | tee %s.out | FileCheck %s
 
 ;CHECK: Infer Widths
 circuit top :
diff --git a/test/passes/initialize-regs/bundle-init.fir b/test/passes/initialize-regs/bundle-init.fir
deleted file mode 100644
index 7e9af8df..00000000
--- a/test/passes/initialize-regs/bundle-init.fir
+++ /dev/null
@@ -1,21 +0,0 @@
-; RUN: firrtl -i %s -o %s.flo -x abcdefghij -p c | tee %s.out | FileCheck %s
-; CHECK: Done!
-circuit top :
-   module A :
-      reg r : { x : UInt, flip y : UInt}
-      wire a : UInt
-      wire b : UInt
-      wire w : { x : UInt, flip y : UInt}
-
-      r.x := a
-      r.y := b
-      on-reset r := w
-      
-; CHECK: reg r : { x, flip y}
-; CHECK: r.x := a
-; CHECK: r.y := b
-; CHECK: when reset :
-; CHECK:    r.x := w.x
-; CHECK:    w.y := r.y
-
-
diff --git a/test/passes/jacktest/MemorySearch.fir b/test/passes/jacktest/MemorySearch.fir
index 955d44f2..60b62ac7 100644
--- a/test/passes/jacktest/MemorySearch.fir
+++ b/test/passes/jacktest/MemorySearch.fir
@@ -3,43 +3,43 @@
 
 circuit MemorySearch : 
   module MemorySearch : 
-    input target : UInt(4)
-    output address : UInt(3)
-    input en : UInt(1)
-    output done : UInt(1)
+    input target : UInt<4>
+    output address : UInt<3>
+    input en : UInt<1>
+    output done : UInt<1>
     
-    node T_35 = UInt(0, 3)
-    reg index : UInt(3)
-    index.init := T_35
-    node T_36 = UInt(0, 1)
-    node T_37 = UInt(4, 3)
-    node T_38 = UInt(15, 4)
-    node T_39 = UInt(14, 4)
-    node T_40 = UInt(2, 2)
-    node T_41 = UInt(5, 3)
-    node T_42 = UInt(13, 4)
-    wire elts : UInt(1)[7]
-    elts.0 := Pad(T_36,?)
-    elts.1 := Pad(T_37,?)
-    elts.2 := Pad(T_38,?)
-    elts.3 := Pad(T_39,?)
-    elts.4 := Pad(T_40,?)
-    elts.5 := Pad(T_41,?)
-    elts.6 := Pad(T_42,?)
+    node T_35 = UInt<3>(0)
+    reg index : UInt<3>
+    on-reset index := T_35
+    node T_36 = UInt<1>(0)
+    node T_37 = UInt<3>(4)
+    node T_38 = UInt<4>(15)
+    node T_39 = UInt<4>(14)
+    node T_40 = UInt<2>(2)
+    node T_41 = UInt<3>(5)
+    node T_42 = UInt<4>(13)
+    wire elts : UInt<1>[7]
+    elts[0] := Pad(T_36,?)
+    elts[1] := Pad(T_37,?)
+    elts[2] := Pad(T_38,?)
+    elts[3] := Pad(T_39,?)
+    elts[4] := Pad(T_40,?)
+    elts[5] := Pad(T_41,?)
+    elts[6] := Pad(T_42,?)
     accessor elt = elts[index]
     node T_43 = bit-not(en)
     node T_44 = eq(Pad(elt,?), Pad(target,?))
-    node T_45 = UInt(7, 3)
+    node T_45 = UInt<3>(7)
     node T_46 = eq(Pad(index,?), Pad(T_45,?))
     node T_47 = bit-or(T_44, T_46)
     node end = bit-and(T_43, T_47)
     when en : 
-      node T_48 = UInt(0, 1)
+      node T_48 = UInt<1>(0)
       index := Pad(T_48,?)
     else : 
       node T_49 = bit-not(end)
       when T_49 : 
-        node T_50 = UInt(1, 1)
+        node T_50 = UInt<1>(1)
         node T_51 = add-wrap(Pad(index,?), Pad(T_50,?))
         index := Pad(T_51,?)
     done := Pad(end,?)
diff --git a/test/passes/jacktest/RegisterVecShift.fir b/test/passes/jacktest/RegisterVecShift.fir
index 733e2036..832bd279 100644
--- a/test/passes/jacktest/RegisterVecShift.fir
+++ b/test/passes/jacktest/RegisterVecShift.fir
@@ -3,35 +3,35 @@
 
 circuit RegisterVecShift : 
   module RegisterVecShift : 
-    input load : UInt(1)
-    output out : UInt(4)
-    input shift : UInt(1)
-    input ins : UInt(4)[4]
+    input load : UInt<1>
+    output out : UInt<4>
+    input shift : UInt<1>
+    input ins : UInt<4>[4]
     
-    reg delays : UInt(4)[4]
+    reg delays : UInt<4>[4]
     when reset : 
-      node T_38 = UInt(0, 4)
-      node T_39 = UInt(0, 4)
-      node T_40 = UInt(0, 4)
-      node T_41 = UInt(0, 4)
-      wire T_42 : UInt(4)[4]
-      T_42.0 := T_38
-      T_42.1 := T_39
-      T_42.2 := T_40
-      T_42.3 := T_41
+      node T_38 = UInt<4>(0)
+      node T_39 = UInt<4>(0)
+      node T_40 = UInt<4>(0)
+      node T_41 = UInt<4>(0)
+      wire T_42 : UInt<4>[4]
+      T_42[0] := T_38
+      T_42[1] := T_39
+      T_42[2] := T_40
+      T_42[3] := T_41
       delays := T_42
-    node T_43 = UInt(5, 3)
+    node T_43 = UInt<3>(5)
     node T_44 = bit-and(Pad(T_43,?), Pad(load,?))
-    node T_45 = UInt(4, 3)
+    node T_45 = UInt<3>(4)
     node T_46 = eq(Pad(T_44,?), Pad(T_45,?))
     when T_46 : 
-      delays.0 := Pad(ins.0,?)
-      delays.1 := Pad(ins.1,?)
-      delays.2 := Pad(ins.2,?)
-      delays.3 := Pad(ins.3,?)
+      delays[0] := Pad(ins[0],?)
+      delays[1] := Pad(ins[1],?)
+      delays[2] := Pad(ins[2],?)
+      delays[3] := Pad(ins[3],?)
     else : when shift : 
-        delays.0 := Pad(ins.0,?)
-        delays.1 := Pad(delays.0,?)
-        delays.2 := Pad(delays.1,?)
-        delays.3 := Pad(delays.2,?)
-    out := Pad(delays.3,?)
+        delays[0] := Pad(ins[0],?)
+        delays[1] := Pad(delays[0],?)
+        delays[2] := Pad(delays[1],?)
+        delays[3] := Pad(delays[2],?)
+    out := Pad(delays[3],?)
diff --git a/test/passes/jacktest/Tlb.fir b/test/passes/jacktest/Tlb.fir
index 35442ac8..b458ac4a 100644
--- a/test/passes/jacktest/Tlb.fir
+++ b/test/passes/jacktest/Tlb.fir
@@ -2,13 +2,13 @@
 ; CHECK: Done!
 circuit Tbl : 
   module Tbl : 
-    output o : UInt(16)
-    input i : UInt(16)
-    input d : UInt(16)
-    input we : UInt(1)
+    output o : UInt<16>
+    input i : UInt<16>
+    input d : UInt<16>
+    input we : UInt<1>
     
-    mem m : UInt(10)[256]
-    node T_12 = UInt(0, 1)
+    mem m : UInt<10>[256]
+    node T_12 = UInt<1>(0)
     o := Pad(T_12,?)
     when we : 
       accessor T_13 = m[i]
diff --git a/test/passes/lower-to-ground/register.fir b/test/passes/lower-to-ground/register.fir
index ecc427d9..918710a5 100644
--- a/test/passes/lower-to-ground/register.fir
+++ b/test/passes/lower-to-ground/register.fir
@@ -15,7 +15,7 @@
       ; CHECK: reg r1$y : SInt
       ; CHECK: wire q$x : UInt
       ; CHECK: wire q$y : SInt
-      ; CHECK: r1$x.init := q$x
-      ; CHECK: q$y := r1$y.init
+      ; CHECK: on-reset r1$x := q$x
+      ; CHECK: on-reset q$y := r1$y   
 
 ; CHECK: Finished Lower To Ground