From 094df27eb8053e8af156a7042c82059d7e76a7e4 Mon Sep 17 00:00:00 2001
From: Parth Sarkar <93054509+parthsarkar17@users.noreply.github.com>
Date: Wed, 18 Dec 2024 15:17:40 -0500
Subject: [PATCH] Regroup assignments outside fsm (#2380)

---
 calyx-backend/src/verilog.rs | 45 ++++--------------------------------
 calyx-ir/src/structure.rs    | 31 ++++++++++---------------
 2 files changed, 16 insertions(+), 60 deletions(-)

diff --git a/calyx-backend/src/verilog.rs b/calyx-backend/src/verilog.rs
index d21ad4af6..6c4bf157e 100644
--- a/calyx-backend/src/verilog.rs
+++ b/calyx-backend/src/verilog.rs
@@ -467,12 +467,8 @@ fn emit_fsm<F: io::Write>(fsm: &RRC<ir::FSM>, f: &mut F) -> io::Result<()> {
     emit_fsm_inlined_reg(&state_reg, &state_next, reg_bitwidth, f)?;
 
     // dump assignments to enable in this state
-    emit_fsm_dependent_assignments(
-        fsm.borrow().merge_assignments(),
-        &state_reg,
-        reg_bitwidth,
-        f,
-    )?;
+    emit_fsm_dependent_assignments(fsm, &state_reg, reg_bitwidth, f)?;
+
 
     // emit fsm in case statement form
     writeln!(f, "always @(*) begin")?;
@@ -493,46 +489,13 @@ fn emit_fsm<F: io::Write>(fsm: &RRC<ir::FSM>, f: &mut F) -> io::Result<()> {
     io::Result::Ok(())
 }
 
-fn emit_fsm_assignments<F: io::Write>(
-    assignments: &Vec<ir::Assignment<Nothing>>,
-    f: &mut F,
-) -> io::Result<()> {
-    // dump assignments to enable in this state
-    for assign in assignments.into_iter() {
-        let dst_ref = &assign.dst;
-        let guard_unmet_value = if is_data_port(dst_ref) {
-            format!("'x")
-        } else {
-            format!("{}'d0", dst_ref.borrow().width)
-        };
-        let destination = if assign.guard.is_true() {
-            format!("{}", VerilogPortRef(&assign.src))
-        } else {
-            format!(
-                "{} ? {} : {}",
-                unflattened_guard(&assign.guard),
-                VerilogPortRef(&assign.src),
-                guard_unmet_value
-            )
-        };
-        writeln!(
-            f,
-            "{}{} = {};",
-            " ".repeat(6),
-            VerilogPortRef(dst_ref),
-            destination
-        )?;
-    }
-    io::Result::Ok(())
-}
-
 fn emit_fsm_dependent_assignments<F: io::Write>(
-    assignments: Vec<Vec<(usize, ir::Assignment<Nothing>)>>,
+    fsm: &RRC<ir::FSM>,
     fsm_out: &String,
     reg_bitwidth: u64,
     f: &mut F,
 ) -> io::Result<()> {
-    for collection in assignments.iter() {
+    for collection in fsm.borrow().merge_assignments().iter() {
         let dst_ref = &collection.first().unwrap().1.dst;
         writeln!(f, "assign {} =", VerilogPortRef(dst_ref))?;
         for (i, (case, assign)) in collection.iter().enumerate() {
diff --git a/calyx-ir/src/structure.rs b/calyx-ir/src/structure.rs
index 0d6a57862..7c264147d 100644
--- a/calyx-ir/src/structure.rs
+++ b/calyx-ir/src/structure.rs
@@ -896,33 +896,26 @@ impl FSM {
     }
 
     pub fn merge_assignments(&self) -> Vec<Vec<(usize, Assignment<Nothing>)>> {
-        let mut gathered_assigns: HashMap<
-            String,
+        let mut assigns_by_port: HashMap<
+            Canonical,
             Vec<(usize, Assignment<Nothing>)>,
         > = HashMap::new();
         for (case, assigns_at_state) in self.assignments.iter().enumerate() {
             for assign in assigns_at_state.iter() {
-                let port = assign.dst.borrow();
-                let dest = match &port.parent {
-                    PortParent::Cell(cell) => {
-                        format!(
-                            "{}_{}",
-                            cell.upgrade().borrow().name,
-                            port.name
-                        )
-                    }
-                    _ => unreachable!(),
-                };
-
-                gathered_assigns
-                    .entry(dest)
-                    .and_modify(|gathered| {
-                        gathered.push((case, assign.clone()));
+                let dest_port = assign.dst.borrow().canonical();
+                assigns_by_port
+                    .entry(dest_port)
+                    .and_modify(|assigns_at_port| {
+                        assigns_at_port.push((case, assign.clone()));
                     })
                     .or_insert(vec![(case, assign.clone())]);
             }
         }
-        gathered_assigns.drain().map(|(_, v)| v).collect()
+        assigns_by_port
+            .into_values()
+            // order by state, for better appearance when emitted
+            .sorted_by(|a, b| a.first().unwrap().0.cmp(&b.first().unwrap().0))
+            .collect()
     }
 }