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(_)
137 | ir::Control::Empty(_)
138 | ir::Control::FSMEnable(_) => {}
139 ir::Control::Invoke(ir::Invoke {
140 comp,
141 comb_group,
142 inputs,
143 outputs,
144 ref_cells,
145 ..
146 }) => {
147 self.handle_invoke(
148 inputs, outputs, ref_cells, comp, comb_group,
149 );
150 }
151 ir::Control::If(ir::If {
152 cond,
153 port,
154 tbranch,
155 fbranch,
156 ..
157 }) => {
158 if let Some(c) = cond {
159 self.cg_to_port
160 .entry(c.borrow().name())
161 .or_default()
162 .push(Rc::clone(port));
163 }
164
165 self.construct(tbranch);
166 self.construct(fbranch);
167 }
168 ir::Control::While(ir::While {
169 cond, port, body, ..
170 }) => {
171 if let Some(c) = cond {
172 self.cg_to_port
173 .entry(c.borrow().name())
174 .or_default()
175 .push(Rc::clone(port));
176 }
177 self.construct(body);
178 }
179 ir::Control::Repeat(ir::Repeat { body, .. }) => {
180 self.construct(body);
181 }
182 ir::Control::Seq(ir::Seq { stmts, .. })
183 | ir::Control::Par(ir::Par { stmts, .. }) => {
184 stmts.iter().for_each(|con| self.construct(con));
185 }
186 ir::Control::Static(sc) => {
187 self.construct_static(sc)
191 }
192 }
193 }
194}
195
196impl<const INVOKE_MAP: bool> From<&ir::Control> for ControlPorts<INVOKE_MAP> {
197 fn from(con: &ir::Control) -> Self {
198 let mut cp = ControlPorts {
199 cg_to_port: HashMap::new(),
200 invoke_map: HashMap::new(),
201 };
202 cp.construct(con);
203 cp.cg_to_port.values_mut().for_each(|v| {
205 *v = v.drain(..).unique_by(|p| p.borrow().canonical()).collect()
206 });
207 if INVOKE_MAP {
209 cp.invoke_map.values_mut().for_each(|v| {
210 *v = v
211 .drain(..)
212 .unique_by(|(cells, ports)| {
213 (
214 cells
215 .clone()
216 .into_iter()
217 .map(|(c, cell)| (c, cell.borrow().name()))
218 .collect_vec(),
219 ports
220 .clone()
221 .into_iter()
222 .map(|(p, v)| (p, v.borrow().canonical()))
223 .collect_vec(),
224 )
225 })
226 .collect()
227 });
228 }
229 cp
230 }
231}