Fix overspilling of MULTI_OBSERVE constraints due to indicator insufficiency (#7)

Author: darth-cy

extensions/native/circuit/src/poseidon2/air.rs

@@ -721,6 +721,7 @@ impl<AB: InteractionBuilder, const SBOX_REGISTERS: usize> Air<AB>
             end_idx,
             aux_after_start,
             aux_before_end,
+            aux_read_enabled,
             read_data,
             write_data,
             data,
@@ -802,7 +803,7 @@ impl<AB: InteractionBuilder, const SBOX_REGISTERS: usize> Air<AB>
                     start_timestamp + i_var * AB::F::TWO - start_idx * AB::F::TWO,
                     &read_data[i],
                 )
-                .eval(builder, aux_after_start[i] * aux_before_end[i]);
+                .eval(builder, multi_observe_row * aux_read_enabled[i]);
 
             self.memory_bridge
                 .write(
@@ -811,21 +812,29 @@ impl<AB: InteractionBuilder, const SBOX_REGISTERS: usize> Air<AB>
                     start_timestamp + i_var * AB::F::TWO - start_idx * AB::F::TWO + AB::F::ONE,
                     &write_data[i],
                 )
-                .eval(builder, aux_after_start[i] * aux_before_end[i]);
+                .eval(builder, multi_observe_row * aux_read_enabled[i]);
         }
 
         for i in 0..(CHUNK - 1) {
             builder
+                .when(multi_observe_row)
                 .when(aux_after_start[i])
                 .assert_one(aux_after_start[i + 1]);
         }
 
         for i in 1..CHUNK {
             builder
+                .when(multi_observe_row)
                 .when(aux_before_end[i])
                 .assert_one(aux_before_end[i - 1]);
         }
 
+        for i in 0..CHUNK {
+            builder
+                .when(multi_observe_row)
+                .assert_eq(aux_after_start[i] * aux_before_end[i], aux_read_enabled[i]);
+        }
+
         builder
             .when(multi_observe_row)
             .when(not(is_first))

extensions/native/circuit/src/poseidon2/chip.rs

@@ -781,6 +781,7 @@ where
                         input_ptr_u32 + input_idx as u32,
                         multi_observe_cols.read_data[j].as_mut(),
                     );
+                    multi_observe_cols.aux_read_enabled[j] = F::ONE;
                     tracing_write_native_inplace(
                         state.memory,
                         state_ptr_u32 + j as u32,

extensions/native/circuit/src/poseidon2/columns.rs

@@ -229,6 +229,7 @@ pub struct MultiObserveCols<T> {
     pub end_idx: T,
     pub aux_after_start: [T; CHUNK],
     pub aux_before_end: [T; CHUNK],
+    pub aux_read_enabled: [T; CHUNK],
 
     // Transcript observation
     pub read_data: [MemoryReadAuxCols<T>; CHUNK],

extensions/native/recursion/tests/recursion.rs

@@ -16,7 +16,8 @@ use openvm_native_compiler::{
     prelude::Usize,
 };
 use openvm_native_recursion::{
-    challenger::duplex::DuplexChallengerVariable, testing_utils::inner::run_recursive_test,
+    challenger::{duplex::DuplexChallengerVariable, CanObserveVariable},
+    testing_utils::inner::run_recursive_test,
 };
 use openvm_stark_backend::{
     config::{Domain, StarkGenericConfig, Val},
@@ -218,7 +219,14 @@ fn build_test_program<C: Config>(builder: &mut Builder<C>) {
     let sample_lens: Vec<usize> = vec![10, 2, 1, 3, 20];
 
     let mut rng = create_seeded_rng();
-    let challenger = DuplexChallengerVariable::new(builder);
+    let mut challenger = DuplexChallengerVariable::new(builder);
+
+    // Observe a setup label
+    let label_f: Vec<u64> = vec![128, 3098, 192, 394, 1662, 928, 374, 281, 598, 182, 475, 729];
+    for n in label_f {
+        let f: Felt<C::F> = builder.constant(C::F::from_canonical_u64(n));
+        challenger.observe(builder, f);
+    }
 
     for l in sample_lens {
         let sample_input: Array<C, Felt<C::F>> = builder.dyn_array(l);