calyx_opt/analysis/
control_id.rs

1use calyx_ir as ir;
2
3const NODE_ID: ir::Attribute =
4    ir::Attribute::Internal(ir::InternalAttr::NODE_ID);
5const BEGIN_ID: ir::Attribute =
6    ir::Attribute::Internal(ir::InternalAttr::BEGIN_ID);
7const END_ID: ir::Attribute = ir::Attribute::Internal(ir::InternalAttr::END_ID);
8
9/// Adding "NODE_ID", "BEGIN_ID", and "END_ID" attribute to control statement
10pub struct ControlId;
11
12impl ControlId {
13    fn compute_unique_ids_static(
14        scon: &mut ir::StaticControl,
15        mut cur_state: u64,
16        two_if_ids: bool,
17    ) -> u64 {
18        match scon {
19            ir::StaticControl::Empty(_) => cur_state,
20            ir::StaticControl::Enable(ir::StaticEnable {
21                attributes, ..
22            })
23            | ir::StaticControl::Invoke(ir::StaticInvoke {
24                attributes, ..
25            }) => {
26                attributes.insert(NODE_ID, cur_state);
27                cur_state + 1
28            }
29            ir::StaticControl::Repeat(ir::StaticRepeat {
30                attributes,
31                body,
32                ..
33            }) => {
34                attributes.insert(NODE_ID, cur_state);
35                cur_state += 1;
36                Self::compute_unique_ids_static(body, cur_state, two_if_ids)
37            }
38            ir::StaticControl::Par(ir::StaticPar {
39                stmts, attributes, ..
40            })
41            | ir::StaticControl::Seq(ir::StaticSeq {
42                stmts, attributes, ..
43            }) => {
44                attributes.insert(NODE_ID, cur_state);
45                cur_state += 1;
46                stmts.iter_mut().for_each(|stmt| {
47                    let new_state = Self::compute_unique_ids_static(
48                        stmt, cur_state, two_if_ids,
49                    );
50                    cur_state = new_state;
51                });
52                cur_state
53            }
54            ir::StaticControl::If(ir::StaticIf {
55                tbranch,
56                fbranch,
57                attributes,
58                ..
59            }) => {
60                if two_if_ids {
61                    attributes.insert(BEGIN_ID, cur_state);
62                    cur_state += 1;
63                    cur_state = Self::compute_unique_ids_static(
64                        tbranch, cur_state, two_if_ids,
65                    );
66                    cur_state = Self::compute_unique_ids_static(
67                        fbranch, cur_state, two_if_ids,
68                    );
69                    attributes.insert(END_ID, cur_state);
70                    cur_state + 1
71                } else {
72                    attributes.insert(NODE_ID, cur_state);
73                    cur_state += 1;
74                    cur_state = Self::compute_unique_ids_static(
75                        tbranch, cur_state, two_if_ids,
76                    );
77                    cur_state = Self::compute_unique_ids_static(
78                        fbranch, cur_state, two_if_ids,
79                    );
80                    cur_state + 1
81                }
82            }
83        }
84    }
85
86    /// Adds the @NODE_ID attribute to all control stmts except emtpy ones.
87    /// If two_if_ids is true, then if statements get a BEGIN_ID and END_ID instead
88    /// of a NODE_ID
89    ///
90    /// ## Example:
91    /// ```
92    /// seq { A; if cond {X} else{Y}; par { C; D; }; E }
93    /// ```
94    ///
95    /// gets the labels (if two_if_ids is):
96    ///
97    /// ```
98    /// @NODE_ID(0)seq {
99    ///   @NODE_ID(1) A;
100    ///   @BEGIN_ID(2) @END_ID(5) if cond {
101    ///     @NODE_ID(3) X
102    ///   }
103    ///   else{
104    ///     @NODE_ID(4) Y
105    ///   }
106    ///   @NODE_ID(6) par {
107    ///     @NODE_ID(7) C;
108    ///     @NODE_ID(8) D;
109    ///   }
110    ///   @NODE_ID(9) E;
111    /// }
112    /// ```
113    /// if two_if_ids were false, the if statement would just get a single NODE_ID
114    pub fn compute_unique_ids(
115        con: &mut ir::Control,
116        mut cur_state: u64,
117        two_if_ids: bool,
118    ) -> u64 {
119        match con {
120            ir::Control::Enable(ir::Enable { attributes, .. })
121            | ir::Control::Invoke(ir::Invoke { attributes, .. })
122            | ir::Control::FSMEnable(ir::FSMEnable { attributes, .. }) => {
123                attributes.insert(NODE_ID, cur_state);
124                cur_state + 1
125            }
126            ir::Control::Par(ir::Par {
127                stmts, attributes, ..
128            })
129            | ir::Control::Seq(ir::Seq {
130                stmts, attributes, ..
131            }) => {
132                attributes.insert(NODE_ID, cur_state);
133                cur_state += 1;
134                stmts.iter_mut().for_each(|stmt| {
135                    let new_state =
136                        Self::compute_unique_ids(stmt, cur_state, two_if_ids);
137                    cur_state = new_state;
138                });
139                cur_state
140            }
141            ir::Control::If(ir::If {
142                tbranch,
143                fbranch,
144                attributes,
145                ..
146            }) => {
147                if two_if_ids {
148                    attributes.insert(BEGIN_ID, cur_state);
149                    cur_state += 1;
150                    cur_state = Self::compute_unique_ids(
151                        tbranch, cur_state, two_if_ids,
152                    );
153                    cur_state = Self::compute_unique_ids(
154                        fbranch, cur_state, two_if_ids,
155                    );
156                    attributes.insert(END_ID, cur_state);
157                    cur_state + 1
158                } else {
159                    attributes.insert(NODE_ID, cur_state);
160                    cur_state += 1;
161                    cur_state = Self::compute_unique_ids(
162                        tbranch, cur_state, two_if_ids,
163                    );
164                    cur_state = Self::compute_unique_ids(
165                        fbranch, cur_state, two_if_ids,
166                    );
167                    cur_state + 1
168                }
169            }
170            ir::Control::While(ir::While {
171                body, attributes, ..
172            })
173            | ir::Control::Repeat(ir::Repeat {
174                body, attributes, ..
175            }) => {
176                attributes.insert(NODE_ID, cur_state);
177                cur_state += 1;
178                Self::compute_unique_ids(body, cur_state, two_if_ids)
179            }
180            ir::Control::Static(s) => {
181                Self::compute_unique_ids_static(s, cur_state, two_if_ids)
182            }
183            ir::Control::Empty(_) => cur_state,
184        }
185    }
186
187    // Gets attribute s from c, panics otherwise. Should be used when you know
188    // that c has attribute s.
189    pub fn get_guaranteed_attribute<A>(c: &ir::Control, attr: A) -> u64
190    where
191        A: Into<ir::Attribute>,
192    {
193        c.get_attribute(attr.into()).unwrap_or_else(||unreachable!(
194          "called get_guaranteed_attribute, meaning we had to be sure it had the attribute"
195      ))
196    }
197
198    // Gets attribute s from c, panics otherwise. Should be used when you know
199    // that c has attribute s.
200    pub fn get_guaranteed_attribute_static<A>(
201        sc: &ir::StaticControl,
202        attr: A,
203    ) -> u64
204    where
205        A: Into<ir::Attribute>,
206    {
207        sc.get_attribute(attr.into()).unwrap_or_else(||unreachable!(
208          "called get_guaranteed_attribute_static, meaning we had to be sure it had the attribute"
209      ))
210    }
211
212    // Gets attribute NODE_ID from c
213    pub fn get_guaranteed_id(c: &ir::Control) -> u64 {
214        Self::get_guaranteed_attribute(c, NODE_ID)
215    }
216
217    // Gets attribute NODE_ID from c
218    pub fn get_guaranteed_id_static(sc: &ir::StaticControl) -> u64 {
219        Self::get_guaranteed_attribute_static(sc, NODE_ID)
220    }
221
222    // takes in a static control scon, and adds unique id to each static enable.
223    // Returns cur_state, i.e., what the next enable should be labeled as
224    pub fn add_static_enable_ids_static(
225        scon: &mut ir::StaticControl,
226        mut cur_state: u64,
227    ) -> u64 {
228        match scon {
229            ir::StaticControl::Enable(se) => {
230                se.attributes.insert(NODE_ID, cur_state);
231                cur_state + 1
232            }
233            ir::StaticControl::Invoke(_) | ir::StaticControl::Empty(_) => {
234                cur_state
235            }
236            ir::StaticControl::Par(ir::StaticPar { stmts, .. })
237            | ir::StaticControl::Seq(ir::StaticSeq { stmts, .. }) => {
238                for stmt in stmts {
239                    let new_state =
240                        Self::add_static_enable_ids_static(stmt, cur_state);
241                    cur_state = new_state
242                }
243                cur_state
244            }
245            ir::StaticControl::If(ir::StaticIf {
246                tbranch, fbranch, ..
247            }) => {
248                let mut new_state =
249                    Self::add_static_enable_ids_static(tbranch, cur_state);
250                cur_state = new_state;
251                new_state =
252                    Self::add_static_enable_ids_static(fbranch, cur_state);
253                new_state
254            }
255            ir::StaticControl::Repeat(ir::StaticRepeat { body, .. }) => {
256                Self::add_static_enable_ids_static(body, cur_state)
257            }
258        }
259    }
260
261    // takes in ir::Control `con`, and adds unique id to every static enable within it.
262    // returns u64 `cur_state` that says what the next staticenable should be labeled as.
263    pub fn add_static_enable_ids(
264        con: &mut ir::Control,
265        mut cur_state: u64,
266    ) -> u64 {
267        match con {
268            ir::Control::Enable(_)
269            | ir::Control::Invoke(_)
270            | ir::Control::Empty(_) => cur_state,
271            ir::Control::Par(ir::Par { stmts, .. })
272            | ir::Control::Seq(ir::Seq { stmts, .. }) => {
273                for stmt in stmts {
274                    let new_state =
275                        Self::add_static_enable_ids(stmt, cur_state);
276                    cur_state = new_state
277                }
278                cur_state
279            }
280            ir::Control::If(ir::If {
281                tbranch, fbranch, ..
282            }) => {
283                let mut new_state =
284                    Self::add_static_enable_ids(tbranch, cur_state);
285                cur_state = new_state;
286                new_state = Self::add_static_enable_ids(fbranch, cur_state);
287                new_state
288            }
289            ir::Control::While(ir::While { body, .. })
290            | ir::Control::Repeat(ir::Repeat { body, .. }) => {
291                Self::add_static_enable_ids(body, cur_state)
292            }
293            ir::Control::Static(s) => {
294                Self::add_static_enable_ids_static(s, cur_state)
295            }
296            ir::Control::FSMEnable(_) => todo!(),
297        }
298    }
299
300    // Gets NODE_ID from StaticEnable se, panics otherwise. Should be used when you know
301    // that se has attributes NODE_ID.
302    pub fn get_guaranteed_enable_id(se: &ir::StaticEnable) -> u64 {
303        se.get_attribute(NODE_ID).unwrap_or_else(||unreachable!(
304          "called get_guaranteed_enable_id, meaning we had to be sure it had a NODE_ID attribute"
305      ))
306    }
307}