From 89be826f8f64b7fc566163505a72a277341aa6d4 Mon Sep 17 00:00:00 2001
From: clabby <ben@clab.by>
Date: Wed, 20 Sep 2023 16:51:31 -0400
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20Add=20`merklize=5Fsubtree`=20for=20?=
 =?UTF-8?q?`Memory`?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 crates/mipsevm/src/memory.rs | 48 ++++++++++++++++++++++++++++++++++--
 1 file changed, 46 insertions(+), 2 deletions(-)

diff --git a/crates/mipsevm/src/memory.rs b/crates/mipsevm/src/memory.rs
index 7d80f27..962bbf1 100644
--- a/crates/mipsevm/src/memory.rs
+++ b/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)]