[FSM] Add loop unroll test

magma/sum_type.py

@@ -249,14 +249,17 @@ def __new__(mcs, name, bases, namespace):
 
         if tag_len:
             tag_map = {}
+            key_map = {}
             T = AnonProduct[{"tag": Bits[tag_len]}]
             for key, value in namespace.items():
                 if isinstance(value, int):
                     namespace[key] = value = T(Bits[tag_len](value))
                     tag_map[value] = value.tag
+                    key_map[key] = value
             namespace['_magma_T_'] = T
             namespace['_magma_Ts_'] = {}
             namespace['_tag_map'] = tag_map
+            namespace['_key_map'] = key_map
 
         return type.__new__(mcs, name, bases, namespace)
 

magma/syntax/fsm.py

@@ -11,7 +11,9 @@ def reset_fsm_context():
 
 
 class FSMState(Enum2):
-    pass
+    @classmethod
+    def get_states(cls):
+        return cls._key_map.values()
 
 
 class FSMContext(MatchContext):

tests/test_syntax/gold/test_fsm_loop_unroll.mlir

@@ -0,0 +1,53 @@
+module attributes {circt.loweringOptions = "locationInfoStyle=none,omitVersionComment"} {
+    hw.module @Foo(in %CLK: i1, out O: i2) {
+        %1 = hw.struct_create (%0) : !hw.struct<tag: i2>
+        %3 = sv.reg name "Register_inst0" : !hw.inout<!hw.struct<tag: i2>>
+        sv.alwaysff(posedge %CLK) {
+            sv.passign %3, %1 : !hw.struct<tag: i2>
+        }
+        %5 = hw.constant 0 : i2
+        %4 = hw.struct_create (%5) : !hw.struct<tag: i2>
+        sv.initial {
+            sv.bpassign %3, %4 : !hw.struct<tag: i2>
+        }
+        %2 = sv.read_inout %3 : !hw.inout<!hw.struct<tag: i2>>
+        %6 = hw.struct_extract %2["tag"] : !hw.struct<tag: i2>
+        %7 = hw.constant 0 : i2
+        %8 = comb.icmp eq %6, %7 : i2
+        %9 = hw.constant 0 : i2
+        %10 = hw.constant 1 : i2
+        %11 = hw.constant 1 : i2
+        %12 = comb.icmp eq %6, %11 : i2
+        %13 = hw.constant 1 : i2
+        %14 = hw.constant 2 : i2
+        %15 = hw.constant 2 : i2
+        %16 = comb.icmp eq %6, %15 : i2
+        %17 = hw.constant 2 : i2
+        %18 = hw.constant 0 : i2
+        %19 = hw.constant 3 : i2
+        %21 = sv.reg : !hw.inout<i2>
+        %20 = sv.read_inout %21 : !hw.inout<i2>
+        %22 = sv.reg : !hw.inout<i2>
+        %0 = sv.read_inout %22 : !hw.inout<i2>
+        sv.alwayscomb {
+            sv.if %8 {
+                sv.bpassign %21, %9 : i2
+                sv.bpassign %22, %10 : i2
+            } else {
+                sv.if %12 {
+                    sv.bpassign %21, %13 : i2
+                    sv.bpassign %22, %14 : i2
+                } else {
+                    sv.if %16 {
+                        sv.bpassign %21, %17 : i2
+                        sv.bpassign %22, %18 : i2
+                    } else {
+                        sv.bpassign %21, %19 : i2
+                        sv.bpassign %22, %10 : i2
+                    }
+                }
+            }
+        }
+        hw.output %20 : i2
+    }
+}

tests/test_syntax/test_fsm.py

@@ -32,3 +32,23 @@ class Foo(m.Circuit):
 
     m.compile("build/test_fsm_wait", Foo, output="mlir")
     assert check_gold(__file__, "test_fsm_wait.mlir")
+
+
+def test_fsm_loop_unroll():
+    class State(m.FSMState):
+        INIT = 0
+        RUN = 1
+        WAIT = 2
+        DONE = 3
+
+    class Foo(m.Circuit):
+        io = m.IO(O=m.Out(m.Bits[2]))
+        with m.fsm(State, init=State.INIT) as state:
+            states = list(State.get_states())
+            for i, x in enumerate(states):
+                with m.case(x):
+                    io.O @= i
+                    state.next @= states[(i + 1) % 3]
+
+    m.compile("build/test_fsm_loop_unroll", Foo, output="mlir")
+    assert check_gold(__file__, "test_fsm_loop_unroll.mlir")