cherry pick

autoprecompiles/src/trace_handler.rs

@@ -2,9 +2,13 @@ use itertools::Itertools;
 use rayon::prelude::*;
 use std::cmp::Eq;
 use std::collections::HashMap;
-use std::hash::Hash;
+use std::ops::{Add, Mul, Neg, Sub};
 
+use crate::adapter::Adapter;
+use crate::expression::RowEvaluator;
 use crate::Apc;
+use crate::InstructionHandler;
+use crate::SymbolicBusInteraction;
 
 /// Returns data needed for constructing the APC trace.
 pub struct TraceHandlerData<'a, F> {
@@ -108,3 +112,48 @@ impl<F> Trace<F> {
         Self { values, width }
     }
 }
+
+pub struct InteractionEvaluator<
+    'a,
+    F: Add<Output = F> + Sub<Output = F> + Mul<Output = F> + Neg<Output = F> + Copy,
+> {
+    pub row_evaluator: RowEvaluator<'a, F>,
+}
+
+impl<'a, F: Add<Output = F> + Sub<Output = F> + Mul<Output = F> + Neg<Output = F> + Copy>
+    InteractionEvaluator<'a, F>
+{
+    pub fn new(row_evaluator: RowEvaluator<'a, F>) -> Self {
+        Self { row_evaluator }
+    }
+
+    pub fn evaluate_bus_interactions(
+        &self,
+        bus_interactions: &Vec<&SymbolicBusInteraction<F>>,
+        filter_by: impl Fn(&SymbolicBusInteraction<F>) -> bool,
+    ) -> Vec<ConcreteBusInteraction<F>> {
+        bus_interactions
+            .iter()
+            .filter(|&bus_interaction| filter_by(bus_interaction))
+            .map(|bus_interaction| {
+                let mult = self.row_evaluator.eval_expr(&bus_interaction.mult);
+                let args = bus_interaction
+                    .args
+                    .iter()
+                    .map(|arg| self.row_evaluator.eval_expr(arg))
+                    .collect_vec();
+                ConcreteBusInteraction {
+                    id: bus_interaction.id,
+                    mult,
+                    args,
+                }
+            })
+            .collect()
+    }
+}
+
+pub struct ConcreteBusInteraction<F> {
+    pub id: u64,
+    pub mult: F,
+    pub args: Vec<F>,
+}

openvm/src/powdr_extension/chip.rs

@@ -217,33 +217,6 @@ impl<F: PrimeField32> From<powdr_autoprecompiles::SymbolicBusInteraction<F>>
     }
 }
 
