Now when expanding ConnectFrom/ToIndex, create a node for the index so it isn't duplicated for all the whens

3 files changed, 34 insertions, 21 deletions

diff --git a/src/main/stanza/passes.stanza b/src/main/stanza/passes.stanza
index b1410fa4..1b752c27 100644
--- a/src/main/stanza/passes.stanza
+++ b/src/main/stanza/passes.stanza
@@ -862,24 +862,38 @@ defn lower-to-ground (c:Circuit) -> Circuit :
 defn expand-connect-indexed-stmt (s: Stmt) -> Stmt :
    defn equality (e1:Expression,e2:Expression) -> Expression :
       DoPrim(EQUAL-UU-OP,list(e1,e2),List(),UIntType(UnknownWidth()))
+   defn get-name (e:Expression) -> Symbol :
+      match(e) :
+         (e:WRef) : symbol-join([name(e) `#])
+         (e:WSubfield) : symbol-join([get-name(exp(e)) `. name(e) `#])
+         (e:WIndex) : symbol-join([get-name(exp(e)) `. to-symbol(value(e)) `#])
+         (e) : `T
    match(s) :
       (s:ConnectToIndexed) : Begin $
          if length(locs(s)) == 0 : list(EmptyStmt())
          else :
-            List(Connect(head(locs(s)),exp(s)), to-list $
-               for (i in 1 to false, l in tail(locs(s))) stream : Conditionally(
-                  equality(index(s),UIntValue(i,UnknownWidth())),
-                  Connect(l,exp(s)),
-                  EmptyStmt())
+            val ref = WRef(firrtl-gensym(get-name(exp(s))),type(index(s)),NodeKind(),UNKNOWN-GENDER)
+            append(
+               list(Connect(head(locs(s)),exp(s)),DefNode(name(ref),index(s)))
+               to-list $
+                  for (i in 1 to false, l in tail(locs(s))) stream : Conditionally(
+                     equality(ref,UIntValue(i,UnknownWidth())),
+                     Connect(l,exp(s)),
+                     EmptyStmt()
+                  )
             )
       (s:ConnectFromIndexed) : Begin $
          if length(exps(s)) == 0 : list(EmptyStmt())
          else :
-            List(Connect(loc(s),head(exps(s))), to-list $
-               for (i in 1 to false, e in tail(exps(s))) stream : Conditionally(
-                  equality(index(s),UIntValue(i,UnknownWidth())),
-                  Connect(loc(s),e),
-                  EmptyStmt())
+            val ref = WRef(firrtl-gensym(get-name(loc(s))),type(index(s)),NodeKind(),UNKNOWN-GENDER)
+            append(
+               list(Connect(loc(s),head(exps(s))),DefNode(name(ref),index(s)))
+               to-list $
+                  for (i in 1 to false, e in tail(exps(s))) stream : Conditionally(
+                     equality(ref,UIntValue(i,UnknownWidth())),
+                     Connect(loc(s),e),
+                     EmptyStmt()
+                  )
             )
       (s) : map(expand-connect-indexed-stmt,s)
   
diff --git a/test/chisel3/Mul.fir b/test/chisel3/Mul.fir
index ec991197..46a8c9b9 100644
--- a/test/chisel3/Mul.fir
+++ b/test/chisel3/Mul.fir
@@ -1,12 +1,10 @@
 ; RUN: firrtl -i %s -o %s.flo -x X -p c | tee %s.out | FileCheck %s
 ; CHECK: Done!
-
 circuit Mul : 
   module Mul : 
     input x : UInt<2>
-    output z : UInt<4>
-    output a : UInt<4>
     input y : UInt<2>
+    output z : UInt<4>
     
     wire tbl : UInt<4>[16]
     tbl[0] := UInt<4>(0)
@@ -25,8 +23,7 @@ circuit Mul :
     tbl[13] := UInt<4>(3)
     tbl[14] := UInt<4>(6)
     tbl[15] := UInt<4>(9)
-    node T_43 = shl(x, 2)
-    node ad = bit-or(Pad(T_43,?), Pad(y,?))
-    a := Pad(ad,?)
-    accessor T_44 = tbl[ad]
+    node T_42 = shl(x, 2)
+    node T_43 = bit-or(Pad(T_42,?), Pad(y,?))
+    accessor T_44 = tbl[T_43]
     z := Pad(T_44,?)
diff --git a/test/passes/expand-connect-indexed/bundle-vecs.fir b/test/passes/expand-connect-indexed/bundle-vecs.fir
index 35b45b6c..c8235ade 100644
--- a/test/passes/expand-connect-indexed/bundle-vecs.fir
+++ b/test/passes/expand-connect-indexed/bundle-vecs.fir
@@ -16,11 +16,13 @@ circuit top :
       ; CHECK: wire b$x : UInt<32>
       ; CHECK: wire b$y : UInt<32>
       ; CHECK: b$x := a$0$x
-      ; CHECK: when eq-uu(i, UInt(1)) :
-      ; CHECK:    b$x := a$1$x
+      ; CHECK: node b$x#0 = i
+      ; CHECK: when eq-uu(b$x#0, UInt(1)) :
+      ; CHECK:    b$x := a$1$x      
       ; CHECK: a$0$y := b$y
-      ; CHECK: when eq-uu(i, UInt(1)) :
-      ; CHECK:    a$1$y := b$y
+      ; CHECK: node b$y#0 = i
+      ; CHECK: when eq-uu(b$y#0, UInt(1)) :
+      ; CHECK:    a$1$y := b$y      
       j := b
       
 ; CHECK: Finished Expand Indexed Connects