powdr changes (#5)

* balance accumulator tree

* avoid eval_fri_log_up_phase

* Tests & Dummy airs

* ColumnsAir

---------

Co-authored-by: Leo Alt <leo@powdrlabs.com>
Co-authored-by: Gaston Zanitti <gzanitti@gmail.com>
Co-authored-by: Georg Wiese <georgwiese@gmail.com>

Author: 4 people

crates/stark-backend/src/air_builders/symbolic/mod.rs

@@ -407,14 +407,15 @@ impl<F: Field> InteractionPhaseAirBuilder for SymbolicRapBuilder<F> {
                 self.interaction_partitions.replace(interaction_partitions);
                 let perm_width = num_chunks + 1;
                 self.after_challenge = Self::new_after_challenge(&[perm_width]);
-            }
 
-            let phases_shapes = self.rap_phase_seq_kind.shape();
-            let phase_shape = phases_shapes.first().unwrap();
+                let phases_shapes = self.rap_phase_seq_kind.shape();
+
+                let phase_shape = phases_shapes.first().unwrap();
 
-            self.challenges = Self::new_challenges(&[phase_shape.num_challenges]);
-            self.exposed_values_after_challenge =
-                Self::new_exposed_values_after_challenge(&[phase_shape.num_exposed_values]);
+                self.challenges = Self::new_challenges(&[phase_shape.num_challenges]);
+                self.exposed_values_after_challenge =
+                    Self::new_exposed_values_after_challenge(&[phase_shape.num_exposed_values]);
+            }
         }
     }
 

crates/stark-backend/src/interaction/mod.rs

@@ -254,6 +254,7 @@ pub enum RapPhaseSeqKind {
     /// Up to one phase with prover/verifier given by [[fri_log_up::FriLogUpPhase]] and
     /// constraints given by [[fri_log_up::eval_fri_log_up_phase]].
     FriLogUp,
+    None,
 }
 
 impl RapPhaseSeqKind {
@@ -264,6 +265,7 @@ impl RapPhaseSeqKind {
                 num_exposed_values: STARK_LU_NUM_EXPOSED_VALUES,
                 extra_opening_rots: vec![],
             }],
+            RapPhaseSeqKind::None => vec![],
         }
     }
 }

crates/stark-backend/src/interaction/rap.rs

@@ -38,6 +38,7 @@ where
                         builder.max_constraint_degree(),
                     );
                 }
+                RapPhaseSeqKind::None => {}
             }
         }
     }

crates/stark-backend/src/rap.rs

@@ -73,20 +73,32 @@ Rap<SymbolicRapBuilder<Val<SC>>> // for keygen to extract fixed data about the R
     + for<'a> Rap<DebugConstraintBuilder<'a, SC>> // for debugging
     + BaseAirWithPublicValues<Val<SC>>
     + PartitionedBaseAir<Val<SC>>
+    + ColumnsAir<Val<SC>>
     + Send + Sync
 {
     fn as_any(&self) -> &dyn Any;
     /// Name for display purposes
     fn name(&self) -> String;
 }
 
+/// Trait for AIRs that can provide column names
+pub trait ColumnsAir<F>: BaseAir<F> {
+    /// If available, returns the names of columns used in this AIR.
+    /// If the result is `Some(names)`, `names.len() == air.width()` should always
+    /// be true
+    fn columns(&self) -> Option<Vec<String>> {
+        None
+    }
+}
+
 impl<SC, T> AnyRap<SC> for T
 where
     SC: StarkGenericConfig,
     T: Rap<SymbolicRapBuilder<Val<SC>>>
         + for<'a> Rap<DebugConstraintBuilder<'a, SC>>
         + BaseAirWithPublicValues<Val<SC>>
         + PartitionedBaseAir<Val<SC>>
+        + ColumnsAir<Val<SC>>
         + Send
         + Sync
         + 'static,

crates/stark-backend/src/verifier/folder.rs

@@ -1,6 +1,6 @@
 use std::{
     marker::PhantomData,
-    ops::{AddAssign, MulAssign},
+    ops::{Add, AddAssign, MulAssign},
 };
 
 use p3_field::{ExtensionField, Field, FieldAlgebra};
@@ -59,15 +59,27 @@ where
         // We do a simple serial evaluation in topological order.
         // This can be parallelized if necessary.
         let exprs = self.eval_nodes(&dag.nodes);
-        for &idx in &dag.constraint_idx {
-            self.assert_zero(exprs[idx].clone());
-        }
+        let v: Vec<Expr> = dag
+            .constraint_idx
+            .iter()
+            .map(|idx| exprs[*idx].clone())
+            .rev()
+            .scan(F::ONE.into(), |state: &mut Expr, next_elem| {
+                let r = next_elem * state.clone();
+                *state *= self.alpha;
+                Some(r)
+            })
+            .collect();
+        self.accumulator = balanced_sum_rec(&v);
     }
