Correctly do when expansion, minus enables and outputting lowered form

7 files changed, 232 insertions, 205 deletions

diff --git a/notes/notes.03.18.15.txt b/notes/notes.03.18.15.txt
index 7fdb1c3b..f9acf7f5 100644
--- a/notes/notes.03.18.15.txt
+++ b/notes/notes.03.18.15.txt
@@ -45,7 +45,22 @@ if w is a wire:
 merge {r=>x, w=>y} with {r=>x} under p : {r=>svmux(p,x,x), w=>y}
 
 if s is a reg:
-merge {r=>x,s=>y} with {r=>x} under p : {r=>svmux(p,x,x), s=>svmux(p,y,void)}
+merge {r=>x,s=>y} with {r=>x} under p : {r=>svmux(p,x,x), s=>svmux(p,y,void)} ;this is to correctly calculate the ENABLE signal! when actually calculating the input, we will reduce it
+
+wire r     {r=>VOID}
+r := x     {r=>x}
+when p     {r=>x}
+  reg s    {r=>x,s=>VOID}
+  s := y   {r=>x,s=>y}
+  wire w   {r=>x,s=>y,w=>VOID}
+  w := y   {r=>x,s=>y,w=>y}
+else 
+  emptystmt {r=>x}
+;merge table-c with table-a
+;get unique keys of table-a + table-c
+; (r,s,w)
+{r=>(p?x:x),s=>
+
 
 
 
diff --git a/src/main/stanza/firrtl-ir.stanza b/src/main/stanza/firrtl-ir.stanza
index 3d8a4ba9..977c8ca1 100644
--- a/src/main/stanza/firrtl-ir.stanza
+++ b/src/main/stanza/firrtl-ir.stanza
@@ -147,11 +147,12 @@ public defstruct WritePort <: Expression :
    index: Expression
    type: Type [multi => false]
    enable: Expression
+public defstruct Register <: Expression :
+   type: Type [multi => false]
+   value: Expression
+   enable: Expression
 
 public definterface Stmt
-public defstruct LetRec <: Stmt :
-   entries: List<KeyValue<Symbol, Element>>
-   body: Stmt
 public defstruct DefWire <: Stmt :
    name: Symbol
    type: Type
@@ -182,24 +183,6 @@ public defstruct Connect <: Stmt :
    exp: Expression
 public defstruct EmptyStmt <: Stmt
 
-public definterface Element
-public defmulti type (e:Element) -> Type
-
-public defstruct Register <: Element :
-   type: Type [multi => false]
-   value: Expression
-   enable: Expression
-public defstruct Memory <: Element :
-   type: Type [multi => false]
-   writers: List<WritePort>
-public defstruct Node <: Element :
-   type: Type [multi => false]
-   value: Expression
-public defstruct Instance <: Element :
-   type: Type [multi => false]
-   module: Expression
-   ports: List<KeyValue<Symbol,Expression>>
-
 public definterface Type
 public defstruct UIntType <: Type :
    width: Width
diff --git a/src/main/stanza/ir-parser.stanza b/src/main/stanza/ir-parser.stanza
index ae60194c..8f0cc8e3 100644
--- a/src/main/stanza/ir-parser.stanza
+++ b/src/main/stanza/ir-parser.stanza
@@ -128,8 +128,6 @@ rd.defsyntax firrtl :
          DefAccessor(ut(name), source, index)
       ((?body:#comm ...)) :
          Begin(body)
-      (letrec : (?elems:#element ...) in : ?body:#comm) :
-         LetRec(elems, body)
       (?x:#exp := ?y:#exp) :
          Connect(x, y)
       (?c:#comm/when) :
@@ -143,16 +141,6 @@ rd.defsyntax firrtl :
       (when ?pred:#exp : ?conseq:#comm) :
          Conditionally(pred, conseq, EmptyStmt())
 
-   defrule element :
-      (reg ?name:#symbol : ?type:#type = Register (@do ?value:#exp ?en:#exp)) :
-         ut(name) => Register(type, value, en)
-      (node ?name:#symbol : ?type:#type = ?exp:#exp) :
-         ut(name) => Node(type, exp)
-      (inst ?name:#symbol = Instance (@do ?module:#exp
-                               (?names:#symbol => ?values:#exp @...))) :
-         val ports = map({ut(_) => _}, names, values)
-         ut(name) => Instance(UnknownType(), module, ports)
-
    defrule exp :
       (?x:#exp . ?f:#int) :
          Index(x, ut(f), UnknownType())
diff --git a/src/main/stanza/ir-utils.stanza b/src/main/stanza/ir-utils.stanza
index aebaf0f9..4d6c0235 100644
--- a/src/main/stanza/ir-utils.stanza
+++ b/src/main/stanza/ir-utils.stanza
@@ -5,7 +5,7 @@ defpackage firrtl.ir-utils :
 
 ;============== DEBUG STUFF =============================
 
-public defmulti print-debug (o:OutputStream, e:Expression|Stmt|Type|Element|Port|Field) -> False
+public defmulti print-debug (o:OutputStream, e:Expression|Stmt|Type|Port|Field) -> False
 
 ;============== PRINTERS ===================================
 
@@ -144,12 +144,6 @@ defmethod print (o:OutputStream, e:Expression) :
 defmethod print (o:OutputStream, c:Stmt) :
    val io = IndentedStream(o, 3)
    match(c) :
-      (c:LetRec) :
-         println(o, "let : ")
-         for entry in entries(c) do :
-            println-all(io,[key(entry) " = " value(entry)])
-         println(o, "in :")
-         print(io, body(c))
       (c:DefWire) :
          print-all(o,["wire " name(c) " : " type(c)])
       (c:DefRegister) : 
@@ -174,25 +168,9 @@ defmethod print (o:OutputStream, c:Stmt) :
       (c:Connect) :
          print-all(o, [loc(c) " := " exp(c)])
       (c:EmptyStmt) :
-         print(o, "ip")
+         print(o, "$empty$")
    print-debug(o,c)
 
-defmethod print (o:OutputStream, e:Element) :
-   match(e) :
-      (e:Register) :
-         print-all(o, ["Register(" type(e) ", " value(e) ", " enable(e) ")"])
-      (e:Memory) :
-         print-all(o, ["Memory(" type(e) ", "])
-         print-all(o, join(writers(e), ", "))
-         print(o, ")")   
-      (e:Node) :
-         print-all(o, ["Node(" type(e) ", " value(e) ")"])
-      (e:Instance) :
-         print-all(o, ["Instance(" module(e) ", "])
-         print-all(o, join(ports(e), ", "))
-         print(o, ")")
-   print-debug(o,e)
-
 defmethod print (o:OutputStream, t:Type) :
    match(t) :
       (t:UnknownType) :
@@ -258,26 +236,9 @@ defmethod map (f: Expression -> Expression, e:Expression) -> Expression :
       (e:WritePort) : WritePort(f(mem(e)), f(index(e)), type(e), enable(e))
       (e) : e
 
-public defmulti map<?T> (f: Expression -> Expression, e:?T&Element) -> T
-defmethod map (f: Expression -> Expression, e:Element) -> Element :
-   match(e) :
-      (e:Register) :
-         Register(type(e), f(value(e)), f(enable(e)))
-      (e:Memory) : e
-      (e:Node) :
-         Node(type(e), f(value(e)))
-      (e:Instance) :
-         val ports* = for p in ports(e) map :
-                         key(p) => f(value(p))
-         Instance(type(e), f(module(e)), ports*)
-
 public defmulti map<?T> (f: Expression -> Expression, c:?T&Stmt) -> T
 defmethod map (f: Expression -> Expression, c:Stmt) -> Stmt :
    match(c) :
-      (c:LetRec) :
-         val entries* = for entry in entries(c) map :
-                           key(entry) => map(f, value(entry))
-         LetRec(entries*, body(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)))
@@ -288,7 +249,6 @@ defmethod map (f: Expression -> Expression, c:Stmt) -> Stmt :
 public defmulti map<?T> (f: Stmt -> Stmt, c:?T&Stmt) -> T
 defmethod map (f: Stmt -> Stmt, c:Stmt) -> Stmt :
    match(c) :
-      (c:LetRec) : LetRec(entries(c), f(body(c)))
       (c:Conditionally) : Conditionally(pred(c), f(conseq(c)), f(alt(c)))
       (c:Begin) : Begin(map(f, body(c)))
       (c) : c
@@ -296,7 +256,6 @@ defmethod map (f: Stmt -> Stmt, c:Stmt) -> Stmt :
 public defmulti children (c:Stmt) -> List<Stmt>
 defmethod children (c:Stmt) :
    match(c) :
-      (c:LetRec) : list(body(c))
       (c:Conditionally) : list(conseq(c), alt(c))
       (c:Begin) : body(c)
       (c) : List()
diff --git a/src/main/stanza/passes.stanza b/src/main/stanza/passes.stanza
index 324630f0..8380af0a 100644
--- a/src/main/stanza/passes.stanza
+++ b/src/main/stanza/passes.stanza
@@ -156,27 +156,27 @@ defmethod print (o:OutputStream, k:Kind) :
       (k:ReadAccessorKind) : "racc"
       (k:WriteAccessorKind) : "wacc"
 
-defn hasGender (e:Expression|Stmt|Type|Element|Port|Field) :
+defn hasGender (e:Expression|Stmt|Type|Port|Field) :
    e typeof WRef|WSubfield|WIndex|WDefAccessor|WRegInit
 
-defn hasWidth (e:Expression|Stmt|Type|Element|Port|Field) :
+defn hasWidth (e:Expression|Stmt|Type|Port|Field) :
    e typeof UIntType|SIntType|UIntValue|SIntValue|WRegInit
 
-defn hasType (e:Expression|Stmt|Type|Element|Port|Field) :
+defn hasType (e:Expression|Stmt|Type|Port|Field) :
    e typeof Ref|Subfield|Index|DoPrim|WritePort|ReadPort|WRef|WSubfield
       |WIndex|DefWire|DefRegister|DefMemory|Register
-      |Memory|Node|Instance|VectorType|Port|Field|WRegInit
+      |VectorType|Port|Field|WRegInit
 
-defn hasKind (e:Expression|Stmt|Type|Element|Port|Field) :
+defn hasKind (e:Expression|Stmt|Type|Port|Field) :
    e typeof WRef
 
-defn any-debug? (e:Expression|Stmt|Type|Element|Port|Field) : 
+defn any-debug? (e:Expression|Stmt|Type|Port|Field) : 
    (hasGender(e) and PRINT-GENDERS) or
    (hasType(e) and PRINT-TYPES) or
    (hasWidth(e) and PRINT-WIDTHS) or
    (hasKind(e) and PRINT-KINDS)
 
-defmethod print-debug (o:OutputStream, e:Expression|Stmt|Type|Element|Port|Field) :
+defmethod print-debug (o:OutputStream, e:Expression|Stmt|Type|Port|Field) :
    defn wipe-width (t:Type) -> Type : 
      match(t) :
         (t:UIntType) : UIntType(UnknownWidth())
@@ -267,9 +267,6 @@ defn resolve-kinds (c:Circuit) :
    defn find (m:Module, kinds:HashTable<Symbol,Kind>) :
       defn find-stmt (s:Stmt) -> Stmt :
          match(s) :
-            (s:LetRec) : 
-               for e in entries(s) do :
-                  kinds[key(e)] = get-elem-kind(value(e))
             (s:DefWire) : kinds[name(s)] = NodeKind()
             ;TODO add DefNode
             (s:DefRegister) : kinds[name(s)] = RegKind()
@@ -279,11 +276,6 @@ defn resolve-kinds (c:Circuit) :
             (s) : false
          map(find-stmt,s)
 
-      defn get-elem-kind (e:Element) :
-         match(e) :
-            (e: Memory) : StructuralMemKind()
-            (e) : NodeKind()
-
       kinds[name(m)] = ModuleKind()
       for p in ports(m) do :
          kinds[name(p)] = PortKind()
@@ -506,9 +498,8 @@ defn infer-exp-types (e:Expression, l:List<KeyValue<Symbol,Type>>) -> Expression
       (e:WritePort) : WritePort(mem(e),index(e),get-vector-subtype(type(mem(e))),enable(e))
       (e:UIntValue|SIntValue) : e
 
-defn infer-types (s:Stmt, l:List<KeyValue<Symbol,Type>>) -> [Stmt, List<KeyValue<Symbol,Type>>] :
+defn infer-types (s:Stmt, l:List<KeyValue<Symbol,Type>>) -> [Stmt List<KeyValue<Symbol,Type>>] :
    match(map(infer-exp-types{_,l},s)) :
-      (s:LetRec) : [s,l] ;TODO, this is wrong but we might be getting rid of letrecs?
       (s:Begin) : 
          var env = l
          val body* = 
@@ -868,8 +859,13 @@ defn lower-module (m:Module,table:HashTable<Symbol,List<KeyValue<Expression,Flip
       match(stmt) :
          (s:DefWire) : table[name(s)] = get-entries(name(s),type(s))
          (s:DefRegister) : 
-            table[name(s)] = get-entries(name(s),type(s))
-            table[to-symbol("~.init" % [to-string(name(s))])] = get-entries(to-symbol("~#init" % [to-string(name(s))]),type(s)) ; TODO, we are keeping WRegInits around
+            val regs = get-entries(name(s),type(s))
+            val init-sym = symbol-join([name(s),`.init])
+            val init-regs = for r in regs map : 
+               val [e f] = [key(r) value(r)]
+               WRegInit(e,symbol-join([name(e),`.init]),type(e),gender(e)) => f
+            table[name(s)] = regs
+            table[init-sym] = init-regs
          (s:DefInstance) : 
             val r = WRef(name(s),type(module(s)),InstanceKind(),FEMALE)
             val ports = table[name(module(s) as WRef)]
@@ -969,6 +965,24 @@ public defstruct SVMux <: SymbolicValue :
    alt : SymbolicValue
 public defstruct SVNul <: SymbolicValue
 
+defstruct SSV :
+   stmt : Stmt
+   sv : SymbolicValue
+
+defn is-equal (a:SymbolicValue,b:SymbolicValue) -> True|False :
+   defn ex-equal (e:Expression,d:Expression) -> True|False :
+      match(e,d) : 
+         (e:WRef,d:WRef) : name(e) == name(d)
+         (e,d) : false
+   match(a,b) :
+      (a:SVNul,b:SVNul) : true
+      (a:SVExp,b:SVExp) : ex-equal(exp(a), exp(b))
+      (a:SVMux,b:SVMux) : ex-equal(pred(a),pred(b)) and is-equal(conseq(a),conseq(b)) and is-equal(alt(a),alt(b))
+      (a,b) : false
+
+defmethod print (o:OutputStream, ssv:SSV) :
+   print-all(o, ["[ {" stmt(ssv) "} :: {" sv(ssv) "} ]"])
+
 defmethod print (o:OutputStream, sv:SymbolicValue) :
    match(sv) :
       (sv: SVExp) : print(o, exp(sv))
@@ -981,131 +995,187 @@ defmethod map (f: SymbolicValue -> SymbolicValue, sv:SymbolicValue) -> SymbolicV
       (sv: SVMux) : SVMux(pred(sv),f(conseq(sv)),f(alt(sv)))
       (sv) : sv
 
-defn new-vec () -> Vector<KeyValue<Symbol,[Stmt SymbolicValue]>> :
-   Vector<KeyValue<Symbol,[Stmt SymbolicValue]>>()
-
-;defn expand-whens (table:List<KeyValue<Symbol,[Stmt SymbolicValue]>>) -> HashMap<Symbol,SymbolicValue> :
-;   val enables = HashMap<Symbol,SymbolicValue>(symbol-hash)
-;   for x in table do :
-;      val [s sv] = value(x)
-;      val sym = key(x)
-;      val sv* = 
-;         match(s) :
-;            (s:WDefAccessor) : 
-;               switch {_ == gender(s)} :
-;                  FEMALE : 
-;                     defn get-enable (sv:SymbolicValue) -> SymbolicValue :
-;                        match(map(get-enable,sv)) :
-;                           (sv: SVExp) : SVExp(UIntValue(1,1))
-;                           (sv: SVNul) : SVExp(UIntValue(0,1))
-;                           (sv) : sv
-;                     get-enable(sv)
-;                  MALE :
-;                     for y in table do :
-;                        val [ys ysv] = value(y)
-;                        if contains?(key(x),ysv) :
-;                           
-;                     defn get-enable (sym:Symbol,sv:SymbolicValue) -> SymbolicValue :
-;                        defn active (e:Expression) -> 
-;                        match(map(get-enable,sv)) :
-;                           (sv: SVExp) : exp
-;                           (sv: SVNul) : SVExp(UIntValue(0,1))
-;                           (sv) : sv
-                              
-
-defn optimize-table (table:List<KeyValue<Symbol,[Stmt SymbolicValue]>>) -> List<KeyValue<Symbol,[Stmt SymbolicValue]>>:
+defn new-table () -> Vector<KeyValue<Symbol,SSV>> :
+   Vector<KeyValue<Symbol,SSV>>()
+
+defn get-enables (table:List<KeyValue<Symbol,SSV>>) -> HashTable<Symbol,SymbolicValue> :;TODO reset is wrong
+   val zero = UIntValue(0,IntWidth(1))
+   val one = UIntValue(1,IntWidth(1))
+   defn equals? (e1:Expression,e2:UIntValue) -> True|False :;TODO note I don't compare width
+      match(e1) :
+         (e:UIntValue) : 
+            if value(e) == value(e2) : true
+            else : false
+         (e) : false
+   defn one? (e:Expression) -> True|False : 
+      equals?(e,one)
+   defn zero? (e:Expression) -> True|False : ;TODO note I don't compare width
+      equals?(e,zero)
+   defn AND (e1:Expression,e2:Expression) -> Expression :
+      if zero?(e1) : zero
+      else if zero?(e2) : zero
+      else if one?(e1) : e2
+      else if one?(e2) : e1
+      else : DoPrim(BIT-AND-OP,list(e1,e2),list(),UIntType(IntWidth(1)))
+   defn OR (e1:Expression,e2:Expression) -> Expression :
+      if one?(e1) : one
+      else if one?(e2) : one
+      else if zero?(e1) : e2
+      else if zero?(e2) : e1
+      else : DoPrim(BIT-OR-OP,list(e1,e2),list(),UIntType(IntWidth(1)))
+   defn NOT (e1:Expression) -> Expression :
+      DoPrim(EQUAL-UU-OP,list(e1,zero),list(),UIntType(IntWidth(1)))
+   defn reduce-or (l:List<True|False>) -> True|False :
+      if length(l) == 0 : false
+      else : head(l) or reduce-or(tail(l))
+   defn reduce-or (l:List<Expression>) -> Expression :
+      if length(l) == 0 : zero
+      else : OR(head(l) reduce-or(tail(l)))
+      
+   defn get-read-enable (sym:Symbol,sv:SymbolicValue) -> Expression :
+      defn active (e:Expression) -> True|False :
+         match(e) :
+            (e:WRef) : 
+               name(e) == sym
+            (e:WSubfield) : active(exp(e))
+            (e:WRegInit) : active(reg(e))
+            (es:DoPrim) : reduce-or{_} $ for e in args(es) map : active(e)
+            (e:ReadPort) : reduce-or{_} $ map(active,list(mem(e),index(e),enable(e)))
+            (e:WritePort) : reduce-or{_} $ map(active,list(mem(e),index(e),enable(e)))
+            (e:Register) : reduce-or{_} $ map(active,list(value(e),enable(e)))
+            (e) : false
+      val x = match(sv) : 
+         (sv: SVNul) : zero
+         (sv: SVExp) : 
+            if active(exp(sv)) : one
+            else : zero
+         (sv: SVMux) : 
+         ; TODO what if used in predicate?
+            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))
+            OR(e0,OR(AND(pred(sv),e1),AND(NOT(pred(sv)),e2)))
+      ;println-all(["Returning " x " from " sym " with " sv])
+      x
+
+   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 table do :
+      val sym = key(x)
+      val s = stmt(value(x))
+      match(s) :
+         (s:WDefAccessor) : 
+            switch {_ == gender(s)} :
+               FEMALE : enables[sym] = get-write-enable(sv(value(x)))
+               MALE : enables[sym] = SVExp{_} $ reduce-or{_} $ 
+                  for y in table map-append :
+                     list(get-read-enable(sym,sv(value(y))))
+         (s:DefRegister) : 
+            val write-enable = get-write-enable(sv(value(x)))
+            enables[sym] = SVMux(WRef(`reset,UIntType(IntWidth(1)),PortKind(),MALE),reset-enable,write-enable) where :
+               val x = for x in table find : key(x) == symbol-join([name(s),`.init])
+               val reset-enable = match(x) :
+                  (x:False) : SVExp(zero)
+                  (x:KeyValue<Symbol,SSV>) : 
+                     match(sv(value(x))) :
+                        (v:SVNul) : SVExp(zero)
+                        (v) : SVExp(one)
+         (s) : s
+   enables
+            
+defn optimize-table (table:List<KeyValue<Symbol,SSV>>) -> List<KeyValue<Symbol,SSV>>:
    defn optimize (sv:SymbolicValue) -> SymbolicValue :
       match(map(optimize,sv)) :
          (sv:SVMux) : 
-            if conseq(sv) typeof SVNul and alt(sv) typeof SVNul : SVNul()
+            if is-equal(conseq(sv),alt(sv)) : conseq(sv)
             else : sv
          (sv) : sv
    for x in table map :
-      val [k sv] = value(x)
-      key(x) => [k optimize(sv)]
+      key(x) => SSV(stmt(value(x)) optimize(sv(value(x))))
 
-defn build-table (s:Stmt, table:Vector<KeyValue<Symbol,[Stmt SymbolicValue]>>) :
+defn build-table (s:Stmt, table-arg:Vector<KeyValue<Symbol,SSV>>) -> Vector<KeyValue<Symbol,SSV>> : 
+   var table = table-arg
+   println("=====================")
    match(s) :
-      (s:DefWire) : add(table,name(s) => [s SVNul()])
-      (s:DefNode) : add(table,name(s) => [s SVNul()])
-      (s:WDefAccessor) : add(table,name(s) => [s SVNul()]) 
-      (s:DefInstance) : add(table,name(s) => [s SVNul()])
-      (s:DefMemory) : add(table,name(s) => [s SVNul()])
+      (s:DefWire) : add(table,name(s) => SSV(s SVNul()))
+      (s:DefNode) : add(table,name(s) => SSV(s SVNul()))
+      (s:DefRegister) : add(table,name(s) => SSV(s SVNul()))
+      (s:WDefAccessor) : add(table,name(s) => SSV(s SVNul()) )
+      (s:DefInstance) : add(table,name(s) => SSV(s SVNul()))
+      (s:DefMemory) : add(table,name(s) => SSV(s SVNul()))
       (s:Conditionally) :
-         defn in? (i:Symbol,t:Vector<KeyValue<Symbol,[Stmt SymbolicValue]>>) -> True|False :
-            for x in t any? : i == key(x)
-         defn get (i:Symbol,t:Vector<KeyValue<Symbol,[Stmt SymbolicValue]>>) -> [Stmt SymbolicValue] :
+         defn deepcopy (t:Vector<KeyValue<Symbol,SSV>>) -> Vector<KeyValue<Symbol,SSV>> :
+            t0 where :
+               val t0 = Vector<KeyValue<Symbol,SSV>>()
+               for x in t do :
+                  add(t0,key(x) => SSV(stmt(value(x)),sv(value(x))))
+         defn get (t:Vector<KeyValue<Symbol,SSV>>,i:Symbol) -> SSV|False :
             val kv = for x in t find : i == key(x)
             match(kv) :
-               (e:False) : [EmptyStmt() SVNul()]
-               (e:KeyValue<Symbol,[Stmt SymbolicValue]>) : value(e)
-         defn all-kvs (ts:List<Vector<KeyValue<Symbol,[Stmt SymbolicValue]>>>) -> Vector<KeyValue<Symbol,[Stmt SymbolicValue]>> :
-            val t0 = new-vec()
-            for v in ts do : 
-               for t in v do :
-                  add(t0,key(t) => value(t))
-            t0
-
-         val table-c = new-vec()
-         build-table(conseq(s),table-c)
-
-         val table-a = new-vec()
-         build-table(alt(s),table-a)
-
-         val table-m = new-vec()
-         for kv in all-kvs(list(table,table-c,table-a)) do :
-            defn get-sv (i:Symbol,t:Vector<KeyValue<Symbol,[Stmt SymbolicValue]>>) -> SymbolicValue|False :
-               val kv = for x in t find : i == key(x)
-               match(kv) :
-                  (e:KeyValue<Symbol,[Stmt SymbolicValue]>) : 
-                     val [k sv] = value(e)
-                     sv
-                  (e:False) : false
-               
-            val i = key(kv)
-            val [k,sv] = get(i,table)
-            val [k-c,sv-c] = get(i,table-c)
-            val [k-a,sv-a] = get(i,table-a)
-            val [k* sv*] = 
-               match(get-sv(i,table),get-sv(i,table-c),get-sv(i,table-a)) :
-                  (sv?:SymbolicValue,sv-c?:SymbolicValue,sv-a?:SymbolicValue) : [k, SVMux(pred(s),sv-c?,sv-a?)]
-                  (sv?:SymbolicValue,sv-c?:SymbolicValue,sv-a?:False) : [k, SVMux(pred(s),sv-c?,sv)]
-                  (sv?:SymbolicValue,sv-c?:False,sv-a?:SymbolicValue) : [k, SVMux(pred(s),sv,sv-a?)]
-                  (sv?:SymbolicValue,sv-c?:False,sv-a?:False) : [k,sv]
-                  (sv?:False,sv-c?:SymbolicValue,sv-a?:SymbolicValue) : [k, SVMux(pred(s),sv-c?,sv-a?)]
-                  (sv?:False,sv-c?:SymbolicValue,sv-a?:False) : [k-c, SVMux(pred(s),sv-c?,sv)]
-                  (sv?:False,sv-c?:False,sv-a?:SymbolicValue) : [k-a, SVMux(pred(s),sv,sv-a?)]
-                  (sv?:False,sv-c?:False,sv-a?:False) : [k,sv]
-
-            val inserted? = label<True|False> myret :
-               for (index in 0 to false, x in table) do :
-                  if key(x) == i : 
-                     table[index] = i => [k* sv*]
-                     myret(true)
-               myret(false)
-            if not inserted? : add(table,i => [k* sv*])
-               
+               (e:False) : false
+               (e:KeyValue<Symbol,SSV>) : value(e)
+         defn get-unique-keys (ts:List<Vector<KeyValue<Symbol,SSV>>>) -> Vector<Symbol> :
+            t0 where :
+               val t0 = Vector<Symbol>()
+               for v in ts do : 
+                  for t in v do :
+                     val duplicate? = for x in t0 any? : x == key(t)
+                     if not duplicate? : add(t0,key(t))
+
+         println("TABLE")
+         for x in table do : println(x)
+
+         val table1 = table
+         val table2 = deepcopy(table)
+
+         val table-c = build-table(conseq(s),table1)
+         println("TABLE-C")
+         for x in table-c do : println(x)
+
+         val table-a = build-table(alt(s),table2)
+         println("TABLE-A")
+         for x in table-a do : println(x)
+
+         val table-m = new-table()
+         for i in get-unique-keys(list(table-c,table-a)) do :
+            add{table-m,i => _} $ match(get(table-c,i),get(table-a,i)) :
+                  (c:SSV,a:SSV) : SSV(stmt(c),SVMux(pred(s),sv(c),sv(a)))
+                  (c:SSV,a:False) : SSV(stmt(c),SVMux(pred(s),sv(c),SVNul()))
+                  (c:False,a:SSV) : SSV(stmt(a),SVMux(pred(s),SVNul(),sv(a)))
+                  (c:False,a:False) : error("Shouldn't be here")
+         println("TABLE-M")
+         for x in table-m do : println(x)
+         table = table-m
+         println("TABLE")
+         for x in table do : println(x)
       (s:Connect) :
-         var i* = 0
-         var kv* = false
-         for (i in 0 to false, kv in table) do :
-            if name(loc(s) as ?) == key(kv) :
-               i* = i
-               kv* = kv
-         match(kv*) :
-            (kv:False) : add(table,name(loc(s) as ?) => [EmptyStmt() SVExp(exp(s))])
-            (kv:KeyValue<Symbol,[Stmt SymbolicValue]>) : 
-               val [k sv] = value(kv)
-               table[i*] = key(kv) => [k SVExp(exp(s))]
-      (s:Begin) : for s* in body(s) do: build-table(s*,table)
+         val i = for (i in 0 to false, kv in table) find :
+            name(loc(s) as ?) == key(kv)
+         match(i) :
+            (i:False) : add(table,name(loc(s) as ?) => SSV(EmptyStmt() SVExp(exp(s))))
+            (i:Int) : 
+               val kv = table[i]
+               val [k sv] = [stmt(value(kv)) sv(value(kv))]
+               table[i] = key(kv) => SSV(k SVExp(exp(s)))
+      (s:Begin) : for s* in body(s) do: table = build-table(s*,table)
       (s) : s
+   table
          
 defn expand-whens (m:Module) -> Module :
-   var table = new-vec()
-   build-table(body(m),table)
+   val table = build-table(body(m),new-table())
+   println("Original Table")
    for x in table do : println(x)
    val table* = optimize-table(to-list(table))
+   println("Optimized Table")
    for x in table* do : println(x)
+   val enables = get-enables(table*)
+   println("Enable Table")
+   for x in enables do : println(x)
    Module(name(m),ports(m),body(m))
    ;Module(name(m),ports(m),expand-whens-stmt(body(m),table))
 
diff --git a/test/passes/expand-whens/one-when.fir b/test/passes/expand-whens/one-when.fir
index 9745d087..66fc2ef6 100644
--- a/test/passes/expand-whens/one-when.fir
+++ b/test/passes/expand-whens/one-when.fir
@@ -6,7 +6,12 @@ circuit top :
       mem m : UInt(1)[2]
       wire i : UInt(1)
       wire p : UInt(1)
+      wire j : UInt(1)
+      reg r : UInt(1)
+
+      p := j
       when p :
+        r.init := i
         accessor a = m[i] 
         i := a
         accessor b = m[i]
@@ -16,5 +21,12 @@ circuit top :
         i := c
         accessor d = m[i] 
         d := i
+      accessor e = m[i]
+      when p :
+        p := i
+        when e :
+           p := p
+           r.init := p
+        
 
 ; CHECK: Finished Expand Whens
diff --git a/test/passes/lower-to-ground/register.fir b/test/passes/lower-to-ground/register.fir
index 9021d0c2..f270bacb 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#init#x := q#x
-      ; CHECK: q#y := r1#init#y
+      ; CHECK: r1#x.init := q#x
+      ; CHECK: q#y := r1#y.init
 
 ; CHECK: Finished Lower To Ground