air_codegen_masm/
visitor.rs

1use crate::utils::contraint_root_domain;
2use air_ir::{
3    Air, ConstraintDomain, ConstraintRoot, NodeIndex, Operation, PeriodicColumn, TraceSegmentId,
4    Value,
5};
6
7pub trait AirVisitor<'ast> {
8    type Value;
9    type Error;
10
11    fn visit_boundary_constraint(
12        &mut self,
13        constraint: &'ast ConstraintRoot,
14        trace_segment: TraceSegmentId,
15    ) -> Result<Self::Value, Self::Error>;
16
17    fn visit_air(&mut self) -> Result<Self::Value, Self::Error>;
18
19    fn visit_integrity_constraint(
20        &mut self,
21        constraint: &'ast ConstraintRoot,
22        trace_segment: TraceSegmentId,
23    ) -> Result<Self::Value, Self::Error>;
24
25    fn visit_node_index(&mut self, node_index: &'ast NodeIndex)
26        -> Result<Self::Value, Self::Error>;
27
28    fn visit_operation(&mut self, op: &'ast Operation) -> Result<Self::Value, Self::Error>;
29
30    fn visit_periodic_column(
31        &mut self,
32        columns: &'ast PeriodicColumn,
33    ) -> Result<Self::Value, Self::Error>;
34
35    fn visit_value(&mut self, value: &'ast Value) -> Result<Self::Value, Self::Error>;
36}
37
38pub fn walk_periodic_columns<'ast, V: AirVisitor<'ast>>(
39    visitor: &mut V,
40    ir: &'ast Air,
41) -> Result<(), V::Error> {
42    for column in ir.periodic_columns() {
43        visitor.visit_periodic_column(column)?;
44    }
45
46    Ok(())
47}
48
49pub fn walk_boundary_constraints<'ast, V: AirVisitor<'ast>>(
50    visitor: &mut V,
51    ir: &'ast Air,
52    segment: TraceSegmentId,
53    constraint_domain: ConstraintDomain,
54) -> Result<(), V::Error> {
55    let mut constraints: Vec<&'ast ConstraintRoot> =
56        ir.boundary_constraints(segment).iter().collect();
57
58    constraints.sort_by_key(contraint_root_domain);
59
60    for boundary in constraints
61        .iter()
62        .filter(|c| c.domain() == constraint_domain)
63    {
64        visitor.visit_boundary_constraint(boundary, segment)?;
65    }
66
67    Ok(())
68}
69
70pub fn walk_integrity_constraints<'ast, V: AirVisitor<'ast>>(
71    visitor: &mut V,
72    ir: &'ast Air,
73    trace_segment: TraceSegmentId,
74) -> Result<(), V::Error> {
75    for integrity in ir.integrity_constraints(trace_segment) {
76        visitor.visit_integrity_constraint(integrity, trace_segment)?;
77    }
78
79    Ok(())
80}
air_codegen_masm/visitor.rs

air_codegen_masm/
visitor.rs
