vm core: implement run_instruction (#6)

* vm core: implement run_instruction

* rm core for now

Author: tcoratger

Cargo.toml

@@ -42,3 +42,4 @@ p3-koala-bear = { git = "https://github.com/Plonky3/Plonky3.git", rev = "d0c4a36
 
 thiserror = "2.0"
 proptest = "1.7"
+num-traits = "0.2"

crates/leanVm/Cargo.toml

@@ -14,6 +14,7 @@ p3-koala-bear.workspace = true
 p3-field.workspace = true
 
 thiserror.workspace = true
+num-traits.workspace = true
 
 [dev-dependencies]
 proptest.workspace = true

crates/leanVm/src/context/run_context.rs

@@ -3,7 +3,7 @@ use p3_field::PrimeField64;
 use crate::{
     errors::{memory::MemoryError, vm::VirtualMachineError},
     memory::{address::MemoryAddress, manager::MemoryManager, val::MemoryValue},
-    types::instruction::MemOrConstant,
+    types::instruction::{MemOrConstant, MemOrFp, MemOrFpOrConstant},
 };
 
 #[derive(Debug, Default)]
@@ -47,7 +47,56 @@ impl RunContext {
         match operand {
             MemOrConstant::Constant(val) => Ok(MemoryValue::Int(*val)),
             MemOrConstant::MemoryAfterFp { shift } => {
-                let addr = (self.fp + *shift)?;
+                let addr = self.fp.add_usize(*shift)?;
+                memory
+                    .get(addr)
+                    .ok_or_else(|| MemoryError::UninitializedMemory(addr).into())
+            }
+        }
+    }
+
+    /// Resolves a `MemOrFp` operand to its final value.
+    ///
+    /// - If the operand is the frame pointer `Fp`, it returns the `fp` address itself.
+    /// - If it's a memory location, it computes the address relative to `fp` and fetches the value.
+    pub fn get_value_from_mem_or_fp<F>(
+        &self,
+        operand: &MemOrFp,
+        memory: &MemoryManager,
+    ) -> Result<MemoryValue<F>, VirtualMachineError<F>>
+    where
+        F: PrimeField64,
+    {
+        match operand {
+            MemOrFp::Fp => Ok(MemoryValue::Address(self.fp)),
+            MemOrFp::MemoryAfterFp { shift } => {
+                let addr = self.fp.add_usize(*shift)?;
+                memory
+                    .get(addr)
+                    .ok_or_else(|| MemoryError::UninitializedMemory(addr).into())
+            }
+        }
+    }
+
+    /// Resolves a `MemOrFpOrConstant` operand to its final value.
+    ///
+    /// This is a comprehensive resolver that handles all three potential operand types:
+    /// - a constant value,
+    /// - a memory location relative to `fp`,
+    /// - the `fp` register itself.
+    pub fn get_value_from_mem_or_fp_or_constant<F>(
+        &self,
+        operand: &MemOrFpOrConstant<F>,
+        memory: &MemoryManager,
+    ) -> Result<MemoryValue<F>, VirtualMachineError<F>>
+    where
+        F: PrimeField64,
+    {
+        match operand {
+            MemOrFpOrConstant::Constant(val) => Ok(MemoryValue::Int(*val)),
+            MemOrFpOrConstant::Fp => Ok(MemoryValue::Address(self.fp)),
+            MemOrFpOrConstant::MemoryAfterFp { shift } => {
+                let addr = self.fp.add_usize(*shift)?;
                 memory
                     .get(addr)
                     .ok_or_else(|| MemoryError::UninitializedMemory(addr).into())
@@ -159,4 +208,46 @@ mod tests {
             other => panic!("Unexpected error: {other:?}"),
         }
     }
+
+    #[test]
+    fn test_get_value_from_mem_or_fp_or_constant_is_constant() {
+        let ctx = RunContext::new(MemoryAddress::new(0, 0), MemoryAddress::new(1, 0));
+        let operand = MemOrFpOrConstant::Constant(F::from_u64(123));
+        let memory = MemoryManager::default();
+        let result = ctx
+            .get_value_from_mem_or_fp_or_constant(&operand, &memory)
+            .unwrap();
+        assert_eq!(result, MemoryValue::Int(F::from_u64(123)));
+    }
+
+    #[test]
+    fn test_get_value_from_mem_or_fp_or_constant_is_fp() {
+        let fp_addr = MemoryAddress::new(1, 10);
+        let ctx = RunContext::new(MemoryAddress::new(0, 0), fp_addr);
+        let operand = MemOrFpOrConstant::<F>::Fp;
+        let memory = MemoryManager::default();
+        let result = ctx
+            .get_value_from_mem_or_fp_or_constant(&operand, &memory)
+            .unwrap();
+        assert_eq!(result, MemoryValue::Address(fp_addr));
+    }
+
+    #[test]
+    fn test_get_value_from_mem_or_fp_or_constant_is_mem_success() {
+        let mut memory = MemoryManager::default();
+        let fp = memory.add();
+        let addr_to_read = fp.add_usize::<F>(7).unwrap();
+        let expected_val = MemoryValue::<F>::Address(MemoryAddress::new(5, 5));
+        memory
+            .memory
+            .insert(addr_to_read, expected_val.clone())
+            .unwrap();
+
+        let ctx = RunContext::new(MemoryAddress::new(0, 0), fp);
+        let operand = MemOrFpOrConstant::MemoryAfterFp { shift: 7 };
+        let result = ctx
+            .get_value_from_mem_or_fp_or_constant(&operand, &memory)
+            .unwrap();
+        assert_eq!(result, expected_val);
+    }
 }

crates/leanVm/src/core.rs

@@ -50,7 +50,7 @@ impl VirtualMachine {
 
             if is_zero {
                 // **Condition is zero**: The jump is not taken. Advance the `pc` by one.
-                (*self.run_context.pc() + 1)?
+                self.run_context.pc().add_usize(1)?
             } else {
                 // **Condition is non-zero**: Execute the jump.
                 //
@@ -64,7 +64,7 @@ impl VirtualMachine {
             }
         } else {
             // For any instruction other than `JumpIfNotZero`, advance the `pc` by one.
-            (*self.run_context.pc() + 1)?
+            self.run_context.pc().add_usize(1)?
         };
 
         // Update the virtual machine's program counter with the calculated next address.
@@ -86,7 +86,7 @@ impl VirtualMachine {
                 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 addr = self.run_context.fp().add_usize(*shift)?;
                     let value = self
                         .memory_manager
                         .get(addr)
@@ -189,7 +189,11 @@ mod tests {
         let pc = MemoryAddress::new(0, 10);
         let fp = MemoryAddress::new(1, 5);
         // Pre-load memory with a zero value at the address `fp + 1`, which will be our condition.
-        let mut vm = setup_vm(pc, fp, &[((fp + 1).unwrap(), MemoryValue::Int(F::ZERO))]);
+        let mut vm = setup_vm(
+            pc,
+            fp,
+            &[(fp.add_usize::<F>(1).unwrap(), MemoryValue::Int(F::ZERO))],
+        );
         // Define a JNZ instruction where the condition points to the zero value.
         let instruction = Instruction::JumpIfNotZero::<F> {
             condition: MemOrConstant::MemoryAfterFp { shift: 1 },
@@ -215,9 +219,15 @@ mod tests {
             fp,
             &[
                 // The condition value (non-zero).
-                ((fp + 1).unwrap(), MemoryValue::Int(F::from_u64(42))),
+                (
+                    fp.add_usize::<F>(1).unwrap(),
+                    MemoryValue::Int(F::from_u64(42)),
+                ),
                 // The destination address for the jump.
-                ((fp + 2).unwrap(), MemoryValue::Address(jump_target)),
+                (
+                    fp.add_usize::<F>(2).unwrap(),
+                    MemoryValue::Address(jump_target),
+                ),
             ],
         );
         // Define a JNZ instruction pointing to the condition and destination.
@@ -242,7 +252,7 @@ mod tests {
             pc,
             fp,
             &[(
-                (fp + 1).unwrap(),
+                fp.add_usize::<F>(1).unwrap(),
                 MemoryValue::Address(MemoryAddress::new(8, 8)),
             )],
         );
@@ -287,7 +297,7 @@ mod tests {
         let mut vm = setup_vm(
             MemoryAddress::new(0, 0),
             fp,
-            &[((fp + 3).unwrap(), MemoryValue::Address(new_fp))],
+            &[(fp.add_usize::<F>(3).unwrap(), MemoryValue::Address(new_fp))],
         );
         // Define a JNZ instruction where `updated_fp` points to the new address in memory.
         let instruction = Instruction::JumpIfNotZero::<F> {
@@ -309,7 +319,10 @@ mod tests {
         let mut vm = setup_vm(
             MemoryAddress::new(0, 0),
             fp,
-            &[((fp + 3).unwrap(), MemoryValue::Int(F::from_u64(99)))],
+            &[(
+                fp.add_usize::<F>(3).unwrap(),
+                MemoryValue::Int(F::from_u64(99)),
+            )],
         );
         // Define a JNZ instruction where `updated_fp` points to this integer value.
         let instruction = Instruction::JumpIfNotZero::<F> {

crates/leanVm/src/errors/math.rs

@@ -1,9 +1,15 @@
+use p3_field::PrimeField64;
 use thiserror::Error;
 
 use crate::memory::address::MemoryAddress;
 
 #[derive(Debug, Error, Eq, PartialEq)]
-pub enum MathError {
+pub enum MathError<F>
+where
+    F: PrimeField64,
+{
     #[error("Operation failed: {} + {}, maximum offset value exceeded", 0.0, 0.1)]
     MemoryAddressAddUsizeOffsetExceeded(Box<(MemoryAddress, usize)>),
+    #[error("Operation failed: {} + {}, maximum offset value exceeded", 0.0, 0.1)]
+    MemoryAddressAddFieldOffsetExceeded(Box<(MemoryAddress, F)>),
 }

crates/leanVm/src/errors/memory.rs

@@ -1,5 +1,6 @@
 use std::fmt::Debug;
 
+use p3_field::PrimeField64;
 use thiserror::Error;
 
 use super::math::MathError;
@@ -8,7 +9,7 @@ use crate::memory::{address::MemoryAddress, val::MemoryValue};
 #[derive(Debug, Eq, PartialEq, Error)]
 pub enum MemoryError<F>
 where
-    F: Debug,
+    F: PrimeField64,
 {
     /// Error for when an operation targets a memory segment that has not been allocated.
     #[error(
@@ -35,7 +36,7 @@ where
 
     /// Error related to mathematical operations.
     #[error(transparent)]
-    Math(#[from] MathError),
+    Math(#[from] MathError<F>),
 
     /// 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.")]
@@ -47,6 +48,12 @@ where
     #[error("Memory address is expected but we got an integer.")]
     ExpectedMemoryAddress,
 
-    #[error("Inteer is expected but we got a memory address.")]
+    #[error("Integer is expected but we got a memory address.")]
     ExpectedInteger,
+
+    #[error("Operation failed: {} + {}, can't add two address values", 0.0, 0.1)]
+    MemoryAddressAdd(Box<(MemoryAddress, MemoryAddress)>),
+
+    #[error("Operation failed: {} * {}, can't multiply these two values", 0.0, 0.1)]
+    InvalidMul(Box<(MemoryValue<F>, MemoryValue<F>)>),
 }

crates/leanVm/src/errors/vm.rs

@@ -1,16 +1,27 @@
 use std::fmt::Debug;
 
+use p3_field::PrimeField64;
 use thiserror::Error;
 
 use super::{math::MathError, memory::MemoryError};
+use crate::memory::val::MemoryValue;
 
 #[derive(Debug, Error)]
 pub enum VirtualMachineError<F>
 where
-    F: Debug,
+    F: PrimeField64,
 {
     #[error(transparent)]
     Memory(#[from] MemoryError<F>),
     #[error(transparent)]
-    Math(#[from] MathError),
+    Math(#[from] MathError<F>),
+    #[error(
+        "Assertion failed: computed value '{:?}' != expected value '{:?}'.",
+        computed,
+        expected
+    )]
+    AssertEqFailed {
+        computed: MemoryValue<F>,
+        expected: MemoryValue<F>,
+    },
 }