[Feat] Native Multiple Observe for Poseidon2-based Challenger

extensions/native/circuit/src/extension.rs

@@ -200,6 +200,7 @@ impl<F: PrimeField32> VmExtension<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/chip.rs

@@ -9,7 +9,7 @@ use openvm_circuit::{
 use openvm_instructions::{instruction::Instruction, program::DEFAULT_PC_STEP, LocalOpcode};
 use openvm_native_compiler::{
     conversion::AS,
-    Poseidon2Opcode::{COMP_POS2, PERM_POS2},
+    Poseidon2Opcode::{COMP_POS2, PERM_POS2, MULTI_OBSERVE},
     VerifyBatchOpcode::VERIFY_BATCH,
 };
 use openvm_poseidon2_air::{Poseidon2Config, Poseidon2SubAir, Poseidon2SubChip};
@@ -485,6 +485,43 @@ impl<F: PrimeField32, const SBOX_REGISTERS: usize> InstructionExecutor<F>
                     initial_log_height: initial_log_height as usize,
                     top_level,
                 });
+        } else if instruction.opcode == MULTI_OBSERVE.global_opcode() {
+            let &Instruction {
+                a: output_register,
+                b: input_register_1,
+                c: input_register_2,
+                d: input_register_3,
+                e: register_address_space,
+                f: data_address_space,
+                ..
+            } = instruction;
+
+            let (_, sponge_ptr) = memory.read_cell(data_address_space, output_register);
+            let (_, arr_ptr) = memory.read_cell(data_address_space, input_register_2);
+
+            let init_pos_read = memory.read_cell(register_address_space, input_register_1);
+            let mut pos = init_pos_read.1.as_canonical_u32() as usize;
+
+            let len_read = memory.read_cell(register_address_space, input_register_3);
+            let len = len_read.1.as_canonical_u32() as usize;
+
+            for i in 0..len {
+                let mod_pos = pos % CHUNK;
+                let n_read = memory.read_cell(register_address_space, arr_ptr + F::from_canonical_usize(i));
+                let n_f = n_read.1;
+
+                memory.write_cell(register_address_space, sponge_ptr + F::from_canonical_usize(mod_pos), n_f);
+                pos += 1;
+
+                if pos % CHUNK == 0 {
+                    let (_, sponge_state) = memory.read::<{CHUNK * 2}>(register_address_space, sponge_ptr);
+                    let output = self.subchip.permute(sponge_state);
+                    memory.write::<{CHUNK * 2}>(data_address_space, sponge_ptr, std::array::from_fn(|i| output[i]));
+                }
+            }
+
+            let mod_pos = pos % CHUNK;
+            memory.write_cell(register_address_space, input_register_1, F::from_canonical_usize(mod_pos));
         } else {
             unreachable!()
         }
@@ -501,6 +538,8 @@ impl<F: PrimeField32, const SBOX_REGISTERS: usize> InstructionExecutor<F>
             String::from("PERM_POS2")
         } else if opcode == COMP_POS2.global_opcode().as_usize() {
             String::from("COMP_POS2")
+        } else if opcode == MULTI_OBSERVE.global_opcode().as_usize() {
+            String::from("MULTI_OBSERVE")
         } else {
             unreachable!("unsupported opcode: {}", opcode)
         }

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

@@ -434,7 +434,8 @@ fn tester_with_random_poseidon2_ops(num_ops: usize) -> VmChipTester<BabyBearBlak
             }
             PERM_POS2 => {
                 tester.write(e, lhs, data);
-            }
+            },
+            MULTI_OBSERVE => {}
         }
 
         tester.execute(&mut chip, &instruction);
@@ -449,6 +450,7 @@ fn tester_with_random_poseidon2_ops(num_ops: usize) -> VmChipTester<BabyBearBlak
                 let actual = tester.read::<{ 2 * CHUNK }>(e, dst);
                 assert_eq!(hash, actual);
             }
+            MULTI_OBSERVE => {}
         }
     }
     tester.build().load(chip).finalize()

extensions/native/compiler/src/asm/compiler.rs

@@ -489,6 +489,12 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
                 DslIr::HintBitsF(var, len) => {
                     self.push(AsmInstruction::HintBits(var.fp(), len), debug_info);
                 }
