♻️ Make Address 32 bits

crates/mipsevm/src/memory.rs

@@ -61,13 +61,13 @@ impl Memory {
         }
 
         // Find the page and invalidate the address within it.
-        match self.page_lookup(address >> page::PAGE_ADDRESS_SIZE) {
+        match self.page_lookup(address as u64 >> page::PAGE_ADDRESS_SIZE) {
             Some(page) => {
                 let mut page = page.borrow_mut();
                 let prev_valid = !page.is_valid(1);
 
                 // Invalidate the address within the page.
-                page.invalidate(address & page::PAGE_ADDRESS_MASK as u64)?;
+                page.invalidate(address & page::PAGE_ADDRESS_MASK as u32)?;
 
                 // If the page was already invalid before, then nodes to the memory
                 // root will also still be invalid.
@@ -82,7 +82,7 @@ impl Memory {
         }
 
         // Find the generalized index of the first page covering the address
-        let mut g_index = ((1u64 << 32) | address) >> page::PAGE_ADDRESS_SIZE as u64;
+        let mut g_index = ((1u64 << 32) | address as u64) >> page::PAGE_ADDRESS_SIZE;
         // Invalidate all nodes in the branch
         while g_index > 0 {
             self.nodes.insert(g_index, None);
@@ -236,7 +236,7 @@ impl Memory {
             anyhow::bail!("Unaligned memory access: {:x}", address);
         }
 
-        let page_index = address >> page::PAGE_ADDRESS_SIZE as u64;
+        let page_index = address as PageIndex >> page::PAGE_ADDRESS_SIZE as u64;
         let page_address = address as usize & page::PAGE_ADDRESS_MASK;
 
         // Attempt to look up the page.
@@ -271,7 +271,7 @@ impl Memory {
             anyhow::bail!("Unaligned memory access: {:x}", address);
         }
 
-        match self.page_lookup(address >> page::PAGE_ADDRESS_SIZE as u64) {
+        match self.page_lookup(address as u64 >> page::PAGE_ADDRESS_SIZE as u64) {
             Some(page) => {
                 let page_address = address as usize & page::PAGE_ADDRESS_MASK;
                 Ok(u32::from_be_bytes(
@@ -313,7 +313,7 @@ impl Memory {
         let mut address = address;
         let mut data = data;
         loop {
-            let page_index = address >> page::PAGE_ADDRESS_SIZE as u64;
+            let page_index = address as PageIndex >> page::PAGE_ADDRESS_SIZE as u64;
             let page_address = address as usize & page::PAGE_ADDRESS_MASK;
 
             let page = self
@@ -327,7 +327,7 @@ impl Memory {
                     if n == 0 {
                         return Ok(());
                     }
-                    address += n as u64;
+                    address += n as u32;
                 }
                 Err(e) => return Err(e.into()),
             };
@@ -338,11 +338,11 @@ impl Memory {
 pub struct MemoryReader {
     memory: Rc<RefCell<Memory>>,
     address: Address,
-    count: u64,
+    count: u32,
 }
 
 impl MemoryReader {
-    pub fn new(memory: Rc<RefCell<Memory>>, address: Address, count: u64) -> Self {
+    pub fn new(memory: Rc<RefCell<Memory>>, address: Address, count: u32) -> Self {
         Self {
             memory,
             address,
@@ -357,13 +357,13 @@ impl Read for MemoryReader {
             return Ok(0);
         }
 
-        let end_address = self.address + self.count;
+        let end_address = self.address + self.count as Address;
 
-        let page_index = self.address >> page::PAGE_ADDRESS_SIZE as u64;
+        let page_index = self.address as PageIndex >> page::PAGE_ADDRESS_SIZE as u64;
         let start = self.address as usize & page::PAGE_ADDRESS_MASK;
         let mut end = page::PAGE_SIZE;
 
-        if page_index == (end_address >> page::PAGE_ADDRESS_SIZE as u64) {
+        if page_index == (end_address as u64 >> page::PAGE_ADDRESS_SIZE as u64) {
             end = end_address as usize & page::PAGE_ADDRESS_MASK;
         }
         let n = end - start;
@@ -377,8 +377,8 @@ impl Read for MemoryReader {
                 std::io::copy(&mut vec![0; n].as_slice(), &mut buf)?;
             }
         };
-        self.address += n as u64;
-        self.count -= n as u64;
+        self.address += n as u32;
+        self.count -= n as u32;
         Ok(n)
     }
 }
@@ -578,7 +578,7 @@ mod test {
                 .expect("Should not error");
 
             let mut reader =
-                MemoryReader::new(Rc::clone(&memory), 0x1337 - 10, data.len() as u64 + 20);
+                MemoryReader::new(Rc::clone(&memory), 0x1337 - 10, data.len() as u32 + 20);
             let mut buf = Vec::with_capacity(1260);
             reader.read_to_end(&mut buf).unwrap();
 

crates/mipsevm/src/mips/instrumented.rs

@@ -66,7 +66,7 @@ where
     /// - Err(_): An error occurred while processing the instruction step in the MIPS emulator.
     pub fn step(&mut self, proof: bool) -> Result<Option<StepWitness>> {
         self.mem_proof_enabled = proof;
-        self.last_mem_access = !0u32 as u64;
+        self.last_mem_access = !0u32 as Address;
         self.last_preimage_offset = !0u32;
 
         let mut witness = None;
@@ -120,7 +120,6 @@ mod test {
         };
 
         #[test]
-        // #[ignore]
         fn open_mips_tests() {
             let tests_path = PathBuf::from(std::env::current_dir().unwrap())
                 .join("open_mips_tests")

crates/mipsevm/src/mips/mips_vm.rs

@@ -31,8 +31,8 @@ where
     /// - `Err(_)`: An error occurred while fetching the preimage.
     pub(crate) fn read_preimage(&mut self, key: B256, offset: u32) -> Result<(B256, usize)> {
         if key != self.last_preimage_key {
-            self.last_preimage_key = key;
             let data = self.preimage_oracle.get(key)?;
+            self.last_preimage_key = key;
 
             // Add the length prefix to the preimage
             // Resizes the `last_preimage` vec in-place to reduce reallocations.
@@ -60,7 +60,7 @@ where
     /// - A [Result] indicating if the operation was successful.
     pub(crate) fn track_mem_access(&mut self, effective_address: Address) -> Result<()> {
         if self.mem_proof_enabled && self.last_mem_access != effective_address {
-            if self.last_mem_access as u32 != u32::MAX {
+            if self.last_mem_access != Address::MAX {
                 anyhow::bail!("Unexpected diffrent memory access at {:x}, already have access at {:x} buffered", effective_address, self.last_mem_access);
             }
 
@@ -214,7 +214,7 @@ where
         let mut v0 = 0;
         let mut v1 = 0;
 
-        let (a0, a1, a2) = (
+        let (a0, a1, mut a2) = (
             self.state.registers[4],
             self.state.registers[5],
             self.state.registers[6],
@@ -256,6 +256,7 @@ where
                         // Nothing to do; Leave v0 and v1 zero, read nothing, and give no error.
                     }
                     Ok(Fd::PreimageRead) => {
+                        dbg!("moose");
                         let effective_address = (a1 & 0xFFFFFFFC) as Address;
 
                         self.track_mem_access(effective_address)?;
@@ -298,11 +299,8 @@ where
                 },
                 Syscall::Write => match (a0 as u8).try_into() {
                     Ok(fd @ (Fd::Stdout | Fd::StdErr)) => {
-                        let mut reader = MemoryReader::new(
-                            Rc::clone(&self.state.memory),
-                            a1 as Address,
-                            a2 as u64,
-                        );
+                        let mut reader =
+                            MemoryReader::new(Rc::clone(&self.state.memory), a1 as Address, a2);
                         let writer: &mut dyn Write = if matches!(fd, Fd::Stdout) {
                             &mut self.std_out
                         } else {
@@ -312,14 +310,11 @@ where
                         v0 = a2;
                     }
                     Ok(Fd::HintWrite) => {
-                        let mut reader = MemoryReader::new(
-                            Rc::clone(&self.state.memory),
-                            a1 as Address,
-                            a2 as u64,
-                        );
+                        let mut reader =
+                            MemoryReader::new(Rc::clone(&self.state.memory), a1 as Address, a2);
                         let mut hint_data = Vec::with_capacity(a2 as usize);
                         reader.read_to_end(&mut hint_data)?;
-                        self.state.last_hint.extend_from_slice(hint_data.as_slice());
+                        self.state.last_hint.extend(hint_data);
 
                         // Continue processing while there is enough data to check if there are any
                         // hints.
@@ -351,7 +346,6 @@ where
                         let alignment = a1 & 0x3;
                         let space = 4 - alignment;
 
-                        let mut a2 = a2;
                         if space < a2 {
                             a2 = space;
                         }

crates/mipsevm/src/page.rs

@@ -51,7 +51,7 @@ impl CachedPage {
     /// ### Returns
     /// - A [Result] indicating if the operation was successful.
     pub fn invalidate(&mut self, page_addr: Address) -> Result<()> {
-        if page_addr >= PAGE_SIZE as u64 {
+        if page_addr >= PAGE_SIZE as Address {
             anyhow::bail!("Invalid page address: {}", page_addr);
         }
 

crates/mipsevm/src/types.rs

@@ -12,8 +12,8 @@ pub type PageIndex = u64;
 /// A [Gindex] is a generalized index, defined as $2^{\text{depth}} + \text{index}$.
 pub type Gindex = u64;
 
-/// An [Address] is a 64 bit address in the MIPS emulator's memory.
-pub type Address = u64;
+/// An [Address] is a 32 bit address in the MIPS emulator's memory.
+pub type Address = u32;
 
 /// The [VMStatus] is an indicator within the [StateWitness] hash that indicates
 /// the current status of the MIPS emulator.