Add integration test for interactions

stark-middleware/src/verifier/error.rs

@@ -1,4 +1,4 @@
-#[derive(Debug)]
+#[derive(Debug, PartialEq, Eq)]
 pub enum VerificationError {
     InvalidProofShape,
     /// An error occurred while verifying the claimed openings.

stark-middleware/tests/fib_selector_air/air.rs

@@ -2,17 +2,32 @@ use std::borrow::Borrow;
 
 use super::columns::FibonacciSelectorCols;
 use crate::fib_air::columns::{FibonacciCols, NUM_FIBONACCI_COLS};
-use afs_middleware::interaction::Chip;
-use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir, PairBuilder};
+use afs_middleware::interaction::{Chip, Interaction};
+use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir, PairBuilder, VirtualPairCol};
 use p3_field::{AbstractField, Field};
 use p3_matrix::{dense::RowMajorMatrix, Matrix};
 
 pub struct FibonacciSelectorAir {
     pub sels: Vec<bool>,
+    pub enable_interactions: bool,
 }
 
-// No interactions
-impl<F: Field> Chip<F> for FibonacciSelectorAir {}
+impl<F: Field> Chip<F> for FibonacciSelectorAir {
+    fn receives(&self) -> Vec<Interaction<F>> {
+        if self.enable_interactions {
+            vec![Interaction::<F> {
+                fields: vec![VirtualPairCol::<F>::new_main(
+                    vec![(0, F::one()), (1, F::one())],
+                    F::zero(),
+                )],
+                count: VirtualPairCol::<F>::single_preprocessed(0),
+                argument_index: 0,
+            }]
+        } else {
+            vec![]
+        }
+    }
+}
 
 impl<F: Field> BaseAir<F> for FibonacciSelectorAir {
     fn width(&self) -> usize {

stark-middleware/tests/integration_test.rs

@@ -18,6 +18,10 @@ use crate::config::poseidon2::StarkConfigPoseidon2;
 mod config;
 mod fib_air;
 mod fib_selector_air;
+mod sender_air;
+
+#[cfg(test)]
+mod interaction;
 
 #[test]
 fn test_single_fib_stark() {
@@ -99,7 +103,10 @@ fn test_single_fib_selector_stark() {
     let sels: Vec<bool> = (0..n).map(|i| i % 2 == 0).collect();
     let pis = [a, b, get_conditional_fib_number(&sels)].map(BabyBear::from_canonical_u32);
 
-    let air = FibonacciSelectorAir { sels };
+    let air = FibonacciSelectorAir {
+        sels,
+        enable_interactions: false,
+    };
 
     let prep_trace = air.preprocessed_trace();
     let setup = PartitionSetup::new(&config);
@@ -154,7 +161,10 @@ fn test_double_fib_starks() {
     let pis = [a, b, get_fib_number(n)].map(BabyBear::from_canonical_u32);
 
     let air1 = FibonacciAir {};
-    let air2 = FibonacciSelectorAir { sels };
+    let air2 = FibonacciSelectorAir {
+        sels,
+        enable_interactions: false,
+    };
 
     let prep_trace1 = air1.preprocessed_trace();
     let prep_trace2 = air2.preprocessed_trace();

stark-middleware/tests/interaction/mod.rs

@@ -0,0 +1,123 @@
+use afs_middleware::{
+    prover::{trace::TraceCommitter, types::ProvenMultiMatrixAirTrace, PartitionProver},
+    setup::PartitionSetup,
+    verifier::{PartitionVerifier, VerificationError},
+};
+use p3_air::BaseAir;
+use p3_baby_bear::BabyBear;
+use p3_field::AbstractField;
+use p3_matrix::dense::RowMajorMatrix;
+use p3_uni_stark::StarkGenericConfig;
+use tracing_forest::util::LevelFilter;
+use tracing_forest::ForestLayer;
+use tracing_subscriber::layer::SubscriberExt;
+use tracing_subscriber::util::SubscriberInitExt;
+use tracing_subscriber::{EnvFilter, Registry};
+
+use crate::{
+    config::{self, poseidon2::StarkConfigPoseidon2},
+    fib_selector_air, get_conditional_fib_number, sender_air,
+};
+
+type Val = BabyBear;
+
+fn verify_fib_interaction(
+    sender_trace_gen: fn(a: u32, b: u32, n: usize, sels: Vec<bool>) -> RowMajorMatrix<Val>,
+) -> Result<(), VerificationError> {
+    use fib_selector_air::air::FibonacciSelectorAir;
+    use fib_selector_air::trace::generate_trace_rows;
+
+    // Set up tracing:
+    let env_filter = EnvFilter::builder()
+        .with_default_directive(LevelFilter::INFO.into())
+        .from_env_lossy();
+    let _ = Registry::default()
+        .with(env_filter)
+        .with(ForestLayer::default())
+        .try_init();
+
+    let log_trace_degree = 3;
+    let perm = config::poseidon2::random_perm();
+    let config = config::poseidon2::default_config(&perm, log_trace_degree);
+
+    // Public inputs:
+    let a = 0u32;
+    let b = 1u32;
+    let n = 1usize << log_trace_degree;
+
+    let sels: Vec<bool> = (0..n).map(|i| i % 2 == 0).collect();
+    let fib_res = get_conditional_fib_number(&sels);
+    let pis = [a, b, fib_res].map(Val::from_canonical_u32);
+
+    let air = FibonacciSelectorAir {
+        sels: sels.clone(),
+        enable_interactions: true,
+    };
+
+    let sender_air = sender_air::SenderAir {};
+
+    let prep_trace = air.preprocessed_trace();
+    let sender_prep_trace = sender_air.preprocessed_trace();
+    let setup = PartitionSetup::new(&config);
+    let (pk, vk) = setup.setup(vec![prep_trace, sender_prep_trace]);
+
+    let trace = generate_trace_rows::<Val>(a, b, &air.sels);
+    let sender_trace = sender_trace_gen(a, b, n, sels);
+
+    let trace_committer = TraceCommitter::<StarkConfigPoseidon2>::new(config.pcs());
+    let proven_trace = trace_committer.commit(vec![trace, sender_trace]);
+    let proven = ProvenMultiMatrixAirTrace {
+        trace_data: &proven_trace,
+        airs: vec![&air, &sender_air],
+    };
+
+    let prover = PartitionProver::new(config);
+    let mut challenger = config::poseidon2::Challenger::new(perm.clone());
+    let proof = prover.prove(&mut challenger, &pk, vec![proven], &pis);
+
+    // Verify the proof:
+    // Start from clean challenger
+    let mut challenger = config::poseidon2::Challenger::new(perm.clone());
+    let verifier = PartitionVerifier::new(prover.config);
+    verifier.verify(&mut challenger, vk, vec![&air, &sender_air], proof, &pis)
+}
+
+#[test]
+fn test_interaction_stark_happy_path() {
+    verify_fib_interaction(|a, b, _n, sels| {
+        let mut curr_a = a;
+        let mut curr_b = b;
+        let mut vals = vec![];
+        for sel in sels {
+            vals.push(Val::from_bool(sel));
+            if sel {
+                let c = curr_a + curr_b;
+                curr_a = curr_b;
+                curr_b = c;
+            }
+            vals.push(Val::from_canonical_u32(curr_b));
+        }
+        RowMajorMatrix::new(vals, 2)
+    })
+    .expect("Verification failed");
+}
+
+#[test]
+fn test_interaction_stark_neg() {
+    let res = verify_fib_interaction(|a, b, _n, sels| {
+        let mut curr_a = a;
+        let mut curr_b = b;
+        let mut vals = vec![];
+        for sel in sels {
+            vals.push(Val::from_bool(sel));
+            if sel {
+                let c = curr_a + curr_b;
+                curr_a = curr_b;
+                curr_b = c;
+            }
+            vals.push(Val::from_canonical_u32(curr_b + 1));
+        }
+        RowMajorMatrix::new(vals, 2)
+    });
+    assert_eq!(res, Err(VerificationError::NonZeroCumulativeSum));
+}

stark-middleware/tests/sender_air/mod.rs

@@ -0,0 +1,30 @@
+use afs_middleware::interaction::{Chip, Interaction};
+use p3_air::{Air, AirBuilderWithPublicValues, BaseAir, PairBuilder, VirtualPairCol};
+use p3_field::Field;
+use p3_matrix::dense::RowMajorMatrix;
+
+pub struct SenderAir {}
+
+impl<F: Field> Chip<F> for SenderAir {
+    fn sends(&self) -> Vec<Interaction<F>> {
+        vec![Interaction::<F> {
+            fields: vec![VirtualPairCol::<F>::single_main(1)],
+            count: VirtualPairCol::<F>::single_main(0),
+            argument_index: 0,
+        }]
+    }
+}
+
+impl<F: Field> BaseAir<F> for SenderAir {
+    fn width(&self) -> usize {
+        2
+    }
+
+    fn preprocessed_trace(&self) -> Option<RowMajorMatrix<F>> {
+        None
+    }
+}
+
+impl<AB: AirBuilderWithPublicValues + PairBuilder> Air<AB> for SenderAir {
+    fn eval(&self, _builder: &mut AB) {}
+}