From 499d1e9fb0654165149f33ebbb63e8453a42312c Mon Sep 17 00:00:00 2001
From: Jiahong Bi <35024050+HahaLan97@users.noreply.github.com>
Date: Mon, 8 Jul 2024 09:29:16 +0200
Subject: [PATCH] [FSM]New builders for StateOp and TransitionOp. (#6991)

* [FSM]New Builders for StateOp and TransitionOp

1. Create an OutputOp inside a StateOp with the output values instead of an empty OutputOp, which needs to be erased after all.

2. Accept two functions as arguments to create the TransitionOp, which is similar to the creation of sv::IfOp. This would be helpful because one doesn't have to create the blocks and set the insertion point manually, which may cause errors sometime.
---
 include/circt/Dialect/FSM/FSMOps.td | 10 +++++---
 lib/Dialect/FSM/FSMOps.cpp          | 39 +++++++++++++++++++++++------
 2 files changed, 39 insertions(+), 10 deletions(-)

diff --git a/include/circt/Dialect/FSM/FSMOps.td b/include/circt/Dialect/FSM/FSMOps.td
index d1dd516cd6..6d52ac4c2c 100644
--- a/include/circt/Dialect/FSM/FSMOps.td
+++ b/include/circt/Dialect/FSM/FSMOps.td
@@ -222,7 +222,8 @@ def StateOp : FSMOp<"state", [HasParent<"MachineOp">, Symbol, NoTerminator]> {
   }];
 
   let builders = [
-    OpBuilder<(ins "StringRef":$stateName)>
+    OpBuilder<(ins "StringRef":$stateName)>,
+    OpBuilder<(ins "StringRef":$stateName, "ValueRange":$outputs)>
   ];
 
   let skipDefaultBuilders = 1;
@@ -301,8 +302,11 @@ def TransitionOp : FSMOp<"transition", [HasParent<"StateOp">, NoTerminator]> {
   }];
 
   let builders = [
-    OpBuilder<(ins "StringRef":$nextState)>,
-    OpBuilder<(ins "fsm::StateOp":$nextState)>
+    // OpBuilder<(ins "StringRef":$nextState)>,
+    OpBuilder<(ins "fsm::StateOp":$nextState)>,
+    OpBuilder<(ins "StringRef":$nextState,
+                   CArg<"llvm::function_ref<void()>", "{}">:$guardCtor,
+                   CArg<"llvm::function_ref<void()>", "{}">:$actionCtor)>
   ];
 
   let skipDefaultBuilders = 1;
diff --git a/lib/Dialect/FSM/FSMOps.cpp b/lib/Dialect/FSM/FSMOps.cpp
index 0cb42ce4ae..0045d48953 100644
--- a/lib/Dialect/FSM/FSMOps.cpp
+++ b/lib/Dialect/FSM/FSMOps.cpp
@@ -321,6 +321,18 @@ void StateOp::build(OpBuilder &builder, OperationState &state,
   state.addAttribute("sym_name", builder.getStringAttr(stateName));
 }
 
+void StateOp::build(OpBuilder &builder, OperationState &state,
+                    StringRef stateName, ValueRange outputs) {
+  OpBuilder::InsertionGuard guard(builder);
+  Region *output = state.addRegion();
+  output->push_back(new Block());
+  builder.setInsertionPointToEnd(&output->back());
+  builder.create<fsm::OutputOp>(state.location, outputs);
+  Region *transitions = state.addRegion();
+  transitions->push_back(new Block());
+  state.addAttribute("sym_name", builder.getStringAttr(stateName));
+}
+
 SetVector<StateOp> StateOp::getNextStates() {
   SmallVector<StateOp> nextStates;
   llvm::transform(
@@ -402,16 +414,29 @@ LogicalResult OutputOp::verify() {
 //===----------------------------------------------------------------------===//
 
 void TransitionOp::build(OpBuilder &builder, OperationState &state,
-                         StringRef nextState) {
-  state.addRegion(); // guard
-  state.addRegion(); // action
-  state.addAttribute("nextState",
-                     FlatSymbolRefAttr::get(builder.getStringAttr(nextState)));
+                         StateOp nextState) {
+  build(builder, state, nextState.getName());
 }
 
 void TransitionOp::build(OpBuilder &builder, OperationState &state,
-                         StateOp nextState) {
-  build(builder, state, nextState.getName());
+                         StringRef nextState,
+                         llvm::function_ref<void()> guardCtor,
+                         llvm::function_ref<void()> actionCtor) {
+  state.addAttribute("nextState",
+                     FlatSymbolRefAttr::get(builder.getStringAttr(nextState)));
+  OpBuilder::InsertionGuard guard(builder);
+
+  Region *guardRegion = state.addRegion(); // guard
+  if (guardCtor) {
+    builder.createBlock(guardRegion);
+    guardCtor();
+  }
+
+  Region *actionRegion = state.addRegion(); // action
+  if (actionCtor) {
+    builder.createBlock(actionRegion);
+    actionCtor();
+  }
 }
 
 Block *TransitionOp::ensureGuard(OpBuilder &builder) {