memory: add memory manager (#4)

* memory: add memory manager

* fmt

* fmt

* better doc

* make clippy happy

Author: tcoratger

crates/leanVm/src/context/run_context.rs

@@ -1,6 +1,6 @@
 use crate::memory::address::MemoryAddress;
 
-#[derive(Debug)]
+#[derive(Debug, Default)]
 pub struct RunContext {
     /// The address in memory of the current instruction to be executed.
     pub(crate) pc: MemoryAddress,

crates/leanVm/src/core.rs

@@ -0,0 +1,7 @@
+use crate::{context::run_context::RunContext, memory::manager::MemoryManager};
+
+#[derive(Debug, Default)]
+pub struct VirtualMachine {
+    pub(crate) run_context: RunContext,
+    pub memory_manager: MemoryManager,
+}

crates/leanVm/src/lib.rs

@@ -1,3 +1,4 @@
 pub mod context;
+pub mod core;
 pub mod memory;
 pub mod types;

crates/leanVm/src/memory/address.rs

@@ -1,12 +1,41 @@
+use std::{fmt::Display, ops::Add};
+
 #[cfg(test)]
 use proptest::prelude::*;
 
+use crate::types::math_errors::MathError;
+
 #[derive(Eq, Ord, Hash, PartialEq, PartialOrd, Clone, Copy, Debug, Default)]
 pub struct MemoryAddress {
     pub segment_index: usize,
     pub offset: usize,
 }
 
+impl Add<usize> for MemoryAddress {
+    type Output = Result<Self, MathError>;
+
+    fn add(self, other: usize) -> Result<Self, MathError> {
+        // Try to compute the new offset by adding `other` to the current offset.
+        //
+        // This uses `checked_add` to safely detect any potential `usize` overflow.
+        self.offset
+            .checked_add(other)
+            .map(|offset| Self {
+                // Keep the same segment index.
+                segment_index: self.segment_index,
+                // Use the new (safe) offset.
+                offset,
+            })
+            .ok_or_else(|| MathError::MemoryAddressAddUsizeOffsetExceeded(Box::new((self, other))))
+    }
+}
+
+impl Display for MemoryAddress {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}:{}", self.segment_index, self.offset)
+    }
+}
+
 #[cfg(test)]
 impl Arbitrary for MemoryAddress {
     type Parameters = ();
@@ -26,3 +55,73 @@ impl Arbitrary for MemoryAddress {
             .boxed()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use proptest::prelude::*;
+
+    use super::*;
+
+    #[test]
+    fn test_add_usize_success() {
+        let addr = MemoryAddress {
+            segment_index: 2,
+            offset: 100,
+        };
+        let result = addr + 25;
+        assert_eq!(
+            result,
+            Ok(MemoryAddress {
+                segment_index: 2,
+                offset: 125
+            })
+        );
+    }
+
+    #[test]
+    fn test_add_zero_offset() {
+        let addr = MemoryAddress {
+            segment_index: 5,
+            offset: 500,
+        };
+        let result = addr + 0;
+        assert_eq!(result, Ok(addr));
+    }
+
+    #[test]
+    fn test_add_usize_overflow() {
+        let addr = MemoryAddress {
+            segment_index: 1,
+            offset: usize::MAX,
+        };
+        let result = addr + 1;
+        match result {
+            Err(MathError::MemoryAddressAddUsizeOffsetExceeded(boxed)) => {
+                let (original, added) = *boxed;
+                assert_eq!(original.segment_index, 1);
+                assert_eq!(original.offset, usize::MAX);
+                assert_eq!(added, 1);
+            }
+            _ => panic!("Expected overflow error, got: {result:?}"),
+        }
+    }
+
+    proptest! {
+        #[test]
+        fn test_add_does_not_overflow(addr in any::<MemoryAddress>(), delta in 0usize..1_000_000) {
+            let result = addr + delta;
+            // Only test when offset + delta won't overflow
+            if let Some(expected_offset) = addr.offset.checked_add(delta) {
+                prop_assert_eq!(result, Ok(MemoryAddress {
+                    segment_index: addr.segment_index,
+                    offset: expected_offset,
+                }));
+            } else {
+                prop_assert!(matches!(
+                    result,
+                    Err(MathError::MemoryAddressAddUsizeOffsetExceeded(_))
+                ));
+            }
+        }
+    }
+}

crates/leanVm/src/memory/error.rs

@@ -3,6 +3,7 @@ use std::fmt::Debug;
 use thiserror::Error;
 
 use super::{address::MemoryAddress, val::MemoryValue};
+use crate::types::math_errors::MathError;
 
 #[derive(Debug, Eq, PartialEq, Error)]
 pub enum MemoryError<F>
@@ -31,4 +32,8 @@ where
         "Memory overflow: the requested memory address is too large and exceeds the machine's capacity."
     )]
     VecCapacityExceeded,
