perf: intermediate hashing

crates/mipsevm/benches/memory.rs

@@ -3,7 +3,10 @@ use criterion::{criterion_group, criterion_main, Criterion};
 use rand::RngCore;
 
 fn merkle_root(c: &mut Criterion) {
-    c.bench_function("Merkle Root (memory size = 25 MB)", |b| {
+    let mut g = c.benchmark_group("memory");
+    g.sample_size(10);
+
+    g.bench_function("Merkle Root (memory size = 25 MB)", |b| {
         let mut memory = Memory::default();
         let mut data = vec![0u8; 25_000_000];
         rand::thread_rng().fill_bytes(&mut data[..]);
@@ -15,7 +18,7 @@ fn merkle_root(c: &mut Criterion) {
         });
     });
 
-    c.bench_function("Merkle Root (memory size = 50 MB)", |b| {
+    g.bench_function("Merkle Root (memory size = 50 MB)", |b| {
         let mut memory = Memory::default();
         let mut data = vec![0u8; 50_000_000];
         rand::thread_rng().fill_bytes(&mut data[..]);
@@ -27,7 +30,7 @@ fn merkle_root(c: &mut Criterion) {
         });
     });
 
-    c.bench_function("Merkle Root (memory size = 100 MB)", |b| {
+    g.bench_function("Merkle Root (memory size = 100 MB)", |b| {
         let mut memory = Memory::default();
         let mut data = vec![0u8; 100_000_000];
         rand::thread_rng().fill_bytes(&mut data[..]);
@@ -39,7 +42,7 @@ fn merkle_root(c: &mut Criterion) {
         });
     });
 
-    c.bench_function("Merkle Root (memory size = 200 MB)", |b| {
+    g.bench_function("Merkle Root (memory size = 200 MB)", |b| {
         let mut memory = Memory::default();
         let mut data = vec![0u8; 200_000_000];
         rand::thread_rng().fill_bytes(&mut data[..]);

crates/mipsevm/src/memory.rs

@@ -2,7 +2,7 @@
 
 use crate::{
     page::{self, CachedPage},
-    utils::keccak_concat_fixed,
+    utils::keccak_concat_fixed_inplace,
     Address, Gindex, PageIndex,
 };
 use alloy_primitives::B256;
@@ -149,7 +149,8 @@ impl Memory {
 
         let left = self.merkleize_subtree(g_index << 1)?;
         let right = self.merkleize_subtree((g_index << 1) | 1)?;
-        let result = keccak_concat_fixed(*left, *right);
+        let mut buf = [0u8; 64];
+        let result = keccak_concat_fixed_inplace(&left, &right, &mut buf);
 
         self.nodes.insert(g_index, Some(result));
 
@@ -172,15 +173,7 @@ impl Memory {
     /// ### Returns
     /// - The 896 bit merkle proof for the given address.
     pub fn merkle_proof(&mut self, address: Address) -> Result<[u8; 28 << 5]> {
-        let proof = self.traverse_branch(1, address, 0)?;
-        let mut proof_out = [0u8; 28 << 5];
-
-        // Encode the proof
-        (0..28).for_each(|i| {
-            let start = i << 5;
-            proof_out[start..start + 32].copy_from_slice(proof[i].as_slice());
-        });
-
+        let (_, proof_out) = self.traverse_branch(1, address, 0)?;
         Ok(proof_out)
     }
 
@@ -198,11 +191,11 @@ impl Memory {
         parent: Gindex,
         address: Address,
         depth: u8,
-    ) -> Result<Vec<B256>> {
+    ) -> Result<(usize, [u8; 28 * 32])> {
         if depth == 32 - 5 {
-            let mut proof = Vec::with_capacity(32 - 5 + 1);
-            proof.push(self.merkleize_subtree(parent)?);
-            return Ok(proof);
+            let mut proof = [0u8; 28 * 32];
+            proof[..32].copy_from_slice(&*self.merkleize_subtree(parent)?);
+            return Ok((32, proof));
         }
 
         if depth > 32 - 5 {
@@ -215,10 +208,10 @@ impl Memory {
             (local, sibling) = (sibling, local);
         }
 
-        let mut proof = self.traverse_branch(local, address, depth + 1)?;
+        let (free_ptr, mut proof) = self.traverse_branch(local, address, depth + 1)?;
         let sibling_node = self.merkleize_subtree(sibling)?;
-        proof.push(sibling_node);
-        Ok(proof)
+        proof[free_ptr..free_ptr + 32].copy_from_slice(&*sibling_node);
+        Ok((free_ptr + 32, proof))
     }
 
     /// Set a 32 bit value in the [Memory] at a given address.

crates/mipsevm/src/mips/instrumented.rs

@@ -280,7 +280,7 @@ mod test {
             if ins.state.exited {
                 break;
             }
-            ins.step(false).unwrap();
+            ins.step(true).unwrap();
         }
 
         assert!(ins.state.exited, "must exit");

crates/mipsevm/src/page.rs

@@ -1,6 +1,9 @@
 //! This module contains the data structure for a [Page] within the MIPS emulator's [Memory].
 
-use crate::{utils::keccak_concat_fixed, Address, Gindex, Page};
+use crate::{
+    utils::{keccak_concat_fixed, keccak_concat_fixed_inplace},
+    Address, Gindex, Page,
+};
 use alloy_primitives::{keccak256, B256};
 use anyhow::Result;
 use once_cell::sync::Lazy;
@@ -96,34 +99,61 @@ impl CachedPage {
             if self.is_valid(j) {
                 continue;
             }
-            self.cache[j] = *keccak_concat_fixed(self.cache[i], self.cache[i + 1]);
+            let mut buf = [0u8; 64];
+            let (a, b) = self.cache.split_at_mut(i);
+            a[j] = *keccak_concat_fixed_inplace(&b[0], &b[1], &mut buf);
             self.set_valid(j, true);
         }
 
         self.cache[1].into()
     }
 
-    pub fn merkleize_subtree(&mut self, g_index: Gindex) -> Result<B256> {
-        // Fill the cache by computing the merkle root.
-        let _ = self.merkle_root();
+    fn fill_cache(&mut self, g_index: usize) {
+        if self.is_valid(g_index) {
+            return;
+        }
 
-        if g_index >= PAGE_SIZE_WORDS as u64 {
-            if g_index >= (PAGE_SIZE_WORDS * 2) as u64 {
-                anyhow::bail!("Generalized index is too deep: {}", g_index);
-            }
+        if g_index >= PAGE_SIZE_WORDS >> 1 {
+            // This is a leaf node.
+            let data_idx = (g_index - (PAGE_SIZE_WORDS >> 1)) << 6;
+            self.cache[g_index] = *keccak256(&self.data[data_idx..data_idx + 64]);
+        } else {
+            // This is an internal node.
+            let left_child = g_index << 1;
+            let right_child = left_child + 1;
+
+            // Ensure children are hashed.
+            self.fill_cache(left_child);
+            self.fill_cache(right_child);
+
+            let mut buf = [0u8; 64];
+            self.cache[g_index] = *keccak_concat_fixed_inplace(
+                &self.cache[left_child],
+                &self.cache[right_child],
+                &mut buf,
+            );
+        }
+        self.set_valid(g_index, true);
+    }
 
+    pub fn merkleize_subtree(&mut self, g_index: Gindex) -> Result<B256> {
+        if (PAGE_SIZE_WORDS..PAGE_SIZE_WORDS * 2).contains(&(g_index as usize)) {
             let node_index = g_index as usize & (PAGE_ADDRESS_MASK >> 5);
             let start = node_index << 5;
             return Ok(B256::from_slice(&self.data[start..start + 32]));
+        } else if g_index as usize >= PAGE_SIZE_WORDS * 2 {
+            anyhow::bail!("Generalized index is too deep: {}", g_index);
         }
 
+        self.fill_cache(g_index as usize);
         Ok(self.cache[g_index as usize].into())
     }
 
     /// Check if a key is valid within the bitmap.
     ///
     /// ### Takes
     /// - `key`: The key to check.
+    #[inline(always)]
     pub fn is_valid(&self, key: usize) -> bool {
         let flag = 1 << (127 - key);
         self.valid & flag == flag
@@ -134,6 +164,7 @@ impl CachedPage {
     /// ### Takes
     /// - `key`: The key to set.
     /// - `valid`: Whether the key should be set as valid or invalid.
+    #[inline(always)]
     pub fn set_valid(&mut self, key: usize, valid: bool) {
         let flag_offset = 127 - key;
         self.valid &= !(1 << flag_offset);

crates/mipsevm/src/utils.rs

@@ -22,3 +22,30 @@ where
 {
     keccak256(concat_fixed(a, b))
 }
+
+#[inline(always)]
+pub(crate) fn concat_fixed_inplace<T, const N: usize, const M: usize>(
+    a: &[T; N],
+    b: &[T; M],
+    concatenated: &mut [T; N + M],
+) where
+    T: Copy,
+{
+    let (left, right) = concatenated.split_at_mut(N);
+    left.copy_from_slice(a);
+    right.copy_from_slice(b);
+}
+
+#[inline(always)]
+pub(crate) fn keccak_concat_fixed_inplace<T, const N: usize, const M: usize>(
+    a: &[T; N],
+    b: &[T; M],
+    buffer: &mut [T; N + M],
+) -> B256
+where
+    T: Copy,
+    [T; N + M]: AsRef<[u8]>,
+{
+    concat_fixed_inplace(a, b, buffer);
+    keccak256(buffer)
+}