calyx_ir/
rewriter.rs

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

calyx_ir/
rewriter.rs
