calyx_opt/analysis/
control_ports.rs1use super::AssignmentAnalysis;
2use calyx_ir::{self as ir, RRC};
3use itertools::Itertools;
4use std::{
5 collections::{HashMap, HashSet},
6 rc::Rc,
7};
8
9type PortMap = HashMap<ir::Id, Vec<RRC<ir::Port>>>;
10type Binding = (Vec<(ir::Id, RRC<ir::Cell>)>, Vec<(ir::Id, RRC<ir::Port>)>);
11type InvokeMap = HashMap<ir::Id, Vec<Binding>>;
12
13pub struct ControlPorts<const INVOKE_MAP: bool> {
17 cg_to_port: PortMap,
19 invoke_map: InvokeMap,
21}
22
23impl<const INVOKE_MAP: bool> ControlPorts<INVOKE_MAP> {
24 pub fn get(&self, group: &ir::Id) -> Option<&Vec<RRC<ir::Port>>> {
26 self.cg_to_port.get(group)
27 }
28
29 pub fn remove(&mut self, group: &ir::Id) -> Option<Vec<RRC<ir::Port>>> {
31 self.cg_to_port.remove(group)
32 }
33
34 pub fn get_bindings(&self, instance: &ir::Id) -> Option<&Vec<Binding>> {
36 if INVOKE_MAP {
37 self.invoke_map.get(instance)
38 } else {
39 panic!("ControlPorts instance built without invoke_map")
40 }
41 }
42
43 pub fn get_all_bindings(self) -> InvokeMap {
45 if INVOKE_MAP {
46 self.invoke_map
47 } else {
48 panic!("ControlPorts instance built without invoke_map")
49 }
50 }
51}
52
53impl<const INVOKE_MAP: bool> ControlPorts<INVOKE_MAP> {
54 fn handle_invoke(
62 &mut self,
63 inputs: &[(ir::Id, ir::RRC<ir::Port>)],
64 outputs: &[(ir::Id, ir::RRC<ir::Port>)],
65 ref_cells: &[(ir::Id, ir::RRC<ir::Cell>)],
66 comp: &ir::RRC<ir::Cell>,
67 comb_group: &Option<ir::RRC<ir::CombGroup>>,
68 ) {
69 if let Some(c) = comb_group {
70 let cells = c
71 .borrow()
72 .assignments
73 .iter()
74 .analysis()
75 .cell_uses()
76 .map(|cell| cell.borrow().name())
77 .collect::<HashSet<_>>();
78
79 let ports =
81 inputs
82 .iter()
83 .map(|(_, port)| Rc::clone(port))
84 .filter(|port| {
85 cells.contains(&port.borrow().get_parent_name())
86 });
87 self.cg_to_port
88 .entry(c.borrow().name())
89 .or_default()
90 .extend(ports);
91 }
92 if INVOKE_MAP {
93 let name = comp.borrow().name();
94 let bindings =
95 inputs.iter().chain(outputs.iter()).cloned().collect_vec();
96 self.invoke_map
97 .entry(name)
98 .or_default()
99 .push((ref_cells.to_vec(), bindings));
100 }
101 }
102
103 fn construct_static(&mut self, scon: &ir::StaticControl) {
107 match scon {
108 ir::StaticControl::Empty(_) | ir::StaticControl::Enable(_) => (),
109 ir::StaticControl::Repeat(ir::StaticRepeat { body, .. }) => {
110 self.construct_static(body)
111 }
112 ir::StaticControl::Seq(ir::StaticSeq { stmts, .. })
113 | ir::StaticControl::Par(ir::StaticPar { stmts, .. }) => {
114 stmts.iter().for_each(|stmt| self.construct_static(stmt));
115 }
116 ir::StaticControl::If(ir::StaticIf {
117 tbranch, fbranch, ..
118 }) => {
119 self.construct_static(tbranch);
120 self.construct_static(fbranch);
121 }
122 ir::StaticControl::Invoke(ir::StaticInvoke {
123 comp,
124 inputs,
125 outputs,
126 ref_cells,
127 ..
128 }) => {
129 self.handle_invoke(inputs, outputs, ref_cells, comp, &None);
130 }
131 }
132 }
133
134 fn construct(&mut self, con: &ir::Control) {
135 match con {
136 ir::Control::Enable(_) | ir::Control::Empty(_) => {}
137 ir::Control::Invoke(ir::Invoke {
138 comp,
139 comb_group,
140 inputs,
141 outputs,
142 ref_cells,
143 ..
144 }) => {
145 self.handle_invoke(
146 inputs, outputs, ref_cells, comp, comb_group,
147 );
148 }
149 ir::Control::If(ir::If {
150 cond,
151 port,
152 tbranch,
153 fbranch,
154 ..
155 }) => {
156 if let Some(c) = cond {
157 self.cg_to_port
158 .entry(c.borrow().name())
159 .or_default()
160 .push(Rc::clone(port));
161 }
162
163 self.construct(tbranch);
164 self.construct(fbranch);
165 }
166 ir::Control::While(ir::While {
167 cond, port, body, ..
168 }) => {
169 if let Some(c) = cond {
170 self.cg_to_port
171 .entry(c.borrow().name())
172 .or_default()
173 .push(Rc::clone(port));
174 }
175 self.construct(body);
176 }
177 ir::Control::Repeat(ir::Repeat { body, .. }) => {
178 self.construct(body);
179 }
180 ir::Control::Seq(ir::Seq { stmts, .. })
181 | ir::Control::Par(ir::Par { stmts, .. }) => {
182 stmts.iter().for_each(|con| self.construct(con));
183 }
184 ir::Control::Static(sc) => {
185 self.construct_static(sc)
189 }
190 }
191 }
192}
193
194impl<const INVOKE_MAP: bool> From<&ir::Control> for ControlPorts<INVOKE_MAP> {
195 fn from(con: &ir::Control) -> Self {
196 let mut cp = ControlPorts {
197 cg_to_port: HashMap::new(),
198 invoke_map: HashMap::new(),
199 };
200 cp.construct(con);
201 cp.cg_to_port.values_mut().for_each(|v| {
203 *v = v.drain(..).unique_by(|p| p.borrow().canonical()).collect()
204 });
205 if INVOKE_MAP {
207 cp.invoke_map.values_mut().for_each(|v| {
208 *v = v
209 .drain(..)
210 .unique_by(|(cells, ports)| {
211 (
212 cells
213 .clone()
214 .into_iter()
215 .map(|(c, cell)| (c, cell.borrow().name()))
216 .collect_vec(),
217 ports
218 .clone()
219 .into_iter()
220 .map(|(p, v)| (p, v.borrow().canonical()))
221 .collect_vec(),
222 )
223 })
224 .collect()
225 });
226 }
227 cp
228 }
229}