Feat: support multi_observe (#12)

* fix1

* fix2

* execution wip

* fix 3

* fix 4

Author: kunxian-xia

extensions/native/circuit/src/extension/mod.rs

@@ -165,6 +165,7 @@ impl<F: PrimeField32> VmExecutionExtension<F> for Native {
                 VerifyBatchOpcode::VERIFY_BATCH.global_opcode(),
                 Poseidon2Opcode::PERM_POS2.global_opcode(),
                 Poseidon2Opcode::COMP_POS2.global_opcode(),
+                Poseidon2Opcode::MULTI_OBSERVE.global_opcode(),
             ],
         )?;
 

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

@@ -8,7 +8,7 @@ use openvm_circuit_primitives::utils::not;
 use openvm_instructions::LocalOpcode;
 use openvm_native_compiler::{
     conversion::AS,
-    Poseidon2Opcode::{COMP_POS2, PERM_POS2, MULTI_OBSERVE},
+    Poseidon2Opcode::{COMP_POS2, MULTI_OBSERVE, PERM_POS2},
     VerifyBatchOpcode::VERIFY_BATCH,
 };
 use openvm_poseidon2_air::{
@@ -26,8 +26,8 @@ use openvm_stark_backend::{
 use crate::poseidon2::{
     chip::{NUM_INITIAL_READS, NUM_SIMPLE_ACCESSES},
     columns::{
-        InsideRowSpecificCols, NativePoseidon2Cols, SimplePoseidonSpecificCols,
-        TopLevelSpecificCols, MultiObserveCols,
+        InsideRowSpecificCols, MultiObserveCols, NativePoseidon2Cols, SimplePoseidonSpecificCols,
+        TopLevelSpecificCols,
     },
 };
 
@@ -118,7 +118,8 @@ impl<AB: InteractionBuilder, const SBOX_REGISTERS: usize> Air<AB>
         builder.assert_bool(inside_row);
         builder.assert_bool(simple);
         builder.assert_bool(multi_observe_row);
-        let enabled = incorporate_row + incorporate_sibling + inside_row + simple + multi_observe_row;
+        let enabled =
+            incorporate_row + incorporate_sibling + inside_row + simple + multi_observe_row;
         builder.assert_bool(enabled.clone());
         builder.assert_bool(end_inside_row);
         builder.when(end_inside_row).assert_one(inside_row);
@@ -730,18 +731,12 @@ impl<AB: InteractionBuilder, const SBOX_REGISTERS: usize> Air<AB>
             input_register_1,
             input_register_2,
             input_register_3,
-            output_register
+            output_register,
         } = multi_observe_specific;
 
-        builder
-            .when(multi_observe_row)
-            .assert_bool(is_first);
-        builder
-            .when(multi_observe_row)
-            .assert_bool(is_last);
-        builder
-            .when(multi_observe_row)
-            .assert_bool(should_permute);
+        builder.when(multi_observe_row).assert_bool(is_first);
+        builder.when(multi_observe_row).assert_bool(is_last);
+        builder.when(multi_observe_row).assert_bool(should_permute);
 
         self.execution_bridge
             .execute_and_increment_pc(
@@ -799,19 +794,19 @@ impl<AB: InteractionBuilder, const SBOX_REGISTERS: usize> Air<AB>
             let i_var = AB::F::from_canonical_usize(i);
             self.memory_bridge
                 .read(
-                    MemoryAddress::new(self.address_space, input_ptr + curr_len + i_var - start_idx),
+                    MemoryAddress::new(
+                        self.address_space,
+                        input_ptr + curr_len + i_var - start_idx,
+                    ),
                     [data[i]],
                     start_timestamp + i_var * AB::F::TWO - start_idx * AB::F::TWO,
-                    &read_data[i]
+                    &read_data[i],
                 )
                 .eval(builder, aux_after_start[i] * aux_before_end[i]);
-                
+
             self.memory_bridge
                 .write(
-                    MemoryAddress::new(
-                        self.address_space,
-                        state_ptr + i_var,
-                    ),
+                    MemoryAddress::new(self.address_space, state_ptr + i_var),
                     [data[i]],
                     start_timestamp + i_var * AB::F::TWO - start_idx * AB::F::TWO + AB::F::ONE,
                     &write_data[i],
@@ -835,59 +830,56 @@ impl<AB: InteractionBuilder, const SBOX_REGISTERS: usize> Air<AB>
             .when(multi_observe_row)
             .when(not(is_first))
             .assert_eq(
-                aux_after_start[0] 
-                + aux_after_start[1]
-                + aux_after_start[2]
-                + aux_after_start[3]
-                + aux_after_start[4]
-                + aux_after_start[5]
-                + aux_after_start[6]
-                + aux_after_start[7],
-                AB::Expr::from_canonical_usize(CHUNK) - start_idx.into()
+                aux_after_start[0]
+                    + aux_after_start[1]
+                    + aux_after_start[2]
+                    + aux_after_start[3]
+                    + aux_after_start[4]
+                    + aux_after_start[5]
+                    + aux_after_start[6]
+                    + aux_after_start[7],
+                AB::Expr::from_canonical_usize(CHUNK) - start_idx.into(),
             );
 
         builder
             .when(multi_observe_row)
             .when(not(is_first))
             .assert_eq(
                 aux_before_end[0]
-                + aux_before_end[1]
-                + aux_before_end[2]
-                + aux_before_end[3]
-                + aux_before_end[4]
-                + aux_before_end[5]
-                + aux_before_end[6]
-                + aux_before_end[7], 
-                end_idx
+                    + aux_before_end[1]
+                    + aux_before_end[2]
+                    + aux_before_end[3]
+                    + aux_before_end[4]
+                    + aux_before_end[5]
+                    + aux_before_end[6]
+                    + aux_before_end[7],
+                end_idx,
             );
 
-        let full_sponge_input = from_fn::<_, {CHUNK * 2}, _>(|i| local.inner.inputs[i]);
-        let full_sponge_output = from_fn::<_, {CHUNK * 2}, _>(|i| local.inner.ending_full_rounds[BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS - 1].post[i]);
+        let full_sponge_input = from_fn::<_, { CHUNK * 2 }, _>(|i| local.inner.inputs[i]);
+        let full_sponge_output = from_fn::<_, { CHUNK * 2 }, _>(|i| {
+            local.inner.ending_full_rounds[BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS - 1].post[i]
+        });
 
         self.memory_bridge
             .read(
-                MemoryAddress::new(
-                    self.address_space,
-                    state_ptr,
-                ),
+                MemoryAddress::new(self.address_space, state_ptr),
                 full_sponge_input,
-                start_timestamp + end_idx * AB::F::TWO - start_idx * AB::F::TWO,
+                start_timestamp + (end_idx - start_idx) * AB::F::TWO,
                 &read_sponge_state,
             )
             .eval(builder, multi_observe_row * should_permute);
-        
+
         self.memory_bridge
             .write(
-                MemoryAddress::new(
-                    self.address_space,
-                    state_ptr
-                ),
+                MemoryAddress::new(self.address_space, state_ptr),
                 full_sponge_output,
-                start_timestamp + end_idx * AB::F::TWO - start_idx * AB::F::TWO + AB::F::ONE,
+                start_timestamp + (end_idx - start_idx) * AB::F::TWO + AB::F::ONE,
                 &write_sponge_state,
             )
             .eval(builder, multi_observe_row * should_permute);
 
+        /*
         self.memory_bridge
             .write(
                 MemoryAddress::new(
@@ -899,13 +891,14 @@ impl<AB: InteractionBuilder, const SBOX_REGISTERS: usize> Air<AB>
                 &write_final_idx
             )
             .eval(builder, multi_observe_row * is_last);
+        */
 
         // Field transitions
         builder
             .when(next.multi_observe_row)
             .when(not(next_multi_observe_specific.is_first))
             .assert_eq(next_multi_observe_specific.curr_len, multi_observe_specific.curr_len + end_idx - start_idx);
-        
+
         // Boundary conditions
         builder
             .when(multi_observe_row)
@@ -927,7 +920,7 @@ impl<AB: InteractionBuilder, const SBOX_REGISTERS: usize> Air<AB>
             .when(not(is_last))
             .assert_one(next.multi_observe_row);
 
-        // Field consistency
+        // Fields remain same across same instance
         builder
             .when(next.multi_observe_row)
             .when(not(next_multi_observe_specific.is_first))
@@ -951,17 +944,26 @@ impl<AB: InteractionBuilder, const SBOX_REGISTERS: usize> Air<AB>
         builder
             .when(next.multi_observe_row)
             .when(not(next_multi_observe_specific.is_first))
-            .assert_eq(input_register_1, next_multi_observe_specific.input_register_1);
+            .assert_eq(
+                input_register_1,
+                next_multi_observe_specific.input_register_1,
+            );
 
         builder
             .when(next.multi_observe_row)
             .when(not(next_multi_observe_specific.is_first))
-            .assert_eq(input_register_2, next_multi_observe_specific.input_register_2);
+            .assert_eq(
+                input_register_2,
+                next_multi_observe_specific.input_register_2,
+            );
 
         builder
             .when(next.multi_observe_row)
             .when(not(next_multi_observe_specific.is_first))
-            .assert_eq(input_register_3, next_multi_observe_specific.input_register_3);
+            .assert_eq(
+                input_register_3,
+                next_multi_observe_specific.input_register_3,
+            );
 
         builder
             .when(next.multi_observe_row)
@@ -974,10 +976,12 @@ impl<AB: InteractionBuilder, const SBOX_REGISTERS: usize> Air<AB>
             .when(not(next_multi_observe_specific.is_first))
             .assert_eq(very_first_timestamp, next.very_first_timestamp);
 
+        /*
         builder
             .when(next.multi_observe_row)
             .when(not(next_multi_observe_specific.is_first))
             .assert_eq(next.start_timestamp, start_timestamp + is_first * AB::F::from_canonical_usize(4) + (end_idx - start_idx) * AB::F::TWO + should_permute * AB::F::TWO);
+         */
     }
 }