-pub struct RangeCheckerSend<F> {
-    pub mult: AlgebraicExpression<F>,
-    pub value: AlgebraicExpression<F>,
-    pub max_bits: AlgebraicExpression<F>,
-}
-
-impl<F: PrimeField32> TryFrom<&powdr_autoprecompiles::SymbolicBusInteraction<F>>
-    for RangeCheckerSend<F>
-{
-    type Error = ();
-
-    fn try_from(i: &powdr_autoprecompiles::SymbolicBusInteraction<F>) -> Result<Self, Self::Error> {
-        if i.id == 3 {
-            assert_eq!(i.args.len(), 2);
-            let value = &i.args[0];
-            let max_bits = &i.args[1];
-            Ok(Self {
-                mult: i.mult.clone(),
-                value: value.clone(),
-                max_bits: max_bits.clone(),
-            })
-        } else {
-            Err(())
-        }
-    }
-}
-
 impl<F: PrimeField32> PowdrAir<F> {
     pub fn new(apc: Arc<Apc<F, Instr<F>>>) -> Self {
         let (column_index_by_poly_id, columns): (BTreeMap<_, _>, Vec<_>) = apc

openvm/src/powdr_extension/executor/mod.rs

@@ -16,11 +16,12 @@ use crate::{
 use powdr_autoprecompiles::{
     adapter::Adapter,
     expression::RowEvaluator,
-    trace_handler::{Trace, TraceHandler, TraceHandlerData},
+    trace_handler::{
+        ConcreteBusInteraction, InteractionEvaluator, Trace, TraceHandler, TraceHandlerData,
+    },
     Apc,
 };
 
-use super::chip::RangeCheckerSend;
 use itertools::Itertools;
 use openvm_circuit::{
     arch::VmConfig, system::memory::MemoryController, utils::next_power_of_two_or_zero,
@@ -192,7 +193,7 @@ impl<F: PrimeField32> PowdrExecutor<F> {
         } = trace_handler.data(&self.apc);
 
         // precompute the symbolic bus sends to the range checker for each original instruction
-        let range_checker_sends_per_original_instruction: Vec<Vec<RangeCheckerSend<_>>> = self
+        let range_checker_sends_per_original_instruction = self
             .apc
             .instructions()
             .iter()
@@ -201,7 +202,7 @@ impl<F: PrimeField32> PowdrExecutor<F> {
                     .get_instruction_air(instruction)
                     .bus_interactions
                     .iter()
-                    .filter_map(|interaction| interaction.try_into().ok())
+                    .filter(|interaction| interaction.id == 3)
                     .collect_vec()
             })
             .collect_vec();
@@ -228,23 +229,17 @@ impl<F: PrimeField32> PowdrExecutor<F> {
                         .zip_eq(&range_checker_sends_per_original_instruction)
                         .zip_eq(&dummy_trace_index_to_apc_index_by_instruction)
                 {
-                    let evaluator = RowEvaluator::new(dummy_row, None);
+                    let interaction_evaluator =
+                        InteractionEvaluator::new(RowEvaluator::new(dummy_row, None));
 
-                    // first remove the side effects of this row on the main periphery
-                    for range_checker_send in range_checker_sends {
-                        let mult = evaluator
-                            .eval_expr(&range_checker_send.mult)
-                            .as_canonical_u32();
-                        let value = evaluator
-                            .eval_expr(&range_checker_send.value)
-                            .as_canonical_u32();
-                        let max_bits = evaluator
-                            .eval_expr(&range_checker_send.max_bits)
-                            .as_canonical_u32();
-                        for _ in 0..mult {
-                            self.periphery
-                                .range_checker
-                                .remove_count(value, max_bits as usize);
+                    for ConcreteBusInteraction { mult, args, .. } in interaction_evaluator
+                        .evaluate_bus_interactions(range_checker_sends, |_| true)
+                    {
+                        for _ in 0..mult.as_canonical_u32() {
+                            self.periphery.range_checker.remove_count(
+                                args[0].as_canonical_u32(),
+                                args[1].as_canonical_u32() as usize,
+                            );
                         }
                     }
 
@@ -256,20 +251,21 @@ impl<F: PrimeField32> PowdrExecutor<F> {
                 // Set the is_valid column to 1
                 row_slice[is_valid_index] = F::ONE;
 
-                let evaluator = RowEvaluator::new(row_slice, Some(column_index_by_poly_id));
+                let interaction_evaluator =
+                    InteractionEvaluator::new(RowEvaluator::new(row_slice, None));
 
                 // replay the side effects of this row on the main periphery
-                // TODO: this could be done in parallel since `self.periphery` is thread safe, but is it worth it? cc @qwang98
-                for bus_interaction in &bus_interactions {
-                    let mult = evaluator
-                        .eval_expr(&bus_interaction.mult)
-                        .as_canonical_u32();
-                    let args = bus_interaction
-                        .args
-                        .iter()
-                        .map(|arg| evaluator.eval_expr(arg).as_canonical_u32());
-
-                    self.periphery.apply(bus_interaction.id, mult, args);
+                for ConcreteBusInteraction { id, mult, args } in interaction_evaluator
+                    .evaluate_bus_interactions(
+                        &self.apc.machine().bus_interactions.iter().collect_vec(),
+                        |_| true,
+                    )
+                {
+                    self.periphery.apply(
+                        id as u16,
+                        mult.as_canonical_u32(),
+                        args.iter().map(|arg| arg.as_canonical_u32()),
+                    );
                 }
             });