Add integration test for interactions

stark-middleware/tests/fib_selector_air/air.rs

@@ -2,17 +2,27 @@ 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>,
 }
 
-// No interactions
-impl<F: Field> Chip<F> for FibonacciSelectorAir {}
+impl<F: Field> Chip<F> for FibonacciSelectorAir {
+    fn receives(&self) -> Vec<Interaction<F>> {
+        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,
+        }]
+    }
+}
 
 impl<F: Field> BaseAir<F> for FibonacciSelectorAir {
     fn width(&self) -> usize {

stark-middleware/tests/integration_test.rs

@@ -3,9 +3,11 @@ use afs_middleware::{
     setup::PartitionSetup,
     verifier::PartitionVerifier,
 };
+use itertools::Itertools;
 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;
@@ -18,6 +20,7 @@ use crate::config::poseidon2::StarkConfigPoseidon2;
 mod config;
 mod fib_air;
 mod fib_selector_air;
+mod sender_air;
 
 #[test]
 fn test_single_fib_stark() {
@@ -126,6 +129,82 @@ fn test_single_fib_selector_stark() {
         .expect("Verification failed");
 }
 
+#[test]
+fn test_interaction_stark() {
+    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;
+
+    type Val = BabyBear;
+    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(BabyBear::from_canonical_u32);
+
+    let air = FibonacciSelectorAir { sels: sels.clone() };
+
+    let sender_air = sender_air::SenderAir {};
+
+    let prep_trace = air.preprocessed_trace();
+    let setup = PartitionSetup::new(&config);
+    let (pk, vk) = setup.setup(vec![prep_trace]);
+
+    let trace = generate_trace_rows::<Val>(a, b, &air.sels);
+    let trace_committer = TraceCommitter::<StarkConfigPoseidon2>::new(config.pcs());
+    let proven_trace = trace_committer.commit(vec![trace]);
+    let proven = ProvenMultiMatrixAirTrace {
+        trace_data: &proven_trace,
+        airs: vec![&air],
+    };
+
+    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));
+    }
+
+    let sender_trace = RowMajorMatrix::new(vals, 2);
+    let sender_proven = ProvenMultiMatrixAirTrace {
+        trace_data: &trace_committer.commit(vec![sender_trace]),
+        airs: vec![&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, sender_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)
+        .expect("Verification failed");
+}
+
 fn get_fib_number(n: usize) -> u32 {
     let mut a = 0;
     let mut b = 1;

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) {}
+}