calyx_ir/
rewriter.rs

1use crate::control::StaticInvoke;
2use crate::{self as ir, RRC};
3use std::borrow::BorrowMut;
4use std::cell::RefCell;
5use std::collections::HashMap;
6use std::rc::Rc;
7
8/// A rewrite map from [ir::Id] to [T].
9pub type RewriteMap<T> = HashMap<ir::Id, RRC<T>>;
10
11/// Map to rewrite port uses. Maps the canonical name of an old port (generated using
12/// [ir::Port::canonical]) to the new [ir::Port] instance.
13pub type PortRewriteMap = HashMap<ir::Canonical, RRC<ir::Port>>;
14
15#[derive(Default)]
16/// A structure to track rewrite maps for ports. Stores both cell rewrites and direct port
17/// rewrites. Attempts to apply port rewrites first before trying the cell
18/// rewrite.
19pub struct Rewriter {
20    /// Mapping from canonical names of ports to port instances
21    pub port_map: PortRewriteMap,
22    /// Mapping from names of cells to cell instance.
23    pub cell_map: RewriteMap<ir::Cell>,
24    /// Mapping from names of groups to group instance.
25    pub group_map: RewriteMap<ir::Group>,
26    /// Mapping from names of combinational groups to combinational group instance.
27    pub comb_group_map: RewriteMap<ir::CombGroup>,
28    /// Mapping from names of static groups to static group instance.
29    pub static_group_map: RewriteMap<ir::StaticGroup>,
30}
31
32impl Rewriter {
33    /// Return the rewrite for a cell
34    pub fn get_cell_rewrite(&self, cell: &ir::Id) -> Option<RRC<ir::Cell>> {
35        self.cell_map.get(cell).map(Rc::clone)
36    }
37
38    /// Return a cell rewrite for the given port. A cell rewrite will attempt
39    /// to give the port with the same name on the new cell.
40    ///
41    /// For example, given with `cell = a` and `new_cell = b`
42    /// ```
43    /// a.in = a.done ? a.out;
44    /// ```
45    /// is rewritten to
46    /// ```
47    /// b.in = b.done ? b.out;
48    /// ```
49    #[inline]
50    fn get_cell_port_rewrite(
51        &self,
52        port_ref: &RRC<ir::Port>,
53    ) -> Option<RRC<ir::Port>> {
54        if self.cell_map.is_empty() {
55            return None;
56        }
57
58        let port = port_ref.borrow();
59        let new_cell = if let ir::PortParent::Cell(cell_wref) = &port.parent {
60            let cell_ref = cell_wref.upgrade();
61            let cell = cell_ref.borrow();
62            self.cell_map.get(&cell.name())
63        } else {
64            None
65        };
66        // Return port with the same name on the new_cell.
67        new_cell.map(|new_cell| Rc::clone(&new_cell.borrow().get(&port.name)))
68    }
69
70    /// Return a port rewrite if present.
71    #[inline]
72    fn get_port_rewrite(
73        &self,
74        port_ref: &RRC<ir::Port>,
75    ) -> Option<RRC<ir::Port>> {
76        if self.port_map.is_empty() {
77            return None;
78        }
79
80        let port = port_ref.borrow();
81        self.port_map.get(&port.canonical()).map(Rc::clone)
82    }
83
84    /// Get any port rewrite defined for the given port.
85    #[inline]
86    pub fn get(&self, port_ref: &RRC<ir::Port>) -> Option<RRC<ir::Port>> {
87        self.get_port_rewrite(port_ref)
88            .or_else(|| self.get_cell_port_rewrite(port_ref))
89    }
90
91    /// Rewrite assignments in a guard
92    pub fn rewrite_guard<T>(&self, guard: &mut ir::Guard<T>) {
93        match guard {
94            ir::Guard::And(l, r) | ir::Guard::Or(l, r) => {
95                self.rewrite_guard(l.borrow_mut());
96                self.rewrite_guard(r.borrow_mut())
97            }
98            ir::Guard::Not(g) => self.rewrite_guard(g.borrow_mut()),
99            ir::Guard::CompOp(_, l, r) => {
100                if let Some(nl) = self.get(l) {
101                    *l = nl;
102                }
103                if let Some(nr) = self.get(r) {
104                    *r = nr;
105                }
106            }
107            ir::Guard::Port(p) => {
108                if let Some(np) = self.get(p) {
109                    *p = np;
110                }
111            }
112            ir::Guard::Info(_) | ir::Guard::True => (),
113        }
114    }
115
116    /// Rewrite an assignment
117    pub fn rewrite_assign<T>(&self, assign: &mut ir::Assignment<T>) {
118        if let Some(dst) = self.get(&assign.dst) {
119            assign.dst = dst;
120        }
121        if let Some(src) = self.get(&assign.src) {
122            assign.src = src;
123        }
124        self.rewrite_guard(&mut assign.guard);
125    }
126
127    // =========== Control Rewriting Methods =============
128    /// Rewrite a `invoke` node using a [RewriteMap<ir::Cell>] and a [RewriteMap<ir::CombGroup>]
129    pub fn rewrite_invoke(&self, inv: &mut ir::Invoke) {
130        // Rewrite the name of the cell
131        let name = inv.comp.borrow().name();
132        if let Some(new_cell) = &self.get_cell_rewrite(&name) {
133            inv.comp = Rc::clone(new_cell);
134        }
135
136        // Rewrite the combinational group
137        if let Some(cg_ref) = &inv.comb_group {
138            let cg = cg_ref.borrow().name();
139            if let Some(new_cg) = &self.comb_group_map.get(&cg) {
140                inv.comb_group = Some(Rc::clone(new_cg));
141            }
142        }
143
144        // Rewrite the parameters
145        inv.inputs
146            .iter_mut()
147            .chain(inv.outputs.iter_mut())
148            .for_each(|(_, port)| {
149                if let Some(new_port) = self.get(&*port) {
150                    *port = new_port;
151                }
152            });
153    }
154
155    /// Rewrite a `static invoke` node using a [RewriteMap<ir::Cell>] and a [RewriteMap<ir::CombGroup>]
156    pub fn rewrite_static_invoke(&self, inv: &mut StaticInvoke) {
157        // Rewrite the name of the cell
158        let name = inv.comp.borrow().name();
159        if let Some(new_cell) = &self.get_cell_rewrite(&name) {
160            inv.comp = Rc::clone(new_cell);
161        }
162
163        // Rewrite the parameters
164        inv.inputs
165            .iter_mut()
166            .chain(inv.outputs.iter_mut())
167            .for_each(|(_, port)| {
168                if let Some(new_port) = self.get(&*port) {
169                    *port = new_port;
170                }
171            });
172    }
173
174    /// Given a control program, rewrite all uses of cells, groups, and comb groups using the given
175    /// rewrite maps.
176    pub fn rewrite_static_control(&self, sc: &mut ir::StaticControl) {
177        match sc {
178            ir::StaticControl::Empty(_) => (),
179            ir::StaticControl::Enable(sen) => {
180                let g = &sen.group.borrow().name();
181                if let Some(new_group) = self.static_group_map.get(g) {
182                    sen.group = Rc::clone(new_group);
183                }
184            }
185            ir::StaticControl::Repeat(rep) => {
186                self.rewrite_static_control(&mut rep.body)
187            }
188            ir::StaticControl::Seq(ir::StaticSeq { stmts, .. })
189            | ir::StaticControl::Par(ir::StaticPar { stmts, .. }) => stmts
190                .iter_mut()
191                .for_each(|c| self.rewrite_static_control(c)),
192            ir::StaticControl::If(sif) => {
193                // Rewrite port use
194                if let Some(new_port) = self.get(&sif.port) {
195                    sif.port = new_port;
196                }
197                // rewrite branches
198                self.rewrite_static_control(&mut sif.tbranch);
199                self.rewrite_static_control(&mut sif.fbranch);
200            }
201            ir::StaticControl::Invoke(sin) => {
202                self.rewrite_static_invoke(sin);
203            }
204        }
205    }
206
207    /// Given a control program, rewrite all uses of cells, groups, and comb groups using the given
208    /// rewrite maps.
209    pub fn rewrite_control(&self, c: &mut ir::Control) {
210        match c {
211            ir::Control::Empty(_) => (),
212            ir::Control::Enable(en) => {
213                let g = &en.group.borrow().name();
214                if let Some(new_group) = self.group_map.get(g) {
215                    en.group = Rc::clone(new_group);
216                }
217            }
218            ir::Control::Seq(ir::Seq { stmts, .. })
219            | ir::Control::Par(ir::Par { stmts, .. }) => {
220                stmts.iter_mut().for_each(|c| self.rewrite_control(c))
221            }
222            ir::Control::If(ife) => {
223                // Rewrite port use
224                if let Some(new_port) = self.get(&ife.port) {
225                    ife.port = new_port;
226                }
227                // Rewrite conditional comb group if defined
228                if let Some(cg_ref) = &ife.cond {
229                    let cg = cg_ref.borrow().name();
230                    if let Some(new_cg) = &self.comb_group_map.get(&cg) {
231                        ife.cond = Some(Rc::clone(new_cg));
232                    }
233                }
234                // rewrite branches
235                self.rewrite_control(&mut ife.tbranch);
236                self.rewrite_control(&mut ife.fbranch);
237            }
238            ir::Control::While(wh) => {
239                // Rewrite port use
240                if let Some(new_port) = self.get(&wh.port) {
241                    wh.port = new_port;
242                }
243                // Rewrite conditional comb group if defined
244                if let Some(cg_ref) = &wh.cond {
245                    let cg = cg_ref.borrow().name();
246                    if let Some(new_cg) = &self.comb_group_map.get(&cg) {
247                        wh.cond = Some(Rc::clone(new_cg));
248                    }
249                }
250                // rewrite body
251                self.rewrite_control(&mut wh.body);
252            }
253            ir::Control::Repeat(rep) => {
254                // rewrite body
255                self.rewrite_control(&mut rep.body);
256            }
257            ir::Control::Invoke(inv) => self.rewrite_invoke(inv),
258            ir::Control::Static(s) => self.rewrite_static_control(s),
259        }
260    }
261
262    /// Rewrite the component using the given maps
263    pub fn rewrite(&self, comp: &mut ir::Component) {
264        // Rewrite all of the ref cell ports
265        comp.for_each_assignment(|assign| {
266            self.rewrite_assign(assign);
267        });
268        comp.for_each_static_assignment(|assign| {
269            self.rewrite_assign(assign);
270        });
271        self.rewrite_control(&mut RefCell::borrow_mut(
272            comp.control.borrow_mut(),
273        ));
274    }
275}
calyx_ir/rewriter.rs

calyx_ir/
rewriter.rs
