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

* Fix multi_observe constraints

* Remove debug flags

* fmt

Author: darth-cy

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

@@ -698,6 +698,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,
@@ -783,7 +784,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(
@@ -795,21 +796,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/columns.rs

@@ -224,6 +224,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/circuit/src/poseidon2/trace.rs

@@ -473,6 +473,11 @@ impl<F: PrimeField32, const SBOX_REGISTERS: usize> NativePoseidon2Chip<F, SBOX_R
             for i in 0..record.end_idx {
                 specific.aux_before_end[i] = F::ONE;
             }
+            for i in 0..CHUNK {
+                if specific.aux_after_start[i] + specific.aux_before_end[i] >= F::TWO {
+                    specific.aux_read_enabled[i] = F::ONE;
+                }
+            }
             for i in record.start_idx..record.end_idx {
                 let read_data_record = memory.record_by_id(record.read_input_data[i]);
                 let write_data_record = memory.record_by_id(record.write_input_data[i]);

extensions/native/recursion/tests/recursion.rs

@@ -5,7 +5,7 @@ use openvm_native_compiler::{
     asm::{AsmBuilder, AsmCompiler}, ir::Felt,
     conversion::{convert_program, CompilerOptions}, 
 };
-use openvm_native_recursion::{testing_utils::inner::run_recursive_test, challenger::duplex::DuplexChallengerVariable};
+use openvm_native_recursion::{testing_utils::inner::run_recursive_test, challenger::{duplex::DuplexChallengerVariable, CanObserveVariable}};
 use openvm_stark_backend::{
     config::{Domain, StarkGenericConfig},
     p3_commit::PolynomialSpace,
@@ -148,7 +148,14 @@ fn build_test_program<C: Config>(
     let sample_lens: Vec<usize> = vec![10, 2, 0, 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);