From 04b109e39315e99af84e86dfd21016050dcad3fb Mon Sep 17 00:00:00 2001
From: Lenny Truong <leonard@truong.io>
Date: Thu, 21 Mar 2024 13:49:24 -0700
Subject: [PATCH] [Syntax] Init FSM

---
 magma/__init__.py             |  1 +
 magma/syntax/fsm.py           | 13 +++++++++++++
 tests/test_syntax/test_fsm.py | 31 +++++++++++++++++++++++++++++++
 3 files changed, 45 insertions(+)
 create mode 100644 magma/syntax/fsm.py
 create mode 100644 tests/test_syntax/test_fsm.py

diff --git a/magma/__init__.py b/magma/__init__.py
index 3da938d5b..8e7bf3770 100644
--- a/magma/__init__.py
+++ b/magma/__init__.py
@@ -147,6 +147,7 @@ def set_mantle_target(t):
 from magma.debug import magma_helper_function
 import magma.mantle
 from magma.sum_type import Sum, match, case, Enum2
+from magma.syntax.fsm import fsm, FSMState, wait
 
 ################################################################################
 # BEGIN ALIASES
diff --git a/magma/syntax/fsm.py b/magma/syntax/fsm.py
new file mode 100644
index 000000000..ad8bef1e2
--- /dev/null
+++ b/magma/syntax/fsm.py
@@ -0,0 +1,13 @@
+import magma as m
+
+
+class FSMState(m.Enum2):
+    pass
+
+
+def fsm(T: m.Kind):
+    pass
+
+
+def wait():
+    pass
diff --git a/tests/test_syntax/test_fsm.py b/tests/test_syntax/test_fsm.py
new file mode 100644
index 000000000..08d6514d9
--- /dev/null
+++ b/tests/test_syntax/test_fsm.py
@@ -0,0 +1,31 @@
+import magma as m
+from magma.testing.utils import check_gold
+
+
+def test_fsm_wait():
+    class State(m.FSMState):
+        INIT = 0
+        RUN = 1
+        DONE = 2
+
+    class Foo(m.Circuit):
+        io = m.IO(O=m.Out(m.Bits[16]))
+        # Must delay elaboration of state register until total number states
+        # computed (number of wait statements)
+        # Could use a builder here?
+        with m.fsm(State) as state:
+            with m.case(State.INIT):
+                io.O @= 0xFEED
+                state.next @= State.RUN
+            with m.case(State.RUN):
+                io.O @= 0xDEAD
+                m.wait()
+                io.O @= 0xBEEF
+                state.next @= State.DONE
+            with m.case(State.DONE):
+                io.O @= 0xBEEF
+                state.next @= State.DONE
+
+    m.compile("build/test_sum_enum", Foo, output="mlir",
+              flatten_all_tuples=True)
+    assert check_gold(__file__, "test_sum_enum.mlir")