memory: design skeleton with simplest operations (#2)

Author: tcoratger

Cargo.toml

@@ -39,3 +39,6 @@ lean-vm = { path = "crates/leanVm" }
 p3-field = { git = "https://github.com/Plonky3/Plonky3.git", rev = "d0c4a36" }
 p3-baby-bear = { git = "https://github.com/Plonky3/Plonky3.git", rev = "d0c4a36" }
 p3-koala-bear = { git = "https://github.com/Plonky3/Plonky3.git", rev = "d0c4a36" }
+
+thiserror = "2.0"
+proptest = "1.7"

crates/leanVm/Cargo.toml

@@ -12,3 +12,8 @@ workspace = true
 p3-baby-bear.workspace = true
 p3-koala-bear.workspace = true
 p3-field.workspace = true
+
+thiserror.workspace = true
+
+[dev-dependencies]
+proptest.workspace = true

crates/leanVm/src/memory/address.rs

@@ -0,0 +1,28 @@
+#[cfg(test)]
+use proptest::prelude::*;
+
+#[derive(Eq, Ord, Hash, PartialEq, PartialOrd, Clone, Copy, Debug, Default)]
+pub struct MemoryAddress {
+    pub segment_index: usize,
+    pub offset: usize,
+}
+
+#[cfg(test)]
+impl Arbitrary for MemoryAddress {
+    type Parameters = ();
+    type Strategy = BoxedStrategy<Self>;
+
+    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
+        (
+            // segment_index fits in 29 bits
+            0..((1u64 << 29) - 1) as usize,
+            // offset fits in 32 bits
+            0..((1u64 << 32) - 1) as usize,
+        )
+            .prop_map(|(segment_index, offset)| Self {
+                segment_index,
+                offset,
+            })
+            .boxed()
+    }
+}

crates/leanVm/src/memory/cell.rs

@@ -2,6 +2,8 @@ use std::ops::{Deref, DerefMut};
 
 use p3_field::PrimeField64;
 
+use super::{address::MemoryAddress, val::MemoryValue};
+
 /// A memory cell used by the VM for storing 64-bit field elements with metadata.
 ///
 /// Internally, the cell holds a single `u64` value.
