1use std::{
2 collections::BTreeMap,
3 hash::{DefaultHasher, Hash, Hasher},
4};
5
6use air_parser::ast::{Identifier, NamespacedIdentifier};
7use miden_diagnostics::{SourceSpan, Span};
8use pretty_assertions::assert_eq;
9
10use crate::{CompileError, ir::*, passes::Visitor};
11
12pub fn strip_spans(mir: &mut Mir) {
13 let graph = mir.constraint_graph_mut();
14 let mut visitor = StripSpansVisitor::default();
15 match visitor.run(graph) {
16 Ok(_) => {}
17 Err(e) => {
18 panic!("Error stripping spans: {e:?}");
19 }
20 }
21}
22
23#[derive(Default)]
24pub struct StripSpansVisitor {
25 _done: BTreeMap<usize, bool>,
26 work_stack: Vec<Link<Node>>,
27}
28
29pub fn extract_roots(
30 graph: &Graph,
31 include_boundary: bool,
32 include_integrity: bool,
33 include_bus: bool,
34 include_func: bool,
35 include_eval: bool,
36) -> Vec<Link<Node>> {
37 let mut nodes = Vec::new();
38 if include_boundary {
39 let boundary_ref = graph.boundary_constraints_roots.borrow();
40 let boundary = boundary_ref.iter().map(|n| n.as_node());
41 nodes.extend(boundary);
42 }
43 if include_integrity {
44 let integrity_ref = graph.integrity_constraints_roots.borrow();
45 let integrity = integrity_ref.iter().map(|n| n.as_node());
46 nodes.extend(integrity);
47 }
48 if include_bus {
49 let buses = graph.get_bus_nodes();
50 let bus_columns = buses
51 .iter()
52 .flat_map(|b| b.borrow().columns.clone())
53 .map(|n| n.as_node());
54 let bus_latches = buses
55 .iter()
56 .flat_map(|b| b.borrow().latches.clone())
57 .map(|n| n.as_node());
58 nodes.extend(bus_columns);
59 nodes.extend(bus_latches);
60 }
61 if include_func {
62 let funcs = graph.get_function_nodes().into_iter().map(|n| n.as_node());
63 nodes.extend(funcs);
64 }
65 if include_eval {
66 let evals = graph.get_evaluator_nodes().into_iter().map(|n| n.as_node());
67 nodes.extend(evals);
68 }
69 nodes
70}
71
72pub fn extract_all_roots(graph: &Graph) -> Vec<Link<Node>> {
73 extract_roots(graph, true, true, true, true, true)
74}
75
76pub fn extract_boundary_roots(graph: &Graph) -> Vec<Link<Node>> {
77 extract_roots(graph, true, false, false, false, false)
78}
79
80pub fn extract_integrity_roots(graph: &Graph) -> Vec<Link<Node>> {
81 extract_roots(graph, false, true, false, false, false)
82}
83
84pub fn extract_bus_roots(graph: &Graph) -> Vec<Link<Node>> {
85 extract_roots(graph, false, false, true, false, false)
86}
87
88pub fn extract_function_roots(graph: &Graph) -> Vec<Link<Node>> {
89 extract_roots(graph, false, false, false, true, false)
90}
91
92pub fn extract_evaluator_roots(graph: &Graph) -> Vec<Link<Node>> {
93 extract_roots(graph, false, false, false, false, true)
94}
95
96impl Visitor for StripSpansVisitor {
97 fn work_stack(&mut self) -> &mut Vec<Link<Node>> {
98 &mut self.work_stack
99 }
100 fn root_nodes_to_visit(&self, graph: &Graph) -> Vec<Link<Node>> {
101 extract_all_roots(graph)
102 }
103
104 fn visit_function(
105 &mut self,
106 _graph: &mut Graph,
107 function: Link<Root>,
108 ) -> Result<(), CompileError> {
109 let mut function = function.as_function_mut().unwrap();
110 function.span = Default::default();
111 Ok(())
112 }
113
114 fn visit_evaluator(
115 &mut self,
116 _graph: &mut Graph,
117 evaluator: Link<Root>,
118 ) -> Result<(), CompileError> {
119 let mut evaluator = evaluator.as_evaluator_mut().unwrap();
120 evaluator.span = Default::default();
121 Ok(())
122 }
123
124 fn visit_enf(&mut self, _graph: &mut Graph, enf: Link<Op>) -> Result<(), CompileError> {
125 let mut enf = enf.as_enf_mut().unwrap();
126 enf.span = Default::default();
127 Ok(())
128 }
129
130 fn visit_boundary(
131 &mut self,
132 _graph: &mut Graph,
133 boundary: Link<Op>,
134 ) -> Result<(), CompileError> {
135 let mut boundary = boundary.as_boundary_mut().unwrap();
136 boundary.span = Default::default();
137 Ok(())
138 }
139
140 fn visit_add(&mut self, _graph: &mut Graph, add: Link<Op>) -> Result<(), CompileError> {
141 let mut add = add.as_add_mut().unwrap();
142 add.span = Default::default();
143 Ok(())
144 }
145
146 fn visit_sub(&mut self, _graph: &mut Graph, sub: Link<Op>) -> Result<(), CompileError> {
147 let mut sub = sub.as_sub_mut().unwrap();
148 sub.span = Default::default();
149 Ok(())
150 }
151
152 fn visit_mul(&mut self, _graph: &mut Graph, mul: Link<Op>) -> Result<(), CompileError> {
153 let mut mul = mul.as_mul_mut().unwrap();
154 mul.span = Default::default();
155 Ok(())
156 }
157
158 fn visit_exp(&mut self, _graph: &mut Graph, exp: Link<Op>) -> Result<(), CompileError> {
159 let mut exp = exp.as_exp_mut().unwrap();
160 exp.span = Default::default();
161 Ok(())
162 }
163
164 fn visit_if(&mut self, _graph: &mut Graph, if_node: Link<Op>) -> Result<(), CompileError> {
165 let mut if_node = if_node.as_if_mut().unwrap();
166 if_node.span = Default::default();
167 Ok(())
168 }
169
170 fn visit_for(&mut self, _graph: &mut Graph, for_node: Link<Op>) -> Result<(), CompileError> {
171 let mut for_node = for_node.as_for_mut().unwrap();
172 for_node.span = Default::default();
173 Ok(())
174 }
175
176 fn visit_call(&mut self, _graph: &mut Graph, call: Link<Op>) -> Result<(), CompileError> {
177 let mut call = call.as_call_mut().unwrap();
178 call.span = Default::default();
179 Ok(())
180 }
181
182 fn visit_fold(&mut self, _graph: &mut Graph, fold: Link<Op>) -> Result<(), CompileError> {
183 let mut fold = fold.as_fold_mut().unwrap();
184 fold.span = Default::default();
185 Ok(())
186 }
187
188 fn visit_vector(&mut self, _graph: &mut Graph, vector: Link<Op>) -> Result<(), CompileError> {
189 let mut vector = vector.as_vector_mut().unwrap();
190 vector.span = Default::default();
191 Ok(())
192 }
193
194 fn visit_matrix(&mut self, _graph: &mut Graph, matrix: Link<Op>) -> Result<(), CompileError> {
195 let mut matrix = matrix.as_matrix_mut().unwrap();
196 matrix.span = Default::default();
197 Ok(())
198 }
199
200 fn visit_accessor(
201 &mut self,
202 _graph: &mut Graph,
203 accessor: Link<Op>,
204 ) -> Result<(), CompileError> {
205 let mut accessor = accessor.as_accessor_mut().unwrap();
206 accessor.span = Default::default();
207 Ok(())
208 }
209
210 fn visit_bus_op(&mut self, _graph: &mut Graph, bus_op: Link<Op>) -> Result<(), CompileError> {
211 let mut bus_op = bus_op.as_bus_op_mut().unwrap();
212 bus_op.span = Default::default();
213 Ok(())
214 }
215
216 fn visit_parameter(
217 &mut self,
218 _graph: &mut Graph,
219 parameter: Link<Op>,
220 ) -> Result<(), CompileError> {
221 let mut parameter = parameter.as_parameter_mut().unwrap();
222 parameter.span = Default::default();
223 Ok(())
224 }
225
226 fn visit_value(&mut self, _graph: &mut Graph, value: Link<Op>) -> Result<(), CompileError> {
227 let mut value = value.as_value_mut().unwrap();
228 value.value.span = Default::default();
229 match &mut value.value.value {
230 MirValue::Constant(_) => {}
231 MirValue::TraceAccess(_) => {}
232 MirValue::PeriodicColumn(v) => {
233 v.name.module.0 = Span::new(SourceSpan::default(), v.name.module.0.item);
234 match v.name.item {
235 NamespacedIdentifier::Function(f) => {
236 v.name.item = NamespacedIdentifier::Function(Identifier::new(
237 SourceSpan::default(),
238 f.0.item,
239 ));
240 }
241 NamespacedIdentifier::Binding(b) => {
242 v.name.item = NamespacedIdentifier::Binding(Identifier::new(
243 SourceSpan::default(),
244 b.0.item,
245 ));
246 }
247 };
248 }
249 MirValue::PublicInput(v) => {
250 v.name.0 = Span::new(SourceSpan::default(), v.name.0.item);
251 }
252 MirValue::PublicInputTable(v) => {
253 v.table_name.0 = Span::new(SourceSpan::default(), v.table_name.0.item);
254 }
255 MirValue::RandomValue(_) => {}
256 MirValue::TraceAccessBinding(_) => {}
257 MirValue::BusAccess(_) => {}
258 MirValue::Null | MirValue::Unconstrained => {}
259 }
260 Ok(())
261 }
262}
263
264pub fn hash<T: Hash>(val: &T) -> u64 {
265 let mut hasher = DefaultHasher::new();
266 val.hash(&mut hasher);
267 hasher.finish()
268}
269
270fn extract_and_compare_mir(
271 lhs: &mut Mir,
272 rhs: &mut Mir,
273 extract: impl Fn(&Graph) -> Vec<Link<Node>>,
274) -> bool {
275 strip_spans(lhs);
276 strip_spans(rhs);
277 let lhs = extract(lhs.constraint_graph());
278 let rhs = extract(rhs.constraint_graph());
279 hash(&lhs) == hash(&rhs)
280}
281
282pub fn compare_mir(lhs: &mut Mir, rhs: &mut Mir) -> bool {
283 extract_and_compare_mir(lhs, rhs, extract_all_roots)
284}
285
286pub fn compare_boundary(lhs: &mut Mir, rhs: &mut Mir) {
287 extract_and_compare_mir(lhs, rhs, extract_boundary_roots);
288}
289
290pub fn compare_integrity(lhs: &mut Mir, rhs: &mut Mir) {
291 extract_and_compare_mir(lhs, rhs, extract_integrity_roots);
292}
293
294pub fn compare_bus(lhs: &mut Mir, rhs: &mut Mir) {
295 extract_and_compare_mir(lhs, rhs, extract_bus_roots);
296}
297
298pub fn compare_function(lhs: &mut Mir, rhs: &mut Mir) {
299 extract_and_compare_mir(lhs, rhs, extract_function_roots);
300}
301
302pub fn compare_evaluator(lhs: &mut Mir, rhs: &mut Mir) {
303 extract_and_compare_mir(lhs, rhs, extract_evaluator_roots);
304}
305
306fn extract_and_assert_mir_eq(
307 lhs: &mut Mir,
308 rhs: &mut Mir,
309 extract: impl Fn(&Graph) -> Vec<Link<Node>>,
310) {
311 strip_spans(lhs);
312 strip_spans(rhs);
313 let lhs = extract(lhs.constraint_graph());
314 let rhs = extract(rhs.constraint_graph());
315 assert_eq!(lhs, rhs);
316}
317
318pub fn assert_mir_eq(lhs: &mut Mir, rhs: &mut Mir) {
319 extract_and_assert_mir_eq(lhs, rhs, extract_all_roots);
320}
321
322pub fn assert_boundary_eq(lhs: &mut Mir, rhs: &mut Mir) {
323 extract_and_assert_mir_eq(lhs, rhs, extract_boundary_roots);
324}
325
326pub fn assert_integrity_eq(lhs: &mut Mir, rhs: &mut Mir) {
327 extract_and_assert_mir_eq(lhs, rhs, extract_integrity_roots);
328}
329
330pub fn assert_bus_eq(lhs: &mut Mir, rhs: &mut Mir) {
331 extract_and_assert_mir_eq(lhs, rhs, extract_bus_roots);
332}
333
334pub fn assert_function_eq(lhs: &mut Mir, rhs: &mut Mir) {
335 extract_and_assert_mir_eq(lhs, rhs, extract_function_roots);
336}
337
338pub fn assert_evaluator_eq(lhs: &mut Mir, rhs: &mut Mir) {
339 extract_and_assert_mir_eq(lhs, rhs, extract_evaluator_roots);
340}