vm core: update context (#5)

* vm core: update context

* small fix

Author: tcoratger

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<F>(
+        &self,
+        operand: &MemOrConstant<F>,
+        memory: &MemoryManager,
+    ) -> Result<MemoryValue<F>, VirtualMachineError<F>>
+    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<F> = 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:?}"),
+        }
+    }
 }

crates/leanVm/src/core.rs

@@ -1,7 +1,108 @@
-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,
+    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<F>(
+        &mut self,
+        instruction: &Instruction<F>,
+    ) -> Result<(), VirtualMachineError<F>>
+    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.
+            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.
+                (*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<F>(
+        &mut self,
+        instruction: &Instruction<F>,
+    ) -> Result<(), VirtualMachineError<F>>
+    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(())
+    }
+}

crates/leanVm/src/errors/memory.rs

@@ -36,4 +36,17 @@ 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 address is expected but we got an integer.")]
+    ExpectedMemoryAddress,
+
+    #[error("Inteer is expected but we got a memory address.")]
+    ExpectedInteger,
 }

crates/leanVm/src/errors/mod.rs

@@ -1,2 +1,3 @@
 pub mod math;
 pub mod memory;
+pub mod vm;

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<F>
+where
+    F: Debug,
+{
+    #[error(transparent)]
+    Memory(#[from] MemoryError<F>),
+    #[error(transparent)]
+    Math(#[from] MathError),
+}

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<F> TryFrom<MemoryValue<F>> for MemoryAddress
+where
+    F: PrimeField64,
+{
+    type Error = MemoryError<F>;
+
+    fn try_from(value: MemoryValue<F>) -> Result<Self, Self::Error> {
+        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<F> = MemoryValue::Address(addr);
+
+        // Try converting it into a MemoryAddress
+        let result: Result<MemoryAddress, MemoryError<F>> = 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<F> = MemoryValue::Int(field_elem);
+
+        // Try converting it into a MemoryAddress
+        let result: Result<MemoryAddress, MemoryError<F>> = val.try_into();
+
+        // Assert it fails
+        assert!(result.is_err());
+
+        // Assert the specific error is ExpectedMemoryAddress
+        assert_eq!(result.unwrap_err(), MemoryError::ExpectedMemoryAddress);
+    }
 }