✨ Add merklize_subtree for Memory

crates/mipsevm/src/memory.rs

@@ -1,7 +1,10 @@
 //! The memory module contains the memory data structures and functionality for the emulator.
 
-use crate::page::{self, CachedPage};
-use alloy_primitives::B256;
+use crate::{
+    page::{self, CachedPage},
+    utils::concat_fixed,
+};
+use alloy_primitives::{keccak256, B256};
 use anyhow::Result;
 use std::{cell::RefCell, collections::BTreeMap, rc::Rc};
 
@@ -89,6 +92,47 @@ impl Memory {
             None
         }
     }
+
+    fn merklize_subtree(&mut self, g_index: u64) -> Result<B256> {
+        // Fetch the amount of bits required to represent the generalized index
+        let bits = 128 - g_index.leading_zeros();
+        if bits > 28 {
+            anyhow::bail!("Gindex is too deep")
+        }
+
+        if bits > page::PAGE_KEY_SIZE as u32 {
+            let depth_into_page = bits - 1 - page::PAGE_KEY_SIZE as u32;
+            let page_index = (g_index >> depth_into_page) & page::PAGE_KEY_MASK as u64;
+            return self.pages.get(&page_index).map_or(
+                Ok(page::ZERO_HASHES[28 - bits as usize]),
+                |page| {
+                    let page_g_index =
+                        (1 << depth_into_page) | (g_index & ((1 << depth_into_page) - 1));
+                    page.borrow_mut().merklize_subtree(page_g_index as usize)
+                },
+            );
+        }
+
+        if bits > page::PAGE_KEY_SIZE as u32 + 1 {
+            anyhow::bail!("Cannot jump into intermediate node of page")
+        }
+
+        if let Some(node) = self.nodes.get(&g_index) {
+            if let Some(node) = node {
+                return Ok(*node);
+            }
+        } else {
+            return Ok(page::ZERO_HASHES[28 - bits as usize]);
+        }
+
+        let left = self.merklize_subtree(g_index << 1)?;
+        let right = self.merklize_subtree((g_index << 1) | 1)?;
+        let result = keccak256(concat_fixed(left.into(), right.into()));
+
+        self.nodes.insert(g_index, Some(result));
+
+        Ok(result)
+    }
 }
 
 #[cfg(test)]