🚧 Memory WIP

♻ Some small cleaning

crates/mipsevm/Cargo.toml

@@ -9,6 +9,8 @@ authors.workspace = true
 [dependencies]
 alloy-primitives = "0.3.3"
 once_cell = "1.18.0"
+serde = "1.0.188"
+serde_json = "1.0.106"
 
 anyhow = "1.0.70"
 tracing = "0.1.37"

crates/mipsevm/src/lib.rs

@@ -2,6 +2,7 @@
 #![feature(generic_const_exprs)]
 #![allow(incomplete_features, dead_code)]
 
+mod memory;
 mod page;
 mod state;
 mod traits;

crates/mipsevm/src/memory.rs

@@ -0,0 +1,95 @@
+//! The memory module contains the memory data structures and functionality for the emulator.
+
+use crate::page::{self, CachedPage};
+use alloy_primitives::B256;
+use anyhow::Result;
+use std::{cell::RefCell, collections::BTreeMap, rc::Rc};
+
+type PageIndex = u64;
+
+/// The [Memory] struct represents the MIPS emulator's memory.
+struct Memory {
+    /// Map of generalized index -> the merkle root of each index. None if invalidated.
+    nodes: BTreeMap<u64, Option<B256>>,
+    /// Map of page indices to [CachedPage]s.
+    pages: BTreeMap<u64, Rc<RefCell<CachedPage>>>,
+    /// We store two caches upfront; we often read instructions from one page and reserve another
+    /// for scratch memory. This prevents map lookups for each instruction.
+    last_page: [(PageIndex, Option<Rc<RefCell<CachedPage>>>); 2],
+}
+
+impl Default for Memory {
+    fn default() -> Self {
+        Self {
+            nodes: BTreeMap::default(),
+            pages: BTreeMap::default(),
+            last_page: [(!0u64, None), (!0u64, None)],
+        }
+    }
+}
+
+impl Memory {
+    /// Returns the number of allocated pages in memory.
+    pub fn page_count(&self) -> usize {
+        self.pages.len()
+    }
+
+    /// Performs an operation on all pages in the memory.
+    ///
+    /// ### Takes
+    /// - `f`: A function that takes a page index and a mutable reference to a [CachedPage].
+    pub fn for_each_page(&mut self, mut f: impl FnMut(u64, Rc<RefCell<CachedPage>>)) {
+        for (key, page) in self.pages.iter() {
+            f(*key, Rc::clone(page));
+        }
+    }
+
+    /// Invalidate a given memory address
+    ///
+    /// ### Takes
+    /// - `address`: The address to invalidate.
+    pub fn invalidate(&mut self, address: u64) -> Result<()> {
+        if address & 0x3 != 0 {
+            panic!("Unaligned memory access: {:x}", address);
+        }
+
+        // Find the page and invalidate the address within it.
+        if let Some(page) = self.page_lookup(address >> page::PAGE_ADDRESS_SIZE) {
+            let mut page = page.borrow_mut();
+            page.invalidate(address & page::PAGE_ADDRESS_MASK as u64)?;
+            if !page.valid[1] {
+                return Ok(());
+            }
+        } else {
+            // Nothing to invalidate
+            return Ok(());
+        }
+
+        // Find the generalized index of the first page covering the address
+        let mut g_index = ((1u64 << 32) | address) >> page::PAGE_ADDRESS_SIZE as u64;
+        while g_index > 0 {
+            self.nodes.insert(g_index, None);
+            g_index >>= 1;
+        }
+
+        Ok(())
+    }
+
+    fn page_lookup(&mut self, page_index: PageIndex) -> Option<Rc<RefCell<CachedPage>>> {
+        // Check caches before maps
+        if let Some((_, Some(page))) = self.last_page.iter().find(|(key, _)| *key == page_index) {
+            Some(Rc::clone(page))
+        } else if let Some(page) = self.pages.get(&page_index) {
+            // Cache the page
+            self.last_page[1] = self.last_page[0].clone();
+            self.last_page[0] = (page_index, Some(Rc::clone(page)));
+
+            Some(Rc::clone(page))
+        } else {
+            None
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {}

crates/mipsevm/src/page.rs

@@ -3,11 +3,13 @@
 
 use crate::utils::concat_fixed;
 use alloy_primitives::{keccak256, B256};
+use anyhow::Result;
 use once_cell::sync::Lazy;
 
 pub(crate) const PAGE_ADDRESS_SIZE: usize = 12;
 pub(crate) const PAGE_KEY_SIZE: usize = 32 - PAGE_ADDRESS_SIZE;
 pub(crate) const PAGE_SIZE: usize = 1 << PAGE_ADDRESS_SIZE;
+pub(crate) const PAGE_SIZE_WORDS: usize = PAGE_SIZE >> 5;
 pub(crate) const PAGE_ADDRESS_MASK: usize = PAGE_SIZE - 1;
 pub(crate) const MAX_PAGE_COUNT: usize = 1 << PAGE_KEY_SIZE;
 pub(crate) const PAGE_KEY_MASK: usize = MAX_PAGE_COUNT - 1;
@@ -25,29 +27,30 @@ pub(crate) static ZERO_HASHES: Lazy<[B256; 256]> = Lazy::new(|| {
 pub type Page = [u8; PAGE_SIZE];
 
 /// A [CachedPage] is a [Page] with an in-memory cache of intermediate nodes.
+#[derive(Debug, Clone, Copy)]
 pub struct CachedPage {
-    data: Page,
+    pub data: Page,
     /// Storage for intermediate nodes
-    cache: [[u8; 32]; PAGE_SIZE >> 5],
+    pub cache: [[u8; 32]; PAGE_SIZE_WORDS],
     /// Maps to true if the node is valid
     /// TODO(clabby): Use a bitmap / roaring bitmap
-    valid: [bool; PAGE_SIZE >> 5],
+    pub valid: [bool; PAGE_SIZE_WORDS],
 }
 
 impl Default for CachedPage {
     fn default() -> Self {
         Self {
             data: [0; PAGE_SIZE],
-            cache: [[0; 32]; PAGE_SIZE >> 5],
-            valid: [false; PAGE_SIZE >> 5],
+            cache: [[0; 32]; PAGE_SIZE_WORDS],
+            valid: [false; PAGE_SIZE_WORDS],
         }
     }
 }
 
 impl CachedPage {
-    pub fn invalidate(&mut self, page_addr: u32) {
-        if page_addr >= PAGE_SIZE as u32 {
-            panic!("Invalid page address");
+    pub fn invalidate(&mut self, page_addr: u64) -> Result<()> {
+        if page_addr >= PAGE_SIZE as u64 {
+            anyhow::bail!("Invalid page address: {}", page_addr);
         }
 
         // The first cache layer caches nodes that have two 32 byte leaf nodes.
@@ -62,6 +65,8 @@ impl CachedPage {
             self.valid[key as usize] = false;
             key >>= 1;
         }
+
+        Ok(())
     }
 
     pub fn invalidate_full(&mut self) {
@@ -71,7 +76,7 @@ impl CachedPage {
     pub fn merkle_root(&mut self) -> B256 {
         // First, hash the bottom layer.
         for i in (0..PAGE_SIZE).step_by(64) {
-            let j = (PAGE_SIZE >> 6) + (i >> 6);
+            let j = (PAGE_SIZE_WORDS >> 1) + (i >> 6);
             if self.valid[j] {
                 continue;
             }
@@ -81,7 +86,7 @@ impl CachedPage {
         }
 
         // Then, hash the cache layers.
-        for i in (1..=(PAGE_SIZE >> 5) - 2).rev().step_by(2) {
+        for i in (1..=PAGE_SIZE_WORDS - 2).rev().step_by(2) {
             let j = i >> 1;
             if self.valid[j] {
                 continue;
@@ -93,20 +98,22 @@ impl CachedPage {
         self.cache[1].into()
     }
 
-    pub fn merklize_subtree(&mut self, g_index: usize) -> B256 {
+    pub fn merklize_subtree(&mut self, g_index: usize) -> Result<B256> {
         // Fill the cache by computing the merkle root.
         let _ = self.merkle_root();
 
-        if g_index >= PAGE_SIZE >> 5 {
-            if g_index >= (PAGE_SIZE >> 5) * 2 {
-                panic!("Gindex too deep");
+        if g_index >= PAGE_SIZE_WORDS {
+            if g_index >= PAGE_SIZE_WORDS * 2 {
+                anyhow::bail!("Generalized index is too deep: {}", g_index);
             }
 
             let node_index = g_index & (PAGE_ADDRESS_MASK >> 5);
-            return B256::from_slice(&self.data[node_index << 5..(node_index << 5) + 32]);
+            return Ok(B256::from_slice(
+                &self.data[node_index << 5..(node_index << 5) + 32],
+            ));
         }
 
-        self.cache[g_index].into()
+        Ok(self.cache[g_index].into())
     }
 }
 
@@ -120,16 +127,16 @@ mod test {
         page.data[42] = 0xab;
 
         let g_index = ((1 << PAGE_ADDRESS_SIZE) | 42) >> 5;
-        let node = page.merklize_subtree(g_index);
+        let node = page.merklize_subtree(g_index).unwrap();
         let mut expected_leaf = B256::ZERO;
         expected_leaf[10] = 0xab;
         assert_eq!(node, expected_leaf, "Leaf nodes should not be hashed");
 
-        let node = page.merklize_subtree(g_index >> 1);
+        let node = page.merklize_subtree(g_index >> 1).unwrap();
         let expected_parent = keccak256(concat_fixed(ZERO_HASHES[0].into(), expected_leaf.into()));
         assert_eq!(node, expected_parent, "Parent should be correct");
 
-        let node = page.merklize_subtree(g_index >> 2);
+        let node = page.merklize_subtree(g_index >> 2).unwrap();
         let expected_grandparent =
             keccak256(concat_fixed(expected_parent.into(), ZERO_HASHES[1].into()));
         assert_eq!(node, expected_grandparent, "Grandparent should be correct");
@@ -142,22 +149,22 @@ mod test {
             "Pre and post state should be equal until the cache is invalidated"
         );
 
-        page.invalidate(42);
+        page.invalidate(42).unwrap();
         let post_b = page.merkle_root();
         assert_ne!(
             post, post_b,
             "Pre and post state should be different after cache invalidation"
         );
 
         page.data[2000] = 0xef;
-        page.invalidate(42);
+        page.invalidate(42).unwrap();
         let post_c = page.merkle_root();
         assert_eq!(
             post_b, post_c,
             "Local invalidation is not global invalidation."
         );
 
-        page.invalidate(2000);
+        page.invalidate(2000).unwrap();
         let post_d = page.merkle_root();
         assert_ne!(
             post_c, post_d,