rewrote all assignments and conditional transitions within fsm to rea…

…d from / write to std_wire primitives

calyx-opt/src/passes/wire_inliner.rs

@@ -23,10 +23,20 @@ impl Named for WireInliner {
 }
 
 fn rewrite(map: &HoleMapping, port: &RRC<ir::Port>) -> Option<RRC<ir::Cell>> {
-    if let ir::PortParent::Group(g) = &port.borrow().parent {
+    let port_cell = &port.borrow().parent;
+    if let ir::PortParent::Group(g) = port_cell {
         let (go, done) = &map[&g.upgrade().borrow().name()];
         let cell = if port.borrow().name == "go" { go } else { done };
         Some(Rc::clone(cell))
+    } else if let ir::PortParent::FSM(f) = port_cell {
+        let bruh = &f.upgrade().borrow().name();
+        let (start, done) = &map[bruh];
+        let cell = if port.borrow().name == "start" {
+            start
+        } else {
+            done
+        };
+        Some(Rc::clone(cell))
     } else {
         None
     }
@@ -52,7 +62,7 @@ fn rewrite_guard(map: &HoleMapping, guard: &mut ir::Guard<Nothing>) {
         // group wire, if it is dependent on a group's port in the first place
         ir::Guard::Port(p) => {
             if let Some(cell) = rewrite(map, p) {
-                guard.update(|_| ir::Guard::port(cell.borrow().get("out")));
+                let _ = std::mem::replace(p, cell.borrow().get("out"));
             }
         }
         // update the ports of a port-comparison guard to read from appropriate
@@ -116,9 +126,10 @@ impl Visitor for WireInliner {
 
         // assume static groups is empty
         let groups = comp.get_groups_mut().drain().collect_vec();
+        let mut fsms = comp.get_fsms_mut().drain().collect_vec();
         let mut builder = ir::Builder::new(comp, sigs);
         // for each group, instantiate wires to hold its `go` and `done` signals.
-        let hole_map: HoleMapping = groups
+        let mut hole_map: HoleMapping = groups
             .iter()
             .map(|gr| {
                 let name = gr.borrow().name();
@@ -136,6 +147,27 @@ impl Visitor for WireInliner {
             })
             .collect();
 
+        // add wires to hold `start` and `done` signals for each FSM, which we
+        // compile differently
+        let fsm_interface = fsms
+            .iter()
+            .map(|fsm| {
+                let name = fsm.borrow().name();
+                let start = builder.add_primitive(
+                    format!("{}_start", name),
+                    "std_wire",
+                    &[1],
+                );
+                let done = builder.add_primitive(
+                    format!("{}_done", name),
+                    "std_wire",
+                    &[1],
+                );
+                (name, (start, done))
+            })
+            .collect_vec();
+        hole_map.extend(fsm_interface);
+
         // Rewrite all assignments first
         groups.iter().for_each(|gr| {
             // Detach assignment from the group first because the rewrite
@@ -150,15 +182,22 @@ impl Visitor for WireInliner {
             gr.borrow_mut().assignments = assigns;
         });
 
-        comp.get_fsms_mut().iter_mut().for_each(|fsm| {
+        // rewrite the contents of each fsm to use instantiated std_wire prims
+        fsms.iter().for_each(|fsm| {
             // rewrite all assignments at each state within every fsm
-            for assigns_at_state in fsm.borrow_mut().assignments.iter_mut() {
+            let mut assigns =
+                fsm.borrow_mut().assignments.drain(..).collect_vec();
+            for assigns_at_state in assigns.iter_mut() {
                 for asgn in assigns_at_state.iter_mut() {
                     rewrite_assign(&hole_map, asgn);
                 }
             }
+            fsm.borrow_mut().assignments = assigns;
+
             // rewrite all guards in transitions that depend on group port values
-            for trans_at_state in fsm.borrow_mut().transitions.iter_mut() {
+            let mut transitions =
+                fsm.borrow_mut().transitions.drain(..).collect_vec();
+            for trans_at_state in transitions.iter_mut() {
                 if let ir::Transition::Conditional(cond_trans_at_state) =
                     trans_at_state
                 {
@@ -167,10 +206,9 @@ impl Visitor for WireInliner {
                     }
                 }
             }
+            fsm.borrow_mut().transitions = transitions;
         });
 
-        // rewrite all transitions as well
-
         comp.continuous_assignments
             .iter_mut()
             .for_each(|assign| rewrite_assign(&hole_map, assign));
@@ -182,6 +220,12 @@ impl Visitor for WireInliner {
 
         comp.continuous_assignments.append(&mut group_assigns);
 
+        // reinstate fsms into the component
+        let comp_fsms_collection = comp.get_fsms_mut();
+        fsms.drain(..).for_each(|fsm| {
+            comp_fsms_collection.add(fsm);
+        });
+
         // remove group from control
         Ok(Action::change(ir::Control::empty()))
     }