instruction: generate witness for multilinear eval (#12)

* instructions: update the precompiles

* instruction: generate witness for multilinear eval

* some fixes

Author: tcoratger

crates/leanVm/src/bytecode/instruction/dot_product.rs

@@ -4,7 +4,7 @@ use crate::{
     bytecode::operand::{MemOrConstant, MemOrFp},
     constant::{DIMENSION, EF, F},
     context::run_context::RunContext,
-    errors::vm::VirtualMachineError,
+    errors::{memory::MemoryError, vm::VirtualMachineError},
     memory::{address::MemoryAddress, manager::MemoryManager},
 };
 
@@ -51,7 +51,7 @@ impl DotProductInstruction {
                 let addr = (ptr_arg_0 + i)?;
                 let vector_coeffs = memory_manager.memory.get_array_as::<F, DIMENSION>(addr)?;
                 EF::from_basis_coefficients_slice(&vector_coeffs)
-                    .ok_or(VirtualMachineError::InvalidExtensionField)
+                    .ok_or(MemoryError::InvalidExtensionFieldConversion)
             })
             .collect::<Result<_, _>>()?;
 
@@ -61,7 +61,7 @@ impl DotProductInstruction {
                 let addr = (ptr_arg_1 + i)?;
                 let vector_coeffs = memory_manager.memory.get_array_as::<F, DIMENSION>(addr)?;
                 EF::from_basis_coefficients_slice(&vector_coeffs)
-                    .ok_or(VirtualMachineError::InvalidExtensionField)
+                    .ok_or(MemoryError::InvalidExtensionFieldConversion)
             })
             .collect::<Result<_, _>>()?;
 

crates/leanVm/src/bytecode/instruction/multilinear_eval.rs

@@ -5,8 +5,9 @@ use crate::{
     bytecode::operand::{MemOrConstant, MemOrFp},
     constant::{DIMENSION, EF, F},
     context::run_context::RunContext,
-    errors::vm::VirtualMachineError,
+    errors::{memory::MemoryError, vm::VirtualMachineError},
     memory::{address::MemoryAddress, manager::MemoryManager},
+    witness::multilinear_eval::WitnessMultilinearEval,
 };
 
 /// An instruction to evaluate a multilinear polynomial at a point in the extension field.
@@ -64,7 +65,7 @@ impl MultilinearEvalInstruction {
                 let addr = (ptr_point + i)?;
                 let vector_coeffs = memory_manager.memory.get_array_as::<F, DIMENSION>(addr)?;
                 EF::from_basis_coefficients_slice(&vector_coeffs)
-                    .ok_or(VirtualMachineError::InvalidExtensionField)
+                    .ok_or(MemoryError::InvalidExtensionFieldConversion)
             })
             .collect::<Result<_, _>>()?;
 
@@ -76,4 +77,46 @@ impl MultilinearEvalInstruction {
 
         Ok(())
     }
+
+    /// Generates the witness for a `MultilinearEval` instruction execution.
+    ///
+    /// This function reads the necessary data from memory (operands and result)
+    /// to construct a `WitnessMultilinearEval` struct, which captures the complete
+    /// state of the operation at a specific cycle.
+    pub fn generate_witness(
+        &self,
+        cycle: usize,
+        run_context: &RunContext,
+        memory_manager: &MemoryManager,
+    ) -> Result<WitnessMultilinearEval, VirtualMachineError> {
+        // Resolve the memory addresses for the coefficients, point, and result.
+        let addr_coeffs = run_context
+            .value_from_mem_or_constant(&self.coeffs, memory_manager)?
+            .try_into()?;
+        let addr_point = run_context
+            .value_from_mem_or_constant(&self.point, memory_manager)?
+            .try_into()?;
+        let addr_res = run_context
+            .value_from_mem_or_fp(&self.res, memory_manager)?
+            .try_into()?;
+
+        // Read the evaluation point from memory using the helper function.
+        let point = memory_manager
+            .memory
+            .get_vectorized_slice_extension(addr_point, self.n_vars)?;
+
+        // Read the result of the evaluation from memory.
+        let res = memory_manager.memory.get_extension(addr_res)?;
+
+        // Construct and return the witness struct.
+        Ok(WitnessMultilinearEval {
+            cycle,
+            addr_coeffs,
+            addr_point,
+            addr_res,
+            n_vars: self.n_vars,
+            point,
+            res,
+        })
+    }
 }