@@ -34,26 +36,13 @@ impl MemoryCell {
     pub(crate) const NONE_MASK: u64 = 1 << 63;
     /// Flag bit indicating the cell was accessed (bit 62 set).
     pub(crate) const ACCESS_MASK: u64 = 1 << 62;
+    /// Flag bit indicating the cell contains a MemoryAddress (bit 61 set).
+    pub(crate) const ADDRESS_MASK: u64 = 1 << 61;
+    /// A mask to extract only the value bits, ignoring all flags.
+    pub(crate) const VALUE_MASK: u64 = 0x1FFF_FFFF_FFFF_FFFF;
     /// Constant representing an empty cell.
     pub(crate) const NONE: Self = Self(Self::NONE_MASK);
 
-    /// Creates a `MemoryCell` from a field element, using its canonical `u64` representation.
-    ///
-    /// This clears any flag bits and assumes the value is valid.
-    pub(crate) fn from_f<F>(value: F) -> Self
-    where
-        F: PrimeField64,
-    {
-        Self(value.as_canonical_u64())
-    }
-
-    /// Creates a raw `MemoryCell` from a `u64` value.
-    ///
-    /// Caller is responsible for ensuring no flag bits are set unless intentional.
-    pub(crate) const fn from_u64(value: u64) -> Self {
-        Self(value)
-    }
-
     /// Returns true if the cell is marked as empty (`NONE`).
     pub(crate) const fn is_none(self) -> bool {
         self.0 & Self::NONE_MASK == Self::NONE_MASK
@@ -73,67 +62,223 @@ impl MemoryCell {
     pub(crate) const fn is_accessed(self) -> bool {
         self.0 & Self::ACCESS_MASK == Self::ACCESS_MASK
     }
+
+    pub(crate) fn value<F>(self) -> Option<MemoryValue<F>>
+    where
+        MemoryValue<F>: From<Self>,
+    {
+        self.is_some().then(|| self.into())
+    }
+}
+
+impl<F> From<MemoryValue<F>> for MemoryCell
+where
+    F: PrimeField64,
+{
+    fn from(value: MemoryValue<F>) -> Self {
+        match value {
+            // If it's an integer, store its u64 representation.
+            // The ADDRESS_MASK bit will be 0 by default.
+            MemoryValue::Int(f) => Self(f.as_canonical_u64()),
+
+            // If it's an address, pack it into the u64.
+            MemoryValue::Address(addr) => {
+                // Ensure the address components fit within their allocated bit-space.
+                // 29 bits for segment allows for 536+ million segments.
+                // 32 bits for offset allows for 4+ billion items per segment.
+                debug_assert!(
+                    addr.segment_index < (1 << 29),
+                    "Segment index out of bounds"
+                );
+                debug_assert!(addr.offset < (1 << 32), "Offset out of bounds");
+
+                // Pack segment and offset into a single u64, and set the address flag.
+                let segment = (addr.segment_index as u64) << 32;
+                let offset = addr.offset as u64;
+                Self(segment | offset | Self::ADDRESS_MASK)
+            }
+        }
+    }
+}
+
+impl<F> From<MemoryCell> for MemoryValue<F>
+where
+    F: PrimeField64,
+{
+    fn from(cell: MemoryCell) -> Self {
+        // Check the address flag to determine the type of value.
+        if (cell.0 & MemoryCell::ADDRESS_MASK) == MemoryCell::ADDRESS_MASK {
+            // It's an address, so we unpack it.
+            let segment_index = ((cell.0 & MemoryCell::VALUE_MASK) >> 32) as usize;
+            // Mask for lower 32 bits
+            let offset = (cell.0 & 0xFFFF_FFFF) as usize;
+
+            Self::Address(MemoryAddress {
+                segment_index,
+                offset,
+            })
+        } else {
+            // It's an integer. We extract the value bits and convert to a field element.
+            let value_bits = cell.0 & MemoryCell::VALUE_MASK;
+            Self::Int(F::from_u64(value_bits))
+        }
+    }
 }
 
 #[cfg(test)]
 mod tests {
     use p3_baby_bear::BabyBear;
     use p3_field::PrimeCharacteristicRing;
+    use proptest::prelude::*;
 
     use super::*;
 
     type F = BabyBear;
 
-    #[test]
-    fn test_from_f_and_accessors() {
-        let f = F::from_u64(123);
-        let cell = MemoryCell::from_f(f);
-        assert_eq!(*cell, 123);
-        assert!(cell.is_some());
-        assert!(!cell.is_none());
-        assert!(!cell.is_accessed());
-    }
-
-    #[test]
-    fn test_from_u64_and_flags() {
-        let raw = 0xFFFF_FFFF;
-        let cell = MemoryCell::from_u64(raw);
-        assert_eq!(*cell, raw);
-        assert!(cell.is_some());
-        assert!(!cell.is_none());
-    }
-
     #[test]
     fn test_is_none_and_is_some() {
+        // A cell explicitly created as NONE should be none.
         let none_cell = MemoryCell::NONE;
         assert!(none_cell.is_none());
         assert!(!none_cell.is_some());
 
-        let some_cell = MemoryCell::from_u64(42);
+        // A cell with a value (even zero) should be some.
+        let some_cell = MemoryCell::from(MemoryValue::<F>::Int(F::from_u64(42)));
         assert!(!some_cell.is_none());
         assert!(some_cell.is_some());
+
+        let zero_cell = MemoryCell::from(MemoryValue::<F>::Int(F::ZERO));
+        assert!(!zero_cell.is_none());
+        assert!(zero_cell.is_some());
     }
 
     #[test]
-    fn test_mark_accessed_and_is_accessed() {
-        let mut cell = MemoryCell::from_u64(7);
+    fn test_mark_and_check_accessed() {
+        let mut cell = MemoryCell::from(MemoryValue::<F>::Int(F::from_u64(99)));
+
+        // Initially not accessed.
         assert!(!cell.is_accessed());
+
+        // Mark it as accessed.
         cell.mark_accessed();
+
+        // Now it should be accessed.
         assert!(cell.is_accessed());
+        // Should not affect the NONE flag.
+        assert!(cell.is_some());
 
-        // Ensure value bits are still preserved
-        assert_eq!(*cell & 0x3FFF_FFFF_FFFF_FFFF, 7);
+        // The original value should be preserved alongside the flag.
+        let value_without_flags = cell.0 & MemoryCell::VALUE_MASK;
+        assert_eq!(value_without_flags, 99);
     }
 
     #[test]
-    fn test_none_and_access_bits_do_not_conflict() {
-        let mut cell = MemoryCell::NONE;
-        assert!(cell.is_none());
-        assert!(!cell.is_accessed());
+    fn test_flag_interactions() {
+        // Mark a NONE cell as accessed
+        let mut none_cell = MemoryCell::NONE;
+        none_cell.mark_accessed();
+        assert!(none_cell.is_none(), "is_none should be true after access");
+        assert!(none_cell.is_accessed(), "is_accessed should be true");
+        assert_eq!(none_cell.0, MemoryCell::NONE_MASK | MemoryCell::ACCESS_MASK);
 
-        // Mark accessed should not clear the NONE flag
-        cell.mark_accessed();
-        assert!(cell.is_none());
-        assert!(cell.is_accessed());
+        // Mark an ADDRESS cell as accessed
+        let mut addr_cell = MemoryCell::from(MemoryValue::<F>::Address(MemoryAddress {
+            segment_index: 1,
+            offset: 2,
+        }));
+        addr_cell.mark_accessed();
+        assert!(addr_cell.is_some(), "Address cell should be 'some'");
+        assert!(
+            (addr_cell.0 & MemoryCell::ADDRESS_MASK) != 0,
+            "Address flag should be set"
+        );
+        assert!(
+            addr_cell.is_accessed(),
+            "Address cell should be marked accessed"
+        );
+    }
+
+    #[test]
+    fn test_value_method() {
+        // Test on a NONE cell.
+        let none_cell = MemoryCell::NONE;
+        assert_eq!(none_cell.value::<F>(), None);
+
+        // Test on a valid integer cell.
+        let int_val = MemoryValue::Int(F::from_u64(123));
+        let int_cell = MemoryCell::from(int_val.clone());
+        assert_eq!(int_cell.value(), Some(int_val));
+
+        // Test on a valid address cell.
+        let addr_val = MemoryValue::<F>::Address(MemoryAddress {
+            segment_index: 5,
+            offset: 10,
+        });
+        let addr_cell = MemoryCell::from(addr_val.clone());
+        assert_eq!(addr_cell.value(), Some(addr_val));
+    }
+
+    #[test]
+    fn test_conversion_from_int_value() {
+        let val = MemoryValue::Int(F::from_u64(500));
+        let cell = MemoryCell::from(val);
+        // Should just be the raw value, no flags set.
+        assert_eq!(cell.0, 500);
+    }
+
+    #[test]
+    fn test_conversion_from_address_value() {
+        let val = MemoryValue::<F>::Address(MemoryAddress {
+            segment_index: 10,
+            offset: 20,
+        });
+        let cell = MemoryCell::from(val);
+
+        // Expected packed value: 0x2000000A00000014
+        // Bit 61 (ADDRESS_MASK) + segment 10 shifted by 32 + offset 20
+        let expected = (10u64 << 32) | 20u64 | MemoryCell::ADDRESS_MASK;
+        assert_eq!(cell.0, expected);
+    }
+
+    #[test]
+    fn test_conversion_to_int_value() {
+        // Raw u64 for an integer.
+        let int_cell = MemoryCell(42);
+        let val = MemoryValue::<F>::from(int_cell);
+        assert_eq!(val, MemoryValue::Int(F::from_u64(42)));
+
+        // An integer cell can also be marked accessed; the flag should be ignored.
+        let accessed_int_cell = MemoryCell(42 | MemoryCell::ACCESS_MASK);
+        let accessed_val = MemoryValue::<F>::from(accessed_int_cell);
+        assert_eq!(accessed_val, MemoryValue::Int(F::from_u64(42)));
+    }
+
+    #[test]
+    fn test_conversion_to_address_value() {
+        let raw_addr = (50u64 << 32) | 100u64 | MemoryCell::ADDRESS_MASK;
+        let addr_cell = MemoryCell(raw_addr);
+        let val = MemoryValue::<F>::from(addr_cell);
+
+        let expected = MemoryValue::Address(MemoryAddress {
+            segment_index: 50,
+            offset: 100,
+        });
+        assert_eq!(val, expected);
+    }
+
+    proptest! {
+        #[test]
+        fn proptest_roundtrip_conversion(
+            val in any::<MemoryValue<F>>()
+        ) {
+            // Convert the generated MemoryValue to a MemoryCell.
+            let cell = MemoryCell::from(val.clone());
+
+            // Convert the MemoryCell back to a MemoryValue.
+            let roundtrip_val = MemoryValue::<F>::from(cell);
+
+            // Assert that the original and round-tripped values are identical.
+            prop_assert_eq!(val, roundtrip_val);
+        }
     }
 }

crates/leanVm/src/memory/error.rs

@@ -0,0 +1,34 @@
+use std::fmt::Debug;
+
+use thiserror::Error;
+
+use super::{address::MemoryAddress, val::MemoryValue};
+
+#[derive(Debug, Eq, PartialEq, Error)]
+pub enum MemoryError<F>
+where
+    F: Debug,
+{
+    /// Error for when an operation targets a memory segment that has not been allocated.
+    #[error(
+        "Memory access out of bounds: cannot access segment {}, as only {} segments are allocated.",
+        0.0,
+        0.1
+    )]
+    UnallocatedSegment(Box<(usize, usize)>),
+
+    /// Error for attempting to overwrite an existing, different value in a memory cell, violating write-once consistency.
+    #[error(
+        "Write-once violation at address {:?}: cannot overwrite existing value '{:?}' with new value '{:?}'.",
+        0.0,
+        0.1,
+        0.2
+    )]
+    InconsistentMemory(Box<(MemoryAddress, MemoryValue<F>, MemoryValue<F>)>),
+
+    /// Error for when a memory operation would exceed the maximum capacity of a segment vector.
+    #[error(
+        "Memory overflow: the requested memory address is too large and exceeds the machine's capacity."
+    )]
+    VecCapacityExceeded,
+}