vm: reorg memory (#19)

Author: tcoratger

crates/vm/src/bytecode/operand.rs

@@ -5,7 +5,7 @@ use std::{
 
 use p3_field::PrimeCharacteristicRing;
 
-use crate::F;
+use crate::{F, Memory, RunnerError};
 
 /// Represents a value that can either be a constant or a value from memory.
 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -50,6 +50,27 @@ impl MemOrConstant {
     }
 }
 
+impl MemOrConstant {
+    pub fn read_value(&self, memory: &Memory, fp: usize) -> Result<F, RunnerError> {
+        match self {
+            Self::Constant(c) => Ok(*c),
+            Self::MemoryAfterFp { offset } => memory.get(fp + *offset),
+        }
+    }
+
+    #[must_use]
+    pub fn is_value_unknown(&self, memory: &Memory, fp: usize) -> bool {
+        self.read_value(memory, fp).is_err()
+    }
+
+    pub const fn memory_address(&self, fp: usize) -> Result<usize, RunnerError> {
+        match self {
+            Self::Constant(_) => Err(RunnerError::NotAPointer),
+            Self::MemoryAfterFp { offset } => Ok(fp + *offset),
+        }
+    }
+}
+
 impl Display for MemOrConstant {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         match self {
@@ -73,6 +94,28 @@ pub enum MemOrFpOrConstant {
     Constant(F),
 }
 
+impl MemOrFpOrConstant {
+    pub fn read_value(&self, memory: &Memory, fp: usize) -> Result<F, RunnerError> {
+        match self {
+            Self::MemoryAfterFp { offset } => memory.get(fp + *offset),
+            Self::Fp => Ok(F::from_usize(fp)),
+            Self::Constant(c) => Ok(*c),
+        }
+    }
+
+    #[must_use]
+    pub fn is_value_unknown(&self, memory: &Memory, fp: usize) -> bool {
+        self.read_value(memory, fp).is_err()
+    }
+
+    pub const fn memory_address(&self, fp: usize) -> Result<usize, RunnerError> {
+        match self {
+            Self::MemoryAfterFp { offset } => Ok(fp + *offset),
+            Self::Fp | Self::Constant(_) => Err(RunnerError::NotAPointer),
+        }
+    }
+}
+
 impl Display for MemOrFpOrConstant {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         match self {
@@ -112,6 +155,27 @@ impl MemOrFp {
     }
 }
 
+impl MemOrFp {
+    pub fn read_value(&self, memory: &Memory, fp: usize) -> Result<F, RunnerError> {
+        match self {
+            Self::MemoryAfterFp { offset } => memory.get(fp + *offset),
+            Self::Fp => Ok(F::from_usize(fp)),
+        }
+    }
+
+    #[must_use]
+    pub fn is_value_unknown(&self, memory: &Memory, fp: usize) -> bool {
+        self.read_value(memory, fp).is_err()
+    }
+
+    pub const fn memory_address(&self, fp: usize) -> Result<usize, RunnerError> {
+        match self {
+            Self::MemoryAfterFp { offset } => Ok(fp + *offset),
+            Self::Fp => Err(RunnerError::NotAPointer),
+        }
+    }
+}
+
 impl Display for MemOrFp {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         match self {

crates/vm/src/error.rs

@@ -0,0 +1,21 @@
+use thiserror::Error;
+
+use crate::F;
+
+#[derive(Debug, Clone, Error)]
+pub enum RunnerError {
+    #[error("Out of memory")]
+    OutOfMemory,
+    #[error("Memory already set")]
+    MemoryAlreadySet,
+    #[error("Not a pointer")]
+    NotAPointer,
+    #[error("Division by zero")]
+    DivByZero,
+    #[error("Computation invalid: {0} != {1}")]
+    NotEqual(F, F),
+    #[error("Undefined memory access")]
+    UndefinedMemory,
+    #[error("Program counter out of bounds")]
+    PCOutOfBounds,
+}

crates/vm/src/lib.rs

@@ -2,8 +2,12 @@ use p3_field::extension::BinomialExtensionField;
 use p3_koala_bear::KoalaBear;
 
 mod bytecode;
+mod error;
+mod memory;
 mod runner;
 pub use bytecode::*;
+pub use error::*;
+pub use memory::*;
 pub use runner::*;
 
 pub const DIMENSION: usize = 8;

crates/vm/src/memory.rs

@@ -0,0 +1,88 @@
+use p3_field::{BasedVectorSpace, ExtensionField};
+use rayon::prelude::*;
+
+use crate::{DIMENSION, EF, F, RunnerError};
+
+pub(crate) const MAX_MEMORY_SIZE: usize = 1 << 23;
+
+#[derive(Debug, Clone, Default)]
+pub struct Memory(pub Vec<Option<F>>);
+
+impl Memory {
+    pub fn new(public_memory: Vec<F>) -> Self {
+        Self(public_memory.into_par_iter().map(Some).collect())
+    }
+
+    pub fn get(&self, index: usize) -> Result<F, RunnerError> {
+        self.0
+            .get(index)
+            .copied()
+            .flatten()
+            .ok_or(RunnerError::UndefinedMemory)
+    }
+
+    pub fn set(&mut self, index: usize, value: F) -> Result<(), RunnerError> {
+        if index >= self.0.len() {
+            if index >= MAX_MEMORY_SIZE {
+                return Err(RunnerError::OutOfMemory);
+            }
+            self.0.resize(index + 1, None);
+        }
+        if let Some(existing) = &mut self.0[index] {
+            if *existing != value {
+                return Err(RunnerError::MemoryAlreadySet);
+            }
+        } else {
+            self.0[index] = Some(value);
+        }
+        Ok(())
+    }
+
+    pub fn get_vector(&self, index: usize) -> Result<[F; DIMENSION], RunnerError> {
+        Ok(self.get_vectorized_slice(index, 1)?.try_into().unwrap())
+    }
+
+    pub fn get_extension(&self, index: usize) -> Result<EF, RunnerError> {
+        Ok(EF::from_basis_coefficients_slice(&self.get_vector(index)?).unwrap())
+    }
+
+    pub fn get_vectorized_slice(&self, index: usize, len: usize) -> Result<Vec<F>, RunnerError> {
+        let mut vector = Vec::with_capacity(len * DIMENSION);
+        for i in 0..len * DIMENSION {
+            vector.push(self.get(index * DIMENSION + i)?);
+        }
+        Ok(vector)
+    }
+
+    pub fn get_vectorized_slice_extension<EF: ExtensionField<F>>(
+        &self,
+        index: usize,
+        len: usize,
+    ) -> Result<Vec<EF>, RunnerError> {
+        let mut vector = Vec::with_capacity(len);
+        for i in 0..len {
+            let v = self.get_vector(index + i)?;
+            vector.push(EF::from_basis_coefficients_slice(&v).unwrap());
+        }
+        Ok(vector)
+    }
+
+    pub fn slice(&self, index: usize, len: usize) -> Result<Vec<F>, RunnerError> {
+        (0..len).map(|i| self.get(index + i)).collect()
+    }
+
+    pub fn set_vector(&mut self, index: usize, value: [F; DIMENSION]) -> Result<(), RunnerError> {
+        value
+            .iter()
+            .enumerate()
+            .try_for_each(|(i, &v)| self.set(index * DIMENSION + i, v))
+    }
+
+    pub fn set_vectorized_slice(&mut self, index: usize, value: &[F]) -> Result<(), RunnerError> {
+        assert!(value.len().is_multiple_of(DIMENSION));
+        value
+            .iter()
+            .enumerate()
+            .try_for_each(|(i, &v)| self.set(index * DIMENSION + i, v))
+    }
+}