🚧 Memory WIP

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,91 @@
+//! The memory module contains the memory data structures and functionality for the emulator.
+
+use crate::page::{self, CachedPage};
+use alloy_primitives::B256;
+use std::{cell::RefCell, collections::BTreeMap, rc::Rc};
+
+type PageIndex = u32;
+
+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: [(!0u32, None), (!0u32, 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: u32) {
+        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 u32);
+            if !page.valid[1] {
+                return;
+            }
+        } else {
+            // Nothing to invalidate
+            return;
+        }
+
+        // Find the generalized index of the first page covering the address
+        let mut g_index = ((1 << 32) | address) >> page::PAGE_ADDRESS_SIZE as u32;
+        while g_index > 0 {
+            self.nodes.insert(g_index.into(), None);
+            g_index >>= 1;
+        }
+    }
+
+    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) {
+            return Some(Rc::clone(page));
+        } else if let Some(page) = self.pages.get(&(page_index as u64)) {
+            // 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

@@ -25,13 +25,14 @@ 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 >> 5],
     /// 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 >> 5],
 }
 
 impl Default for CachedPage {