+}
 
-    pub fn assert_zero(&mut self, x: impl Into<Expr>) {
-        let x = x.into();
-        self.accumulator *= self.alpha;
-        self.accumulator += x;
+fn balanced_sum_rec<E: Clone + Add<Output = E>>(v: &[E]) -> E {
+    if v.len() == 1 {
+        v[0].clone()
+    } else {
+        let (left, right) = v.split_at(v.len() / 2);
+        balanced_sum_rec(left) + balanced_sum_rec(right)
     }
 }
 

crates/stark-backend/tests/fib_selector_air/air.rs

@@ -3,7 +3,7 @@ use std::borrow::Borrow;
 use openvm_stark_backend::{
     interaction::{InteractionBuilder, LookupBus},
     p3_field::{Field, FieldAlgebra},
-    rap::{BaseAirWithPublicValues, PartitionedBaseAir},
+    rap::{BaseAirWithPublicValues, ColumnsAir, PartitionedBaseAir},
 };
 use openvm_stark_sdk::dummy_airs::fib_air::columns::{FibonacciCols, NUM_FIBONACCI_COLS};
 use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir, PairBuilder};
@@ -41,6 +41,8 @@ impl<F: Field> BaseAir<F> for FibonacciSelectorAir {
     }
 }
 
+impl<F: Field> ColumnsAir<F> for FibonacciSelectorAir {}
+
 impl<F: Field> BaseAirWithPublicValues<F> for FibonacciSelectorAir {
     fn num_public_values(&self) -> usize {
         3

crates/stark-backend/tests/fib_triples_air/air.rs

@@ -1,6 +1,6 @@
 use std::borrow::Borrow;
 
-use openvm_stark_backend::rap::{BaseAirWithPublicValues, PartitionedBaseAir};
+use openvm_stark_backend::rap::{BaseAirWithPublicValues, ColumnsAir, PartitionedBaseAir};
 use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir};
 use p3_matrix::Matrix;
 
@@ -15,6 +15,8 @@ impl<F> BaseAir<F> for FibonacciAir {
     }
 }
 
+impl<F> ColumnsAir<F> for FibonacciAir {}
+
 impl<F> BaseAirWithPublicValues<F> for FibonacciAir {
     fn num_public_values(&self) -> usize {
         3

crates/stark-backend/tests/partitioned_sum_air/air.rs

@@ -6,7 +6,7 @@
 use openvm_stark_backend::{
     air_builders::PartitionedAirBuilder,
     p3_field::FieldAlgebra,
-    rap::{BaseAirWithPublicValues, PartitionedBaseAir},
+    rap::{BaseAirWithPublicValues, ColumnsAir, PartitionedBaseAir},
 };
 use p3_air::{Air, BaseAir};
 use p3_matrix::Matrix;
@@ -29,6 +29,8 @@ impl<F> BaseAir<F> for SumAir {
     }
 }
 
+impl<F> ColumnsAir<F> for SumAir {}
+
 impl<AB: PartitionedAirBuilder> Air<AB> for SumAir {
     fn eval(&self, builder: &mut AB) {
         assert_eq!(builder.cached_mains().len(), 1);

crates/stark-sdk/src/dummy_airs/fib_air/air.rs

@@ -3,7 +3,7 @@ use std::borrow::Borrow;
 use openvm_stark_backend::{
     p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir},
     p3_matrix::Matrix,
-    rap::{BaseAirWithPublicValues, PartitionedBaseAir},
+    rap::{BaseAirWithPublicValues, ColumnsAir, PartitionedBaseAir},
 };
 
 use super::columns::{FibonacciCols, NUM_FIBONACCI_COLS};
@@ -24,6 +24,8 @@ impl<F> BaseAirWithPublicValues<F> for FibonacciAir {
     }
 }
 
+impl<F> ColumnsAir<F> for FibonacciAir {}
+
 impl<AB: AirBuilderWithPublicValues> Air<AB> for FibonacciAir {
     fn eval(&self, builder: &mut AB) {
         let main = builder.main();

crates/stark-sdk/src/dummy_airs/interaction/dummy_interaction_air.rs

@@ -20,7 +20,7 @@ use openvm_stark_backend::{
         hal::TraceCommitter,
         types::{AirProofInput, AirProofRawInput, CommittedTraceData},
     },
-    rap::{AnyRap, BaseAirWithPublicValues, PartitionedBaseAir},
+    rap::{AnyRap, BaseAirWithPublicValues, ColumnsAir, PartitionedBaseAir},
     Chip, ChipUsageGetter,
 };
 
@@ -45,6 +45,8 @@ pub struct DummyInteractionAir {
     pub partition: bool,
 }
 
+impl<F: Field> ColumnsAir<F> for DummyInteractionAir {}
+
 impl DummyInteractionAir {
     pub fn new(field_width: usize, is_send: bool, bus_index: BusIndex) -> Self {
         Self {