+                DslIr::Poseidon2MultiObserve(dst, init_pos, arr_ptr, len) => {
+                    self.push(
+            AsmInstruction::Poseidon2MultiObserve(dst.fp(), init_pos.fp(), arr_ptr.fp(), len.get_var().fp()),
+                        debug_info,
+                    );
+                },
                 DslIr::Poseidon2PermuteBabyBear(dst, src) => match (dst, src) {
                     (Array::Dyn(dst, _), Array::Dyn(src, _)) => self.push(
                         AsmInstruction::Poseidon2Permute(dst.fp(), src.fp()),

extensions/native/compiler/src/asm/instruction.rs

@@ -108,6 +108,11 @@ pub enum AsmInstruction<F, EF> {
     /// Halt.
     Halt,
 
+    /// Absorbs multiple base elements into a duplex transcript with Poseidon2 permutation
+    /// (sponge_state, init_pos, arr_ptr, len)
+    /// Returns the final index position of hash sponge
+    Poseidon2MultiObserve(i32, i32, i32, i32),
+
     /// Perform a Poseidon2 permutation on state starting at address `lhs`
     /// and store new state at `rhs`.
     /// (a, b) are pointers to (lhs, rhs).
@@ -331,6 +336,9 @@ impl<F: PrimeField32, EF: ExtensionField<F>> AsmInstruction<F, EF> {
             AsmInstruction::Trap => write!(f, "trap"),
             AsmInstruction::Halt => write!(f, "halt"),
             AsmInstruction::HintBits(src, len) => write!(f, "hint_bits ({})fp, {}", src, len),
+            AsmInstruction::Poseidon2MultiObserve(dst, init_pos, arr, len) => {
+                write!(f, "poseidon2_multi_observe ({})fp, ({})fp ({})fp ({})fp", dst, init_pos, arr, len)
+            }
             AsmInstruction::Poseidon2Permute(dst, lhs) => {
                 write!(f, "poseidon2_permute ({})fp, ({})fp", dst, lhs)
             }

extensions/native/compiler/src/conversion/mod.rs

@@ -440,6 +440,18 @@ fn convert_instruction<F: PrimeField32, EF: ExtensionField<F>>(
                     AS::Native,
                     AS::Native,
             )],
+        AsmInstruction::Poseidon2MultiObserve(dst, init, arr, len) => vec![
+            Instruction {
+                opcode: options.opcode_with_offset(Poseidon2Opcode::MULTI_OBSERVE),
+                a: i32_f(dst),
+                b: i32_f(init),
+                c: i32_f(arr),
+                d: i32_f(len),
+                e: AS::Native.to_field(),
+                f: AS::Native.to_field(),
+                g: F::ZERO,
+            }
+        ],
         AsmInstruction::Poseidon2Compress(dst, src1, src2) => vec![inst(
             options.opcode_with_offset(Poseidon2Opcode::COMP_POS2),
             i32_f(dst),

extensions/native/compiler/src/ir/instructions.rs

@@ -196,6 +196,13 @@ pub enum DslIr<C: Config> {
     /// Permutes an array of Bn254 elements using Poseidon2 (output = p2_permute(array)). Should only
     /// be used when target is a circuit.
     CircuitPoseidon2Permute([Var<C::N>; 3]),
+    /// Absorbs an array of baby bear elements into a duplex transcript with Poseidon2 permutations (output = p2_multi_observe(array, els)).
+    Poseidon2MultiObserve(
+        Ptr<C::N>,              // sponge_state
+        Var<C::N>,              // initial input_ptr position
+        Ptr<C::N>,              // input array (els)
+        Usize<C::N>,            // len of els
+    ),
 
     // Miscellaneous instructions.
     /// Prints a variable.

extensions/native/compiler/src/ir/poseidon.rs

@@ -1,13 +1,57 @@
 use openvm_native_compiler_derive::iter_zip;
 use openvm_stark_backend::p3_field::FieldAlgebra;
 
+use crate::ir::Variable;
+
 use super::{Array, ArrayLike, Builder, Config, DslIr, Ext, Felt, MemIndex, Ptr, Usize, Var};
 
 pub const DIGEST_SIZE: usize = 8;
 pub const HASH_RATE: usize = 8;
 pub const PERMUTATION_WIDTH: usize = 16;
 
 impl<C: Config> Builder<C> {
+    /// Extends native VM ability to observe multiple base elements in one opcode operation
+    /// Absorbs elements sequentially at the RATE portion of sponge state and performs as many permutations as necessary.
+    /// Returns the index position of the input_ptr and the final sponge state.
+    /// 
+    /// [Reference](https://docs.rs/p3-poseidon2/latest/p3_poseidon2/struct.Poseidon2.html)
+    pub fn poseidon2_multi_observe(
+        &mut self,
+        sponge_state: &Array<C, Felt<C::F>>, 
+        input_ptr: Ptr<C::N>,
+        io_full_ptr: Ptr<C::N>,
+        arr: &Array<C, Felt<C::F>>,
+    ) -> Usize<C::N> {
+        let buffer_size: Var<C::N> = Var::uninit(self);
+        self.assign(&buffer_size, C::N::from_canonical_usize(HASH_RATE));
+
+        match sponge_state {
+            Array::Fixed(_) => {
+                panic!("Poseidon2 permutation is not allowed on fixed arrays");
+            }
+            Array::Dyn(sponge_ptr, _) => {
+                match arr {
+                    Array::Fixed(_) => {
+                        panic!("Base elements input must be dynamic");
+                    }
+                    Array::Dyn(ptr, len) => {
+                        let init_pos: Var<C::N> = Var::uninit(self);
+                        self.assign(&init_pos, buffer_size - (io_full_ptr.address - input_ptr.address));
+
+                        self.operations.push(DslIr::Poseidon2MultiObserve(
+                            *sponge_ptr,
+                            init_pos,
+                            *ptr,
+                            len.clone(),
+                        ));
+
+                        Usize::Var(init_pos)
+                    }
+                }
+            }
+        }
+    }
+
     /// Applies the Poseidon2 permutation to the given array.
     ///
     /// [Reference](https://docs.rs/p3-poseidon2/latest/p3_poseidon2/struct.Poseidon2.html)

extensions/native/compiler/src/lib.rs

@@ -184,6 +184,7 @@ pub enum NativePhantom {
 pub enum Poseidon2Opcode {
     PERM_POS2,
     COMP_POS2,
+    MULTI_OBSERVE,
 }
 
 /// Opcodes for FRI opening proofs.