Visit all bus operations

Author: Leo-Besancon

air/src/tests/buses.rs

@@ -49,6 +49,10 @@ fn buses_in_integrity_constraints() {
         main: [a],
     }
 
+    fn double(a: felt) -> felt {
+        return a+a;
+    }
+
     buses {
         unit p,
         mult q,
@@ -69,7 +73,7 @@ fn buses_in_integrity_constraints() {
     }
 
     integrity_constraints {
-        p.add(1) when 1;
+        p.add(double(a)) when 1;
         p.rem(1) when 1;
         q.add(1, 2) when 1;
         q.add(1, 2) when 1;

mir/src/passes/inlining.rs

@@ -204,6 +204,11 @@ impl Visitor for InliningFirstPass<'_> {
         let evaluators = graph.get_evaluator_nodes();
         let boundary_constraints_roots_ref = graph.boundary_constraints_roots.borrow();
         let integrity_constraints_roots_ref = graph.integrity_constraints_roots.borrow();
+        let bus_roots: Vec<_> = graph
+            .buses
+            .values()
+            .flat_map(|b| b.borrow().clone().columns.into_iter().collect::<Vec<_>>())
+            .collect();
 
         let combined_roots = boundary_constraints_roots_ref
             .clone()
@@ -215,6 +220,7 @@ impl Visitor for InliningFirstPass<'_> {
                     .into_iter()
                     .map(|ic| ic.as_node()),
             )
+            .chain(bus_roots.into_iter().map(|b| b.as_node()))
             .chain(evaluators.into_iter().map(|e| e.as_node()))
             .chain(functions.into_iter().map(|f| f.as_node()));
         combined_roots.collect()

mir/src/passes/unrolling.rs

@@ -846,6 +846,11 @@ impl Visitor for UnrollingFirstPass<'_> {
     fn root_nodes_to_visit(&self, graph: &Graph) -> Vec<Link<Node>> {
         let boundary_constraints_roots_ref = graph.boundary_constraints_roots.borrow();
         let integrity_constraints_roots_ref = graph.integrity_constraints_roots.borrow();
+        let bus_roots: Vec<_> = graph
+            .buses
+            .values()
+            .flat_map(|b| b.borrow().clone().columns.into_iter().collect::<Vec<_>>())
+            .collect();
         let combined_roots = boundary_constraints_roots_ref
             .clone()
             .into_iter()
@@ -855,7 +860,8 @@ impl Visitor for UnrollingFirstPass<'_> {
                     .clone()
                     .into_iter()
                     .map(|ic| ic.as_node()),
-            );
+            )
+            .chain(bus_roots.into_iter().map(|b| b.as_node()));
         combined_roots.collect()
     }
 
@@ -894,7 +900,7 @@ impl Visitor for UnrollingFirstPass<'_> {
             Node::Accessor(a) => {
                 to_link_and(a.clone(), graph, |g, el| self.visit_accessor_bis(g, el))
             }
-            Node::BusOp(_b) => todo!(),
+            Node::BusOp(_b) => Ok(None),
             Node::Parameter(p) => {
                 to_link_and(p.clone(), graph, |g, el| self.visit_parameter_bis(g, el))
             }