Redesign ir::Rewriter interface (#1691)

* Redesign `ir::Rewriter` interface

* update changelog

CHANGELOG.md

@@ -1,5 +1,6 @@
 ## Unreleased
-- Deprecate `Cell::find_with_attr` in favor of `Cell::find_with_unique_attr`. The former is error-prone because pass logic might implicitly assume that there is only one port with a particular attribute.
+- BREAKING: Deprecate `Cell::find_with_attr` in favor of `Cell::find_with_unique_attr`. The former is error-prone because pass logic might implicitly assume that there is only one port with a particular attribute.
+- BREAKING: Redesign the `ir::Rewriter` interface to take all the rewrite maps when constructing the `ir::Rewriter` struct.
 
 ## 0.5.1
 - Change the `calyx` build script to use the `CALYX_PRIMITIVES_DIR` env variable to install primitive libraries. If unset, use `$HOME/.calyx`.

calyx-ir/src/rewriter.rs

@@ -11,22 +11,24 @@ pub type RewriteMap<T> = HashMap<ir::Id, RRC<T>>;
 /// [ir::Port::canonical]) to the new [ir::Port] instance.
 pub type PortRewriteMap = HashMap<ir::Canonical, RRC<ir::Port>>;
 
+#[derive(Default)]
 /// A structure to track rewrite maps for ports. Stores both cell rewrites and direct port
 /// rewrites. Attempts to apply port rewrites first before trying the cell
 /// rewrite.
-pub struct Rewriter<'a> {
-    cell_map: &'a RewriteMap<ir::Cell>,
-    port_map: &'a PortRewriteMap,
+pub struct Rewriter {
+    /// Mapping from canonical names of ports to port instances
+    pub port_map: PortRewriteMap,
+    /// Mapping from names of cells to cell instance.
+    pub cell_map: RewriteMap<ir::Cell>,
+    /// Mapping from names of groups to group instance.
+    pub group_map: RewriteMap<ir::Group>,
+    /// Mapping from names of combinational groups to combinational group instance.
+    pub comb_group_map: RewriteMap<ir::CombGroup>,
+    /// Mapping from names of static groups to static group instance.
+    pub static_group_map: RewriteMap<ir::StaticGroup>,
 }
 
-impl<'a> Rewriter<'a> {
-    pub fn new(
-        cell_map: &'a RewriteMap<ir::Cell>,
-        port_map: &'a PortRewriteMap,
-    ) -> Self {
-        Self { cell_map, port_map }
-    }
-
+impl Rewriter {
     /// Return the rewrite for a cell
     pub fn get_cell_rewrite(&self, cell: &ir::Id) -> Option<RRC<ir::Cell>> {
         self.cell_map.get(cell).map(Rc::clone)
@@ -123,11 +125,7 @@ impl<'a> Rewriter<'a> {
 
     // =========== Control Rewriting Methods =============
     /// Rewrite a `invoke` node using a [RewriteMap<ir::Cell>] and a [RewriteMap<ir::CombGroup>]
-    pub fn rewrite_invoke(
-        &self,
-        inv: &mut ir::Invoke,
-        comb_group_map: &RewriteMap<ir::CombGroup>,
-    ) {
+    pub fn rewrite_invoke(&self, inv: &mut ir::Invoke) {
         // Rewrite the name of the cell
         let name = inv.comp.borrow().name();
         if let Some(new_cell) = &self.get_cell_rewrite(&name) {
@@ -137,7 +135,7 @@ impl<'a> Rewriter<'a> {
         // Rewrite the combinational group
         if let Some(cg_ref) = &inv.comb_group {
             let cg = cg_ref.borrow().name();
-            if let Some(new_cg) = &comb_group_map.get(&cg) {
+            if let Some(new_cg) = &self.comb_group_map.get(&cg) {
                 inv.comb_group = Some(Rc::clone(new_cg));
             }
         }
@@ -174,34 +172,30 @@ impl<'a> Rewriter<'a> {
 
     /// Given a control program, rewrite all uses of cells, groups, and comb groups using the given
     /// rewrite maps.
-    pub fn rewrite_static_control(
-        &self,
-        sc: &mut ir::StaticControl,
-        static_group_map: &RewriteMap<ir::StaticGroup>,
-    ) {
+    pub fn rewrite_static_control(&self, sc: &mut ir::StaticControl) {
         match sc {
             ir::StaticControl::Empty(_) => (),
             ir::StaticControl::Enable(sen) => {
                 let g = &sen.group.borrow().name();
-                if let Some(new_group) = static_group_map.get(g) {
+                if let Some(new_group) = self.static_group_map.get(g) {
                     sen.group = Rc::clone(new_group);
                 }
             }
             ir::StaticControl::Repeat(rep) => {
-                self.rewrite_static_control(&mut rep.body, static_group_map)
+                self.rewrite_static_control(&mut rep.body)
             }
             ir::StaticControl::Seq(ir::StaticSeq { stmts, .. })
             | ir::StaticControl::Par(ir::StaticPar { stmts, .. }) => stmts
                 .iter_mut()
-                .for_each(|c| self.rewrite_static_control(c, static_group_map)),
+                .for_each(|c| self.rewrite_static_control(c)),
             ir::StaticControl::If(sif) => {
                 // Rewrite port use
                 if let Some(new_port) = self.get(&sif.port) {
                     sif.port = new_port;
                 }
                 // rewrite branches
-                self.rewrite_static_control(&mut sif.tbranch, static_group_map);
-                self.rewrite_static_control(&mut sif.fbranch, static_group_map);
+                self.rewrite_static_control(&mut sif.tbranch);
+                self.rewrite_static_control(&mut sif.fbranch);
             }
             ir::StaticControl::Invoke(sin) => {
                 self.rewrite_static_invoke(sin);
@@ -211,31 +205,18 @@ impl<'a> Rewriter<'a> {
 
     /// Given a control program, rewrite all uses of cells, groups, and comb groups using the given
     /// rewrite maps.
-    pub fn rewrite_control(
-        &self,
-        c: &mut ir::Control,
-        group_map: &RewriteMap<ir::Group>,
-        comb_group_map: &RewriteMap<ir::CombGroup>,
-        static_group_map: &RewriteMap<ir::StaticGroup>,
-    ) {
+    pub fn rewrite_control(&self, c: &mut ir::Control) {
         match c {
             ir::Control::Empty(_) => (),
             ir::Control::Enable(en) => {
                 let g = &en.group.borrow().name();
-                if let Some(new_group) = group_map.get(g) {
+                if let Some(new_group) = self.group_map.get(g) {
                     en.group = Rc::clone(new_group);
                 }
             }
             ir::Control::Seq(ir::Seq { stmts, .. })
             | ir::Control::Par(ir::Par { stmts, .. }) => {
-                stmts.iter_mut().for_each(|c| {
-                    self.rewrite_control(
-                        c,
-                        group_map,
-                        comb_group_map,
-                        static_group_map,
-                    )
-                })
+                stmts.iter_mut().for_each(|c| self.rewrite_control(c))
             }
             ir::Control::If(ife) => {
                 // Rewrite port use
@@ -245,23 +226,13 @@ impl<'a> Rewriter<'a> {
                 // Rewrite conditional comb group if defined
                 if let Some(cg_ref) = &ife.cond {
                     let cg = cg_ref.borrow().name();
-                    if let Some(new_cg) = &comb_group_map.get(&cg) {
+                    if let Some(new_cg) = &self.comb_group_map.get(&cg) {
                         ife.cond = Some(Rc::clone(new_cg));
                     }
                 }
                 // rewrite branches
-                self.rewrite_control(
-                    &mut ife.tbranch,
-                    group_map,
-                    comb_group_map,
-                    static_group_map,
-                );
-                self.rewrite_control(
-                    &mut ife.fbranch,
-                    group_map,
-                    comb_group_map,
-                    static_group_map,
-                );
+                self.rewrite_control(&mut ife.tbranch);
+                self.rewrite_control(&mut ife.fbranch);
             }
             ir::Control::While(wh) => {
                 // Rewrite port use
@@ -271,33 +242,19 @@ impl<'a> Rewriter<'a> {
                 // Rewrite conditional comb group if defined
                 if let Some(cg_ref) = &wh.cond {
                     let cg = cg_ref.borrow().name();
-                    if let Some(new_cg) = &comb_group_map.get(&cg) {
+                    if let Some(new_cg) = &self.comb_group_map.get(&cg) {
                         wh.cond = Some(Rc::clone(new_cg));
                     }
                 }
                 // rewrite body
-                self.rewrite_control(
-                    &mut wh.body,
-                    group_map,
-                    comb_group_map,
-                    static_group_map,
-                );
+                self.rewrite_control(&mut wh.body);
             }
             ir::Control::Repeat(rep) => {
                 // rewrite body
-                self.rewrite_control(
-                    &mut rep.body,
-                    group_map,
-                    comb_group_map,
-                    static_group_map,
-                );
-            }
-            ir::Control::Invoke(inv) => {
-                self.rewrite_invoke(inv, comb_group_map)
-            }
-            ir::Control::Static(s) => {
-                self.rewrite_static_control(s, static_group_map)
+                self.rewrite_control(&mut rep.body);
             }
+            ir::Control::Invoke(inv) => self.rewrite_invoke(inv),
+            ir::Control::Static(s) => self.rewrite_static_control(s),
         }
     }
 }

calyx-opt/src/passes/cell_share.rs

@@ -518,8 +518,10 @@ impl Visitor for CellShare {
         }
 
         // Rewrite assignments using the coloring generated.
-        let empty_map: ir::rewriter::PortRewriteMap = HashMap::new();
-        let rewriter = ir::Rewriter::new(&coloring, &empty_map);
+        let rewriter = ir::Rewriter {
+            cell_map: coloring,
+            ..Default::default()
+        };
         comp.for_each_assignment(|assign| {
             assign.for_each_port(|port| rewriter.get(port));
         });
@@ -528,12 +530,7 @@ impl Visitor for CellShare {
         });
 
         // Rewrite control uses of ports
-        rewriter.rewrite_control(
-            &mut comp.control.borrow_mut(),
-            &HashMap::new(),
-            &HashMap::new(),
-            &HashMap::new(),
-        );
+        rewriter.rewrite_control(&mut comp.control.borrow_mut());
 
         Ok(Action::Stop)
     }

calyx-opt/src/passes/comb_prop.rs

@@ -1,7 +1,6 @@
 use crate::traversal::{Action, ConstructVisitor, Named, VisResult, Visitor};
 use calyx_ir::{self as ir, RRC};
 use itertools::Itertools;
-use std::collections::HashMap;
 use std::rc::Rc;
 
 /// A data structure to track rewrites of ports with added functionality to declare
@@ -369,14 +368,11 @@ impl Visitor for CombProp {
             }
         });
 
-        let cell_rewrites = HashMap::new();
-        let rewriter = ir::Rewriter::new(&cell_rewrites, &rewrites);
-        rewriter.rewrite_control(
-            &mut comp.control.borrow_mut(),
-            &HashMap::new(),
-            &HashMap::new(),
-            &HashMap::new(),
-        );
+        let rewriter = ir::Rewriter {
+            port_map: rewrites,
+            ..Default::default()
+        };
+        rewriter.rewrite_control(&mut comp.control.borrow_mut());
 
         Ok(Action::Stop)
     }

calyx-opt/src/passes/compile_ref.rs

@@ -164,13 +164,11 @@ impl Visitor for CompileRef {
         _comps: &[ir::Component],
     ) -> VisResult {
         log::debug!("compile-ref: {}", comp.name);
-        self.ref_cells = dump_ports::dump_ports_to_signature(
-            comp,
-            is_external_cell,
-            true,
-            &mut self.port_names,
-            &mut self.removed,
-        );
+        let dump_ports::DumpResults { cells, rewrites } =
+            dump_ports::dump_ports_to_signature(comp, is_external_cell, true);
+        self.removed.extend(rewrites.clone());
+        self.port_names.insert(comp.name, rewrites);
+        self.ref_cells = cells;
 
         // For all subcomponents that had a `ref` cell in them, we need to
         // update their cell to have the new ports added from inlining the
@@ -237,21 +235,19 @@ impl Visitor for CompileRef {
         _sigs: &LibrarySignatures,
         _comps: &[ir::Component],
     ) -> VisResult {
+        let port_map = std::mem::take(&mut self.removed);
         // Rewrite all of the ref cell ports
-        let empty = HashMap::new();
-        let rw = ir::Rewriter::new(&empty, &self.removed);
+        let rw = ir::Rewriter {
+            port_map,
+            ..Default::default()
+        };
         comp.for_each_assignment(|assign| {
             rw.rewrite_assign(assign);
         });
         comp.for_each_static_assignment(|assign| {
             rw.rewrite_assign(assign);
         });
-        rw.rewrite_control(
-            &mut comp.control.borrow_mut(),
-            &HashMap::new(),
-            &HashMap::new(),
-            &HashMap::new(),
-        );
+        rw.rewrite_control(&mut comp.control.borrow_mut());
         Ok(Action::Continue)
     }
 }

calyx-opt/src/passes/component_iniliner.rs

@@ -290,21 +290,25 @@ impl ComponentInliner {
 
         // Rewrites to inline the interface.
         let interface_map = Self::inline_interface(builder, comp, name);
-        let rewrite = ir::Rewriter::new(&cell_map, &interface_map);
+        let mut rewrite = ir::Rewriter {
+            cell_map,
+            port_map: interface_map,
+            ..Default::default()
+        };
 
         // For each group, create a new group and rewrite all assignments within
         // it using the `rewrite_map`.
-        let group_map: rewriter::RewriteMap<ir::Group> = comp
+        rewrite.group_map = comp
             .get_groups()
             .iter()
             .map(|gr| Self::inline_group(builder, &rewrite, gr))
             .collect();
-        let static_group_map: rewriter::RewriteMap<ir::StaticGroup> = comp
+        rewrite.static_group_map = comp
             .get_static_groups()
             .iter()
             .map(|gr| Self::inline_static_group(builder, &rewrite, gr))
             .collect();
-        let comb_group_map: rewriter::RewriteMap<ir::CombGroup> = comp
+        rewrite.comb_group_map = comp
             .comb_groups
             .iter()
             .map(|gr| Self::inline_comb_group(builder, &rewrite, gr))
@@ -320,17 +324,12 @@ impl ComponentInliner {
 
         // Generate a control program associated with this instance
         let mut con = ir::Cloner::control(&comp.control.borrow());
-        rewrite.rewrite_control(
-            &mut con,
-            &group_map,
-            &comb_group_map,
-            &static_group_map,
-        );
+        rewrite.rewrite_control(&mut con);
 
         // Generate interface map for use in the parent cell.
         // Return as an iterator because it's immediately merged into the global rewrite map.
         let rev_interface_map =
-            interface_map.into_iter().map(move |(cp, pr)| {
+            rewrite.port_map.into_iter().map(move |(cp, pr)| {
                 let ir::Canonical(_, p) = cp;
                 let port = pr.borrow();
                 let np = match port.name.id.as_str() {