instruction: fix poseidon instructions

tcoratger · tcoratger · commit 9c6308f0e2f7 · 2025-08-17T22:50:53.000+02:00
diff --git a/crates/leanVm/src/bytecode/instruction/poseidon16.rs b/crates/leanVm/src/bytecode/instruction/poseidon16.rs
@@ -1,6 +1,8 @@
+use p3_field::BasedVectorSpace;
 use p3_symmetric::Permutation;
 
 use crate::{
+    bytecode::operand::{MemOrConstant, MemOrFp},
     constant::{DIMENSION, F},
     context::run_context::RunContext,
     errors::vm::VirtualMachineError,
@@ -11,21 +13,21 @@ use crate::{
 /// Poseidon2 permutation over 16 field elements (2 inputs, 2 outputs).
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub struct Poseidon2_16Instruction {
-    /// The starting offset `s` from `fp`. The instruction reads 4 pointers from `m[fp+s]` to `m[fp+s+3]`.
-    pub shift: usize,
+    /// A pointer to the first 8-element input vector.
+    pub arg_a: MemOrConstant,
+    /// A pointer to the second 8-element input vector.
+    pub arg_b: MemOrConstant,
+    /// A pointer to the location for the two 8-element output vectors.
+    pub res: MemOrFp,
 }
 
 impl Poseidon2_16Instruction {
     /// Executes the `Poseidon2_16` precompile instruction.
     ///
-    /// Reads four pointers from memory starting at `fp + shift`, representing:
-    /// - two input vector addresses (`ptr_arg_0`, `ptr_arg_1`)
-    /// - two output vector addresses (`ptr_res_0`, `ptr_res_1`)
-    ///
-    /// Each input is an 8-element vector of `F`. The two inputs are concatenated,
-    /// permuted using Poseidon2, and written back to the output locations.
-    ///
-    /// The operation is: `Poseidon2(m_vec[ptr_0], m_vec[ptr_1]) -> (m_vec[ptr_2], m_vec[ptr_3])`
+    /// This function resolves pointers from its operands to find the memory locations for
+    /// two 8-element input vectors and two 8-element output vectors. It reads the inputs,
+    /// concatenates them, applies the permutation, and writes the two resulting vectors
+    /// back to their designated output locations.
     pub fn execute<Perm>(
         &self,
         run_context: &RunContext,
@@ -35,42 +37,50 @@ impl Poseidon2_16Instruction {
     where
         Perm: Permutation<[F; 2 * DIMENSION]>,
     {
-        // Read Pointers from Memory
+        // Pointer Resolution
         //
-        // The instruction specifies 4 consecutive pointers starting at `fp + shift`.
-        let base_ptr_addr = (run_context.fp + self.shift)?;
-        let ptrs: [MemoryAddress; 4] = memory_manager.memory.get_array_as(base_ptr_addr)?;
-
-        // Convert the `MemoryValue` pointers to `MemoryAddress`.
-        let [ptr_arg_0, ptr_arg_1, ptr_res_0, ptr_res_1] = ptrs;
+        // Resolve the pointer to the first input vector from the `arg_a` operand.
+        let ptr_arg_a: MemoryAddress = run_context
+            .value_from_mem_or_constant(&self.arg_a, memory_manager)?
+            .try_into()?;
+        // Resolve the pointer to the second input vector from the `arg_b` operand.
+        let ptr_arg_b: MemoryAddress = run_context
+            .value_from_mem_or_constant(&self.arg_b, memory_manager)?
+            .try_into()?;
+        // Resolve the pointer for the start of the output block from the `res` operand.
+        let ptr_res: MemoryAddress = run_context
+            .value_from_mem_or_fp(&self.res, memory_manager)?
+            .try_into()?;
+        // The second output vector will be stored immediately after the first one.
+        let ptr_res_b = (ptr_res + 1)?;
 
-        // Read Input Vectors
+        // Data Reading
         //
-        // Read the 8-element vectors from the locations pointed to by `ptr_arg_0` and `ptr_arg_1`.
+        // Read the first 8-element input vector from memory.
         let arg0 = memory_manager
             .memory
-            .get_array_as::<F, DIMENSION>(ptr_arg_0)?;
+            .get_array_as::<F, DIMENSION>(ptr_arg_a)?;
+        // Read the second 8-element input vector from memory.
         let arg1 = memory_manager
             .memory
-            .get_array_as::<F, DIMENSION>(ptr_arg_1)?;
+            .get_array_as::<F, DIMENSION>(ptr_arg_b)?;
 
-        // Perform Hashing
+        // Hashing
         //
         // Concatenate the two input vectors into a single 16-element array for the permutation.
         let mut state = [arg0, arg1].concat().try_into().unwrap();
-
         // Apply the Poseidon2 permutation to the state.
         perm.permute_mut(&mut state);
 
-        // Write Output Vectors
+        // Data Writing
         //
         // Split the permuted state back into two 8-element output vectors.
         let res0: [F; DIMENSION] = state[..DIMENSION].try_into().unwrap();
         let res1: [F; DIMENSION] = state[DIMENSION..].try_into().unwrap();
-
-        // Write the output vectors to the memory locations pointed to by `ptr_res_0` and `ptr_res_1`.
-        memory_manager.load_data(ptr_res_0, &res0)?;
-        memory_manager.load_data(ptr_res_1, &res1)?;
+        // Write the first output vector to its memory location.
+        memory_manager.load_data(ptr_res, &res0)?;
+        // Write the second output vector to its memory location.
+        memory_manager.load_data(ptr_res_b, &res1)?;
 
         Ok(())
     }
@@ -82,36 +92,51 @@ impl Poseidon2_16Instruction {
         run_context: &RunContext,
         memory_manager: &MemoryManager,
     ) -> Result<WitnessPoseidon16, VirtualMachineError> {
-        // Read the four pointers (input_a, input_b, output_a, output_b) from memory.
-        let base_ptr_addr = (run_context.fp + self.shift)?;
-        let [addr_input_a, addr_input_b, addr_output_a, addr_output_b]: [MemoryAddress; 4] =
-            memory_manager.memory.get_array_as(base_ptr_addr)?;
+        // Pointer Resolution
+        //
+        // Resolve the pointer to the first input vector from the `arg_a` operand.
+        let addr_input_a: MemoryAddress = run_context
+            .value_from_mem_or_constant(&self.arg_a, memory_manager)?
+            .try_into()?;
+        // Resolve the pointer to the second input vector from the `arg_b` operand.
+        let addr_input_b: MemoryAddress = run_context
+            .value_from_mem_or_constant(&self.arg_b, memory_manager)?
+            .try_into()?;
+        // Resolve the pointer for the start of the output block from the `res` operand.
+        let addr_output: MemoryAddress = run_context
+            .value_from_mem_or_fp(&self.res, memory_manager)?
+            .try_into()?;
 
-        // Read the two 8-element input vectors from their respective addresses.
+        // Data Reading
+        //
+        // Read the first 8-element input vector from its respective address.
         let value_a = memory_manager
             .memory
             .get_array_as::<F, DIMENSION>(addr_input_a)?;
+        // Read the second 8-element input vector from its respective address.
         let value_b = memory_manager
             .memory
             .get_array_as::<F, DIMENSION>(addr_input_b)?;
-
-        // Read the two 8-element output vectors.
-        let output_a = memory_manager
+        // Read the full 16-element output from memory, starting at the output address.
+        let output = memory_manager
             .memory
-            .get_array_as::<F, DIMENSION>(addr_output_a)?;
-        let output_b = memory_manager
-            .memory
-            .get_array_as::<F, DIMENSION>(addr_output_b)?;
+            .get_vectorized_slice_extension(addr_output, 2)?;
+        let output_coeffs: Vec<F> = output
+            .iter()
+            .flat_map(BasedVectorSpace::as_basis_coefficients_slice)
+            .copied()
+            .collect();
 
+        // Witness Construction
+        //
         // Construct and return the witness struct.
         Ok(WitnessPoseidon16 {
             cycle: Some(cycle),
             addr_input_a,
             addr_input_b,
-            // The output address is the start of the first output vector.
-            addr_output: addr_output_a,
+            addr_output,
             input: [value_a, value_b].concat().try_into().unwrap(),
-            output: [output_a, output_b].concat().try_into().unwrap(),
+            output: output_coeffs.try_into().unwrap(),
         })
     }
 }
diff --git a/crates/leanVm/src/bytecode/instruction/poseidon24.rs b/crates/leanVm/src/bytecode/instruction/poseidon24.rs
@@ -1,32 +1,33 @@
+use p3_field::BasedVectorSpace;
 use p3_symmetric::Permutation;
 
 use crate::{
+    bytecode::operand::{MemOrConstant, MemOrFp},
     constant::{DIMENSION, F},
     context::run_context::RunContext,
     errors::vm::VirtualMachineError,
     memory::{address::MemoryAddress, manager::MemoryManager},
+    witness::poseidon::WitnessPoseidon24,
 };
 
 /// Poseidon2 permutation over 24 field elements (3 inputs, 3 outputs).
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub struct Poseidon2_24Instruction {
-    /// The starting offset from `fp`. The instruction reads 6 pointers from `m[fp+shift]` to `m[fp+shift+5]`.
-    pub shift: usize,
+    /// A pointer to the first two 8-element input vectors.
+    pub arg_a: MemOrConstant,
+    /// A pointer to the third 8-element input vector.
+    pub arg_b: MemOrConstant,
+    /// A pointer to the location of the third 8-element output vector.
+    pub res: MemOrFp,
 }
 
 impl Poseidon2_24Instruction {
     /// Executes the `Poseidon2_24` precompile instruction.
     ///
-    /// Reads six pointers from memory starting at `fp + shift`, representing:
-    /// - three input vector addresses (`ptr_arg_0`, `ptr_arg_1`, `ptr_arg_2`)
-    /// - three output vector addresses (`ptr_res_0`, `ptr_res_1`, `ptr_res_2`)
-    ///
-    /// Each input is an 8-element vector of `F`. The three inputs are concatenated into
-    /// a single 24-element state, permuted using Poseidon2, and written back to the three
-    /// output locations as three 8-element vectors.
-    ///
-    /// The operation is:
-    /// `Poseidon2(m_vec[ptr_0], m_vec[ptr_1], m_vec[ptr_2]) -> (m_vec[ptr_3], m_vec[ptr_4], m_vec[ptr_5])`
+    /// This function resolves pointers from its operands to find the memory locations for
+    /// three 8-element input vectors and one 8-element output vector. It reads the inputs,
+    /// concatenates them, applies the permutation, and writes the third resulting vector
+    /// back to its designated output location.
     pub fn execute<Perm>(
         &self,
         run_context: &RunContext,
@@ -36,55 +37,109 @@ impl Poseidon2_24Instruction {
     where
         Perm: Permutation<[F; 3 * DIMENSION]>,
     {
-        // Read Pointers from Memory
+        // Pointer Resolution
         //
-        // The instruction specifies 6 consecutive pointers starting at `fp + shift`.
-        let ptr_addr = (run_context.fp + self.shift)?;
-        let ptrs: [MemoryAddress; 6] = memory_manager.memory.get_array_as(ptr_addr)?;
+        // Resolve the pointer to the first block of two input vectors from the `arg_a` operand.
+        let ptr_arg_a: MemoryAddress = run_context
+            .value_from_mem_or_constant(&self.arg_a, memory_manager)?
+            .try_into()?;
+        // The second input vector is located immediately after the first one.
+        let ptr_arg_b = (ptr_arg_a + 1)?;
+        // Resolve the pointer to the third input vector from the `arg_b` operand.
+        let ptr_arg_c: MemoryAddress = run_context
+            .value_from_mem_or_constant(&self.arg_b, memory_manager)?
+            .try_into()?;
+        // Resolve the pointer to the third output vector from the `res` operand.
+        let ptr_res: MemoryAddress = run_context
+            .value_from_mem_or_fp(&self.res, memory_manager)?
+            .try_into()?;
 
-        // Convert the raw memory values into memory addresses.
-        let [
-            ptr_arg_0,
-            ptr_arg_1,
-            ptr_arg_2,
-            ptr_res_0,
-            ptr_res_1,
-            ptr_res_2,
-        ] = ptrs;
-
-        // Read Input Vectors
+        // Data Reading
         //
-        // Each is an 8-element array of field elements.
+        // Read the first 8-element input vector from memory.
         let arg0 = memory_manager
             .memory
-            .get_array_as::<F, DIMENSION>(ptr_arg_0)?;
+            .get_array_as::<F, DIMENSION>(ptr_arg_a)?;
+        // Read the second 8-element input vector from memory.
         let arg1 = memory_manager
             .memory
-            .get_array_as::<F, DIMENSION>(ptr_arg_1)?;
+            .get_array_as::<F, DIMENSION>(ptr_arg_b)?;
+        // Read the third 8-element input vector from memory.
         let arg2 = memory_manager
             .memory
-            .get_array_as::<F, DIMENSION>(ptr_arg_2)?;
+            .get_array_as::<F, DIMENSION>(ptr_arg_c)?;
 
-        // Perform Hashing
+        // Hashing
         //
         // Concatenate the three input vectors into a single 24-element array for the permutation.
         let mut state = [arg0, arg1, arg2].concat().try_into().unwrap();
-
         // Apply the Poseidon2 permutation to the state.
         perm.permute_mut(&mut state);
 
-        // Write Output Vectors
+        // Data Writing
         //
-        // Split the permuted state back into three 8-element output vectors.
-        let res0: [F; DIMENSION] = state[..DIMENSION].try_into().unwrap();
-        let res1: [F; DIMENSION] = state[DIMENSION..2 * DIMENSION].try_into().unwrap();
-        let res2: [F; DIMENSION] = state[2 * DIMENSION..].try_into().unwrap();
-
-        // Write the output vectors to the memory locations pointed to by the result pointers.
-        memory_manager.load_data(ptr_res_0, &res0)?;
-        memory_manager.load_data(ptr_res_1, &res1)?;
-        memory_manager.load_data(ptr_res_2, &res2)?;
+        // Extract the last 8 elements of the permuted state, which is the result.
+        let res: [F; DIMENSION] = state[2 * DIMENSION..].try_into().unwrap();
+        // Write the result vector to the memory location pointed to by the result pointer.
+        memory_manager.load_data(ptr_res, &res)?;
 
         Ok(())
     }
+
+    /// Generates the witness for a `Poseidon2_24` instruction execution.
+    pub fn generate_witness(
+        &self,
+        cycle: usize,
+        run_context: &RunContext,
+        memory_manager: &MemoryManager,
+    ) -> Result<WitnessPoseidon24, VirtualMachineError> {
+        // Pointer Resolution
+        //
+        // Resolve the pointer to the first block of two input vectors from the `arg_a` operand.
+        let addr_input_a: MemoryAddress = run_context
+            .value_from_mem_or_constant(&self.arg_a, memory_manager)?
+            .try_into()?;
+        // Resolve the pointer to the third input vector from the `arg_b` operand.
+        let addr_input_b: MemoryAddress = run_context
+            .value_from_mem_or_constant(&self.arg_b, memory_manager)?
+            .try_into()?;
+        // Resolve the pointer to the third output vector from the `res` operand.
+        let addr_output: MemoryAddress = run_context
+            .value_from_mem_or_fp(&self.res, memory_manager)?
+            .try_into()?;
+
+        // Data Reading
+        //
+        // Read the first two input vectors (a slice of 2 EF elements, total 16 F elements) from memory.
+        let value_a = memory_manager
+            .memory
+            .get_vectorized_slice_extension(addr_input_a, 2)?;
+        // Read the third input vector (a single EF element, total 8 F elements) from memory.
+        let value_b = memory_manager.memory.get_extension(addr_input_b)?;
+        // Read the third output vector (a single EF element, total 8 F elements) from memory.
+        let output = memory_manager.memory.get_extension(addr_output)?;
+
+        // Witness Construction
+        //
+        // Get the F coefficients from the first two EF elements (value_a).
+        let input_part1: Vec<F> = value_a
+            .iter()
+            .flat_map(BasedVectorSpace::as_basis_coefficients_slice)
+            .copied()
+            .collect();
+        // Get the F coefficients from the third EF element (value_b).
+        let input_part2: Vec<F> = value_b.as_basis_coefficients_slice().to_vec();
+        // Concatenate the parts to form the full 24-element input for the witness.
+        let full_input = [input_part1, input_part2].concat().try_into().unwrap();
+
+        // Construct and return the final witness struct.
+        Ok(WitnessPoseidon24 {
+            cycle: Some(cycle),
+            addr_input_a,
+            addr_input_b,
+            addr_output,
+            input: full_input,
+            output: output.as_basis_coefficients_slice().try_into().unwrap(),
+        })
+    }
 }
diff --git a/crates/leanVm/src/witness/poseidon.rs b/crates/leanVm/src/witness/poseidon.rs
@@ -23,11 +23,11 @@ pub struct WitnessPoseidon24 {
     /// The CPU cycle at which this operation is initiated, if applicable.
     pub cycle: Option<usize>,
     /// The memory address (vectorized pointer, of size 2) of the first two 8-element input vectors.
-    pub addr_input_a: usize,
+    pub addr_input_a: MemoryAddress,
     /// The memory address (vectorized pointer, of size 1) of the third 8-element input vector.
-    pub addr_input_b: usize,
+    pub addr_input_b: MemoryAddress,
     /// The memory address (vectorized pointer, of size 1) where the relevant 8-element output vector is stored.
-    pub addr_output: usize,
+    pub addr_output: MemoryAddress,
     /// The full 24-element input state for the permutation.
     pub input: [F; 24],
     /// The last 8 elements of the 24-element output state from the permutation.