crates/leanVm/src/errors/memory.rs

@@ -52,4 +52,7 @@ pub enum MemoryError {
 
     #[error("Operation failed: {} * {}, can't multiply these two values", 0.0, 0.1)]
     InvalidMul(Box<(MemoryValue, MemoryValue)>),
+
+    #[error("Invalid extensionn field conversion.")]
+    InvalidExtensionFieldConversion,
 }

crates/leanVm/src/memory/mem.rs

@@ -1,7 +1,12 @@
 use std::mem::MaybeUninit;
 
+use p3_field::BasedVectorSpace;
+
 use super::{address::MemoryAddress, cell::MemoryCell, val::MemoryValue};
-use crate::errors::memory::MemoryError;
+use crate::{
+    constant::{DIMENSION, EF, F},
+    errors::memory::MemoryError,
+};
 
 #[derive(Debug, Default, Clone)]
 pub struct Memory {
@@ -202,6 +207,40 @@ impl Memory {
         // SAFETY: All elements have been initialized above in the loop
         Ok(unsafe { std::mem::transmute_copy::<_, [V; DIM]>(&out) })
     }
+
+    /// Retrieves a slice of extension field elements starting from a given vectorized address.
+    ///
+    /// This function reads `len` consecutive vectors of size `DIMENSION` and converts each
+    /// into an extension field element.
+    pub(crate) fn get_vectorized_slice_extension(
+        &self,
+        start_address: MemoryAddress,
+        len: usize,
+    ) -> Result<Vec<EF>, MemoryError> {
+        (0..len)
+            .map(|i| {
+                // Calculate the address for the i-th vector.
+                let current_addr = (start_address + i).map_err(MemoryError::Math)?;
+
+                // Read the DIMENSION base field elements for the i-th vector.
+                let vector_coeffs = self.get_array_as::<F, DIMENSION>(current_addr)?;
+
+                // Convert the base field elements into an extension field element.
+                EF::from_basis_coefficients_slice(&vector_coeffs)
+                    .ok_or(MemoryError::InvalidExtensionFieldConversion)
+            })
+            .collect()
+    }
+
+    /// Retrieves a single extension field element from a vectorized address.
+    pub(crate) fn get_extension(&self, address: MemoryAddress) -> Result<EF, MemoryError> {
+        // Read the DIMENSION base field elements for the vector.
+        let vector_coeffs = self.get_array_as::<F, DIMENSION>(address)?;
+
+        // Convert the base field elements into an extension field element.
+        EF::from_basis_coefficients_slice(&vector_coeffs)
+            .ok_or(MemoryError::InvalidExtensionFieldConversion)
+    }
 }
 
 #[cfg(test)]

crates/leanVm/src/witness/mod.rs

@@ -1,2 +1,3 @@
 pub mod dot_product;
+pub mod multilinear_eval;
 pub mod poseidon;

crates/leanVm/src/witness/multilinear_eval.rs

@@ -0,0 +1,12 @@
+use crate::{constant::EF, memory::address::MemoryAddress};
+
+#[derive(Debug)]
+pub struct WitnessMultilinearEval {
+    pub cycle: usize,
+    pub addr_coeffs: MemoryAddress, // vectorized pointer, of size 8.2^size
+    pub addr_point: MemoryAddress,  // vectorized pointer, of size `size`
+    pub addr_res: MemoryAddress,    // vectorized pointer
+    pub n_vars: usize,
+    pub point: Vec<EF>,
+    pub res: EF,
+}