+
+    /// Error related to mathematical operations.
+    #[error(transparent)]
+    Math(#[from] MathError),
 }

crates/leanVm/src/memory/manager.rs

@@ -0,0 +1,238 @@
+use p3_field::PrimeField64;
+
+use super::{address::MemoryAddress, error::MemoryError, mem::Memory, val::MemoryValue};
+
+/// A high level manager for the memory.
+#[derive(Debug, Default)]
+pub struct MemoryManager {
+    pub memory: Memory,
+}
+
+impl MemoryManager {
+    /// Returns the number of currently allocated segments in memory.
+    ///
+    /// This reflects the actual physical segment count, starting from segment index 0.
+    ///
+    /// # Returns
+    /// * `usize` — the number of segments allocated in `self.memory`.
+    #[must_use]
+    pub fn num_segments(&self) -> usize {
+        self.memory.data.len()
+    }
+
+    /// Adds a new, empty segment to memory and returns its starting address.
+    ///
+    /// This operation appends an empty segment to the memory, which starts at offset 0.
+    /// The returned `MemoryAddress` corresponds to the beginning of that segment.
+    ///
+    /// # Returns
+    /// * `MemoryAddress` — the starting address (segment_index, offset=0) of the new segment.
+    pub fn add(&mut self) -> MemoryAddress {
+        // Allocate a new, empty segment at the end of the list.
+        self.memory.data.push(Vec::new());
+
+        // Compute the index of the newly created segment.
+        let new_segment_index = self.memory.data.len() - 1;
+
+        // Return the starting address of the new segment (offset always 0).
+        MemoryAddress {
+            segment_index: new_segment_index,
+            offset: 0,
+        }
+    }
+
+    /// Loads a slice of data into memory starting from a given address.
+    ///
+    /// The function writes each value in `data` to consecutive addresses starting
+    /// from `ptr`. The write is done in reverse order to ensure that any required
+    /// memory extension happens once at the end rather than multiple times.
+    ///
+    /// If all values are written successfully, the function returns the first
+    /// address **after** the last inserted value.
+    ///
+    /// # Type Parameters
+    /// * `F`: A finite field, used as the scalar type.
+    ///
+    /// # Arguments
+    /// * `ptr`: Starting address where the data should be written.
+    /// * `data`: A slice of memory values representing the values to be stored.
+    ///
+    /// # Returns
+    /// * `Ok(MemoryAddress)` — the address immediately following the last written value.
+    /// * `Err(MemoryError<F>)` — if writing fails due to:
+    ///     - Memory cell already initialized with a different value.
+    ///     - Overflow when computing addresses.
+    ///     - Exceeding vector capacity.
+    pub fn load_data<F>(
+        &mut self,
+        ptr: MemoryAddress,
+        data: &[MemoryValue<F>],
+    ) -> Result<MemoryAddress, MemoryError<F>>
+    where
+        F: PrimeField64,
+    {
+        // Iterate over the data values in reverse order, with indices.
+        //
+        // This reverse order allows any required memory segment resizing
+        // (e.g., length extension or capacity reservation) to occur *once*
+        // at the highest offset instead of repeatedly during writes.
+        for (num, value) in data.iter().enumerate().rev() {
+            // Compute the target address: ptr + num.
+            //
+            // This operation may fail if it causes overflow.
+            let addr = (ptr + num).map_err(MemoryError::Math)?;
+
+            // Attempt to write the value into memory at the computed address.
+            //
+            // This enforces the write-once rule — it will fail if the cell is already
+            // initialized with a different value.
+            self.memory.insert(addr, value.clone())?;
+        }
+
+        // After writing all values, compute and return the address after the last item.
+        //
+        // This is simply ptr + data.len(), and it may also fail on overflow.
+        (ptr + data.len()).map_err(MemoryError::Math)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use p3_baby_bear::BabyBear;
+    use p3_field::PrimeCharacteristicRing;
+
+    use super::*;
+    use crate::{memory::cell::MemoryCell, types::math_errors::MathError};
+
+    type F = BabyBear;
+
+    #[test]
+    fn test_add_segment_returns_correct_address() {
+        // Create a new empty memory manager.
+        let mut manager = MemoryManager::default();
+
+        // Initially, there should be no segments.
+        assert_eq!(manager.num_segments(), 0);
+
+        // Add the first memory segment.
+        let addr1 = manager.add();
+
+        // The first segment should have index 0, starting at offset 0.
+        assert_eq!(addr1.segment_index, 0);
+        assert_eq!(addr1.offset, 0);
+
+        // After adding, the total number of segments should be 1.
+        assert_eq!(manager.num_segments(), 1);
+
+        // Add another segment.
+        let addr2 = manager.add();
+
+        // The second segment should have index 1, also starting at offset 0.
+        assert_eq!(addr2.segment_index, 1);
+        assert_eq!(addr2.offset, 0);
+
+        // Now there should be exactly 2 segments.
+        assert_eq!(manager.num_segments(), 2);
+    }
+
+    #[test]
+    fn test_load_data_successful() {
+        // Create a new memory manager.
+        let mut manager = MemoryManager::default();
+
+        // Add a memory segment and get its starting address.
+        let base_addr = manager.add(); // segment_index = 0, offset = 0
+
+        // Prepare a list of memory values to load.
+        let values = vec![
+            MemoryValue::Int(F::from_u64(10)),
+            MemoryValue::Int(F::from_u64(20)),
+            MemoryValue::Int(F::from_u64(30)),
+        ];
+
+        // Load the data into memory starting at base_addr.
+        let end_addr = manager.load_data(base_addr, &values).unwrap();
+
+        // The returned end address should be immediately after the last inserted value.
+        assert_eq!(end_addr.segment_index, base_addr.segment_index);
+        assert_eq!(end_addr.offset, base_addr.offset + values.len());
+
+        // Verify that each value was inserted correctly at its expected offset.
+        for (i, expected) in values.iter().enumerate() {
+            let addr = MemoryAddress {
+                segment_index: base_addr.segment_index,
+                offset: base_addr.offset + i,
+            };
+            assert_eq!(manager.memory.get(addr), Some(expected.clone()));
+        }
+    }
+
+    #[test]
+    fn test_load_data_returns_math_error_on_address_overflow() {
+        // Create a new memory manager.
+        let mut manager = MemoryManager::default();
+
+        // Define a starting address where the offset is at the maximum `usize` value.
+        let base_addr = MemoryAddress {
+            segment_index: 0,
+            offset: usize::MAX,
+        };
+
+        // Manually push an empty segment to allow writing into segment 0.
+        manager.memory.data.push(Vec::new());
+
+        // Create two values to write: this will cause an overflow.
+        let values = vec![
+            MemoryValue::Int(F::from_u64(1)),
+            MemoryValue::Int(F::from_u64(2)),
+        ];
+
+        // Try to load data starting at MAX offset.
+        //
+        // This should fail with an overflow error.
+        let err = manager.load_data(base_addr, &values).unwrap_err();
+
+        // Confirm that the error is a MathError due to offset overflow.
+        match err {
+            MemoryError::Math(MathError::MemoryAddressAddUsizeOffsetExceeded(boxed)) => {
+                let (addr, delta) = *boxed;
+                // original address
+                assert_eq!(addr, base_addr);
+                // the amount that caused overflow
+                assert_eq!(delta, 1);
+            }
+            other => panic!("Unexpected error: {other:?}"),
+        }
+    }
+
+    #[test]
+    fn test_load_data_partial_write_does_not_corrupt_memory() {
+        // Create a new memory manager.
+        let mut manager = MemoryManager::default();
+
+        // Add a segment (segment_index = 0).
+        let _ = manager.add();
+
+        // Construct values to insert. The second will cause a failure.
+        let values = vec![
+            MemoryValue::Int(F::from_u64(1)),
+            MemoryValue::Int(F::from_u64(2)),
+        ];
+
+        // Set the starting address such that adding even one to it will cause overflow.
+        let failing_addr = MemoryAddress {
+            segment_index: 0,
+            offset: usize::MAX - 1,
+        };
+
+        // Simulate a memory segment with a small preallocated length.
+        manager.memory.data[0] = vec![MemoryCell::NONE; 4];
+
+        // Attempt to load the data starting at the failing address.
+        // This should fail due to memory limitations.
+        let err = manager.load_data(failing_addr, &values).unwrap_err();
+
+        // Confirm the error is due to exceeding vector capacity.
+        assert!(matches!(err, MemoryError::VecCapacityExceeded));
+    }
+}

crates/leanVm/src/memory/mod.rs

@@ -1,5 +1,6 @@
 pub mod address;
 pub mod cell;
 pub mod error;
+pub mod manager;
 pub mod mem;
 pub mod val;