From c6537a1e2c4cb7ea2c48ecdb7df77e1eca06de00 Mon Sep 17 00:00:00 2001 From: Thomas Coratger Date: Sun, 27 Jul 2025 20:35:13 +0200 Subject: [PATCH 1/2] vm core: update context --- crates/leanVm/src/context/run_context.rs | 136 ++++++++++++++++++++++- crates/leanVm/src/core.rs | 108 +++++++++++++++++- crates/leanVm/src/errors/memory.rs | 10 ++ crates/leanVm/src/errors/mod.rs | 1 + crates/leanVm/src/errors/vm.rs | 16 +++ crates/leanVm/src/memory/manager.rs | 9 ++ crates/leanVm/src/memory/val.rs | 65 +++++++++++ crates/leanVm/src/types/instruction.rs | 4 +- 8 files changed, 345 insertions(+), 4 deletions(-) create mode 100644 crates/leanVm/src/errors/vm.rs diff --git a/crates/leanVm/src/context/run_context.rs b/crates/leanVm/src/context/run_context.rs index 4623fb7a..71bcca8c 100644 --- a/crates/leanVm/src/context/run_context.rs +++ b/crates/leanVm/src/context/run_context.rs @@ -1,4 +1,10 @@ -use crate::memory::address::MemoryAddress; +use p3_field::PrimeField64; + +use crate::{ + errors::{memory::MemoryError, vm::VirtualMachineError}, + memory::{address::MemoryAddress, manager::MemoryManager, val::MemoryValue}, + types::instruction::MemOrConstant, +}; #[derive(Debug, Default)] pub struct RunContext { @@ -25,4 +31,132 @@ impl RunContext { pub const fn fp(&self) -> &MemoryAddress { &self.fp } + + /// Resolves a `MemOrConstant` operand to its final value. + /// + /// - If the operand is a constant, it returns the constant. + /// - If it's a memory location, it computes the address relative to `fp` and fetches the value from memory. + pub fn get_value( + &self, + operand: &MemOrConstant, + memory: &MemoryManager, + ) -> Result, VirtualMachineError> + where + F: PrimeField64, + { + match operand { + MemOrConstant::Constant(val) => Ok(MemoryValue::Int(*val)), + MemOrConstant::MemoryAfterFp { shift } => { + let addr = (self.fp + *shift)?; + memory + .get(addr) + .ok_or_else(|| MemoryError::UninitializedMemory(addr).into()) + } + } + } +} + +#[cfg(test)] +mod tests { + use p3_baby_bear::BabyBear; + use p3_field::PrimeCharacteristicRing; + + use super::*; + + type F = BabyBear; + + #[test] + fn test_get_value_constant() { + // Create a dummy RunContext with pc and fp. + let ctx = RunContext::new( + MemoryAddress { + segment_index: 0, + offset: 0, + }, + MemoryAddress { + segment_index: 1, + offset: 0, + }, + ); + + // A constant operand with field element 42. + let operand = MemOrConstant::Constant(F::from_u64(42)); + + // Run `get_value` with an unused memory manager (memory is not needed for constants). + let memory = MemoryManager::default(); + + // It should return the wrapped constant as a MemoryValue::Int. + let result = ctx.get_value(&operand, &memory).unwrap(); + assert_eq!(result, MemoryValue::Int(F::from_u64(42))); + } + + #[test] + fn test_get_value_memory_after_fp_success() { + let mut memory = MemoryManager::default(); + + // Add a segment that will be used for `fp`. + let fp = memory.add(); // segment_index = 0, offset = 0 + + // Shift = 2, so address to read is fp + 2 => offset 2 in the same segment. + let addr_to_read = MemoryAddress { + segment_index: fp.segment_index, + offset: fp.offset + 2, + }; + + // Insert a value at that address manually. + let expected_val = MemoryValue::Int(F::from_u64(99)); + memory + .memory + .insert(addr_to_read, expected_val.clone()) + .unwrap(); + + // Create a RunContext with that fp. + let ctx = RunContext::new( + MemoryAddress { + segment_index: 0, + offset: 0, + }, // dummy pc + fp, + ); + + // The operand asks to read memory at fp + 2. + let operand = MemOrConstant::MemoryAfterFp { shift: 2 }; + + // Call get_value, which should fetch the value we inserted. + let result = ctx.get_value(&operand, &memory).unwrap(); + assert_eq!(result, expected_val); + } + + #[test] + fn test_get_value_memory_after_fp_uninitialized_memory() { + let mut memory = MemoryManager::default(); + + // Create a segment and set fp to its base. + let fp = memory.add(); // segment_index = 0, offset = 0 + + // We won't insert anything, so all memory is uninitialized. + + // Shift = 1 → fp + 1 points to offset 1 (which is uninitialized). + let operand: MemOrConstant = MemOrConstant::MemoryAfterFp { shift: 1 }; + + // Set up context. + let ctx = RunContext::new( + MemoryAddress { + segment_index: 0, + offset: 0, + }, // dummy pc + fp, + ); + + // Calling get_value should return a VirtualMachineError::MemoryError::UninitializedMemory. + let err = ctx.get_value(&operand, &memory).unwrap_err(); + + match err { + VirtualMachineError::Memory(MemoryError::UninitializedMemory(addr)) => { + assert_eq!(addr.segment_index, fp.segment_index); + assert_eq!(addr.offset, fp.offset + 1); + } + other => panic!("Unexpected error: {other:?}"), + } + } } diff --git a/crates/leanVm/src/core.rs b/crates/leanVm/src/core.rs index 5d8339e9..acae9441 100644 --- a/crates/leanVm/src/core.rs +++ b/crates/leanVm/src/core.rs @@ -1,7 +1,113 @@ -use crate::{context::run_context::RunContext, memory::manager::MemoryManager}; +use p3_field::PrimeField64; + +use crate::{ + context::run_context::RunContext, + errors::{memory::MemoryError, vm::VirtualMachineError}, + memory::{manager::MemoryManager, val::MemoryValue}, + types::instruction::{Instruction, MemOrFp}, +}; #[derive(Debug, Default)] pub struct VirtualMachine { pub(crate) run_context: RunContext, pub memory_manager: MemoryManager, } + +impl VirtualMachine { + /// Advances the program counter (`pc`) to the next instruction. + /// + /// This function embodies the control flow logic of the zkVM. For most instructions, + /// it performs a regular increment of the **`pc`**. However, for the `JumpIfNotZero` + /// instruction (`JUZ`), it implements conditional branching. + /// + /// ### `JumpIfNotZero` Logic + /// + /// When a `JumpIfNotZero` instruction is processed: + /// 1. The `condition` operand is resolved to a field element. + /// 2. If this value is **zero**, the program continues sequentially, and the **`pc`** is incremented by 1. + /// 3. If the value is **non-zero**, a jump is executed. The `dest` operand is resolved to find the + /// target `MemoryAddress`, which then becomes the new **`pc`**. + pub fn update_pc( + &mut self, + instruction: &Instruction, + ) -> Result<(), VirtualMachineError> + where + F: PrimeField64, + { + // Determine the next program counter `pc` by checking if the instruction is a conditional jump. + let next_pc = if let Instruction::JumpIfNotZero { + condition, dest, .. + } = instruction + { + // For a `JumpIfNotZero` instruction, resolve the `condition` operand from memory or constants. + // This will return an error if the memory location is uninitialized. + let condition_val = self + .run_context + .get_value(condition, &self.memory_manager)?; + + // A jump condition must be a field element. + // + // An address is considered non-zero by convention. + let is_zero = match condition_val { + MemoryValue::Int(felt) => felt.is_zero(), + MemoryValue::Address(_) => false, + }; + + if is_zero { + // **Condition is zero**: The jump is not taken. Advance the `pc` by one. + (*self.run_context.pc() + 1)? + } else { + // **Condition is non-zero**: Execute the jump. + // + // First, resolve the `dest` operand to get the target address value. + let dest_val = self.run_context.get_value(dest, &self.memory_manager)?; + + // The resolved destination value must be a valid address. + // + // Convert it and set it as the new `pc`. + dest_val.try_into()? + } + } else { + // For any instruction other than `JumpIfNotZero`, advance the `pc` by one. + (*self.run_context.pc() + 1)? + }; + + // Update the virtual machine's program counter with the calculated next address. + self.run_context.pc = next_pc; + Ok(()) + } + + /// Updates the frame pointer (`fp`) based on the executed instruction. + pub fn update_fp( + &mut self, + instruction: &Instruction, + ) -> Result<(), VirtualMachineError> + where + F: PrimeField64, + { + if let Instruction::JumpIfNotZero { updated_fp, .. } = instruction { + let new_fp = match updated_fp { + // The instruction specifies keeping the same `fp`. + MemOrFp::Fp => self.run_context.fp, + // The instruction specifies updating `fp` to a value from memory. + MemOrFp::MemoryAfterFp { shift } => { + let addr = (*self.run_context.fp() + *shift)?; + let value = self + .memory_manager + .get(addr) + .ok_or(MemoryError::UninitializedMemory(addr))?; + + // The fetched value must be a valid memory address to become the new `fp`. + value.try_into()? + } + }; + self.run_context.fp = new_fp; + } + + // For the other instructions, we do nothing for now. + // + // To be checked in the future. + + Ok(()) + } +} diff --git a/crates/leanVm/src/errors/memory.rs b/crates/leanVm/src/errors/memory.rs index 4a879dbb..efcad66c 100644 --- a/crates/leanVm/src/errors/memory.rs +++ b/crates/leanVm/src/errors/memory.rs @@ -36,4 +36,14 @@ where /// Error related to mathematical operations. #[error(transparent)] Math(#[from] MathError), + + /// Error when a memory value is expected to be an integer, but it is an address to another memory location. + #[error("Memory value should be an integer.")] + ValueNotInteger, + + #[error("Memory at address {0:?} is uninitialized.")] + UninitializedMemory(MemoryAddress), + + #[error("Memory addresses must be relocatable")] + AddressNotRelocatable, } diff --git a/crates/leanVm/src/errors/mod.rs b/crates/leanVm/src/errors/mod.rs index 71d2a089..96c5eccd 100644 --- a/crates/leanVm/src/errors/mod.rs +++ b/crates/leanVm/src/errors/mod.rs @@ -1,2 +1,3 @@ pub mod math; pub mod memory; +pub mod vm; diff --git a/crates/leanVm/src/errors/vm.rs b/crates/leanVm/src/errors/vm.rs new file mode 100644 index 00000000..64f5e804 --- /dev/null +++ b/crates/leanVm/src/errors/vm.rs @@ -0,0 +1,16 @@ +use std::fmt::Debug; + +use thiserror::Error; + +use super::{math::MathError, memory::MemoryError}; + +#[derive(Debug, Error)] +pub enum VirtualMachineError +where + F: Debug, +{ + #[error(transparent)] + Memory(#[from] MemoryError), + #[error(transparent)] + Math(#[from] MathError), +} diff --git a/crates/leanVm/src/memory/manager.rs b/crates/leanVm/src/memory/manager.rs index e8d7cd55..41ca2926 100644 --- a/crates/leanVm/src/memory/manager.rs +++ b/crates/leanVm/src/memory/manager.rs @@ -95,6 +95,15 @@ impl MemoryManager { // This is simply ptr + data.len(), and it may also fail on overflow. (ptr + data.len()).map_err(MemoryError::Math) } + + /// Retrieves the value stored at a given memory address. + #[must_use] + pub fn get(&self, address: MemoryAddress) -> Option> + where + F: PrimeField64, + { + self.memory.get(address) + } } #[cfg(test)] diff --git a/crates/leanVm/src/memory/val.rs b/crates/leanVm/src/memory/val.rs index 094f5d81..706782e3 100644 --- a/crates/leanVm/src/memory/val.rs +++ b/crates/leanVm/src/memory/val.rs @@ -1,7 +1,9 @@ +use p3_field::PrimeField64; #[cfg(test)] use proptest::prelude::*; use super::address::MemoryAddress; +use crate::errors::memory::MemoryError; #[derive(Eq, Ord, Hash, PartialEq, PartialOrd, Clone, Debug)] pub enum MemoryValue { @@ -9,6 +11,20 @@ pub enum MemoryValue { Int(F), } +impl TryInto for MemoryValue +where + F: PrimeField64, +{ + type Error = MemoryError; + + fn try_into(self) -> Result { + match self { + Self::Address(addr) => Ok(addr), + Self::Int(_) => Err(MemoryError::AddressNotRelocatable), + } + } +} + #[cfg(test)] impl Arbitrary for MemoryValue where @@ -27,3 +43,52 @@ where .boxed() } } + +#[cfg(test)] +mod tests { + use p3_baby_bear::BabyBear; + use p3_field::PrimeCharacteristicRing; + + use super::*; + + type F = BabyBear; + + #[test] + fn test_try_into_memory_address_ok() { + // Construct a MemoryAddress. + let addr = MemoryAddress { + segment_index: 3, + offset: 42, + }; + + // Wrap it in a MemoryValue::Address variant + let val: MemoryValue = MemoryValue::Address(addr); + + // Try converting it into a MemoryAddress + let result: Result> = val.try_into(); + + // Assert it succeeds + assert!(result.is_ok()); + + // Assert the returned address is equal to the original + assert_eq!(result.unwrap(), addr); + } + + #[test] + fn test_try_into_memory_address_err_on_int() { + // Create an integer value + let field_elem = F::from_u64(17); + + // Wrap it in a MemoryValue::Int variant + let val: MemoryValue = MemoryValue::Int(field_elem); + + // Try converting it into a MemoryAddress + let result: Result> = val.try_into(); + + // Assert it fails + assert!(result.is_err()); + + // Assert the specific error is AddressNotRelocatable + assert_eq!(result.unwrap_err(), MemoryError::AddressNotRelocatable); + } +} diff --git a/crates/leanVm/src/types/instruction.rs b/crates/leanVm/src/types/instruction.rs index bbbbaaca..9744049a 100644 --- a/crates/leanVm/src/types/instruction.rs +++ b/crates/leanVm/src/types/instruction.rs @@ -104,8 +104,8 @@ pub enum Instruction { /// The value that the result of the double dereference must be equal to. res: MemOrFpOrConstant, }, - /// A conditional jump, called `JUZ` (Jump Unless Zero) - /// . + /// A conditional jump, called `JUZ` (Jump Unless Zero). + /// /// Changes the `pc` if `condition` is non-zero. JumpIfNotZero { /// The value to check. The jump is taken if this value is not zero. From 2599896029916d3192f212aeccd6b26ace7f3517 Mon Sep 17 00:00:00 2001 From: Thomas Coratger Date: Mon, 28 Jul 2025 14:42:04 +0200 Subject: [PATCH 2/2] small fix --- crates/leanVm/src/core.rs | 9 +---- crates/leanVm/src/errors/memory.rs | 7 +++- crates/leanVm/src/memory/address.rs | 61 ++++++++++++++++++++++++++++- crates/leanVm/src/memory/val.rs | 52 ++++++++++++------------ 4 files changed, 92 insertions(+), 37 deletions(-) diff --git a/crates/leanVm/src/core.rs b/crates/leanVm/src/core.rs index acae9441..410252aa 100644 --- a/crates/leanVm/src/core.rs +++ b/crates/leanVm/src/core.rs @@ -3,7 +3,7 @@ use p3_field::PrimeField64; use crate::{ context::run_context::RunContext, errors::{memory::MemoryError, vm::VirtualMachineError}, - memory::{manager::MemoryManager, val::MemoryValue}, + memory::manager::MemoryManager, types::instruction::{Instruction, MemOrFp}, }; @@ -46,12 +46,7 @@ impl VirtualMachine { .get_value(condition, &self.memory_manager)?; // A jump condition must be a field element. - // - // An address is considered non-zero by convention. - let is_zero = match condition_val { - MemoryValue::Int(felt) => felt.is_zero(), - MemoryValue::Address(_) => false, - }; + let is_zero = condition_val.to_f()?.is_zero(); if is_zero { // **Condition is zero**: The jump is not taken. Advance the `pc` by one. diff --git a/crates/leanVm/src/errors/memory.rs b/crates/leanVm/src/errors/memory.rs index efcad66c..588192ac 100644 --- a/crates/leanVm/src/errors/memory.rs +++ b/crates/leanVm/src/errors/memory.rs @@ -44,6 +44,9 @@ where #[error("Memory at address {0:?} is uninitialized.")] UninitializedMemory(MemoryAddress), - #[error("Memory addresses must be relocatable")] - AddressNotRelocatable, + #[error("Memory address is expected but we got an integer.")] + ExpectedMemoryAddress, + + #[error("Inteer is expected but we got a memory address.")] + ExpectedInteger, } diff --git a/crates/leanVm/src/memory/address.rs b/crates/leanVm/src/memory/address.rs index da988018..d39d00af 100644 --- a/crates/leanVm/src/memory/address.rs +++ b/crates/leanVm/src/memory/address.rs @@ -1,9 +1,11 @@ use std::{fmt::Display, ops::Add}; +use p3_field::PrimeField64; #[cfg(test)] use proptest::prelude::*; -use crate::errors::math::MathError; +use super::val::MemoryValue; +use crate::errors::{math::MathError, memory::MemoryError}; #[derive(Eq, Ord, Hash, PartialEq, PartialOrd, Clone, Copy, Debug, Default)] pub struct MemoryAddress { @@ -36,6 +38,20 @@ impl Display for MemoryAddress { } } +impl TryFrom> for MemoryAddress +where + F: PrimeField64, +{ + type Error = MemoryError; + + fn try_from(value: MemoryValue) -> Result { + match value { + MemoryValue::Address(addr) => Ok(addr), + MemoryValue::Int(_) => Err(MemoryError::ExpectedMemoryAddress), + } + } +} + #[cfg(test)] impl Arbitrary for MemoryAddress { type Parameters = (); @@ -58,10 +74,14 @@ impl Arbitrary for MemoryAddress { #[cfg(test)] mod tests { + use p3_baby_bear::BabyBear; + use p3_field::PrimeCharacteristicRing; use proptest::prelude::*; use super::*; + type F = BabyBear; + #[test] fn test_add_usize_success() { let addr = MemoryAddress { @@ -124,4 +144,43 @@ mod tests { } } } + + #[test] + fn test_try_into_memory_address_ok() { + // Construct a MemoryAddress. + let addr = MemoryAddress { + segment_index: 3, + offset: 42, + }; + + // Wrap it in a MemoryValue::Address variant + let val: MemoryValue = MemoryValue::Address(addr); + + // Try converting it into a MemoryAddress + let result: Result> = val.try_into(); + + // Assert it succeeds + assert!(result.is_ok()); + + // Assert the returned address is equal to the original + assert_eq!(result.unwrap(), addr); + } + + #[test] + fn test_try_into_memory_address_err_on_int() { + // Create an integer value + let field_elem = F::from_u64(17); + + // Wrap it in a MemoryValue::Int variant + let val: MemoryValue = MemoryValue::Int(field_elem); + + // Try converting it into a MemoryAddress + let result: Result> = val.try_into(); + + // Assert it fails + assert!(result.is_err()); + + // Assert the specific error is ExpectedMemoryAddress + assert_eq!(result.unwrap_err(), MemoryError::ExpectedMemoryAddress); + } } diff --git a/crates/leanVm/src/memory/val.rs b/crates/leanVm/src/memory/val.rs index 706782e3..9679125a 100644 --- a/crates/leanVm/src/memory/val.rs +++ b/crates/leanVm/src/memory/val.rs @@ -11,16 +11,14 @@ pub enum MemoryValue { Int(F), } -impl TryInto for MemoryValue +impl MemoryValue where F: PrimeField64, { - type Error = MemoryError; - - fn try_into(self) -> Result { + pub const fn to_f(&self) -> Result> { match self { - Self::Address(addr) => Ok(addr), - Self::Int(_) => Err(MemoryError::AddressNotRelocatable), + Self::Address(_) => Err(MemoryError::ExpectedInteger), + Self::Int(f) => Ok(*f), } } } @@ -54,41 +52,41 @@ mod tests { type F = BabyBear; #[test] - fn test_try_into_memory_address_ok() { - // Construct a MemoryAddress. - let addr = MemoryAddress { - segment_index: 3, - offset: 42, - }; + fn test_to_f_ok() { + // Create an integer value + let field_elem = F::from_u64(12345); - // Wrap it in a MemoryValue::Address variant - let val: MemoryValue = MemoryValue::Address(addr); + // Wrap it in a MemoryValue::Int variant + let val: MemoryValue = MemoryValue::Int(field_elem); - // Try converting it into a MemoryAddress - let result: Result> = val.try_into(); + // Call to_f() + let result = val.to_f(); // Assert it succeeds assert!(result.is_ok()); - // Assert the returned address is equal to the original - assert_eq!(result.unwrap(), addr); + // Assert the returned value is equal to the original + assert_eq!(result.unwrap(), field_elem); } #[test] - fn test_try_into_memory_address_err_on_int() { - // Create an integer value - let field_elem = F::from_u64(17); + fn test_to_f_err_on_address() { + // Construct a MemoryAddress. + let addr = MemoryAddress { + segment_index: 1, + offset: 99, + }; - // Wrap it in a MemoryValue::Int variant - let val: MemoryValue = MemoryValue::Int(field_elem); + // Wrap it in a MemoryValue::Address variant + let val: MemoryValue = MemoryValue::Address(addr); - // Try converting it into a MemoryAddress - let result: Result> = val.try_into(); + // Call to_f() + let result = val.to_f(); // Assert it fails assert!(result.is_err()); - // Assert the specific error is AddressNotRelocatable - assert_eq!(result.unwrap_err(), MemoryError::AddressNotRelocatable); + // Assert the specific error is ExpectedInteger + assert_eq!(result.unwrap_err(), MemoryError::ExpectedInteger); } }