Add get function

morph-dev · Oct 5, 2023 · 7cfd68f · 7cfd68f
1 parent 81ec064
commit 7cfd68f
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 0 deletions.
diff --git a/merkle/src/nodes/branch.rs b/merkle/src/nodes/branch.rs
@@ -23,6 +23,13 @@ impl BranchNode {
         nodes_length + value_length
     }
 
+    pub fn get(&self, path: &Nibbles, db: &dyn Db<B256, Vec<u8>>) -> Option<Vec<u8>> {
+        match path.first() {
+            None => self.value.clone(),
+            Some(nibble) => self.nodes[nibble].clone().get(&path.skip_head(1), db),
+        }
+    }
+
     pub fn update(&self, path: Nibbles, value: Vec<u8>, db: &mut dyn Db<B256, Vec<u8>>) -> Node {
         let mut nodes = self.nodes.clone();
 

diff --git a/merkle/src/nodes/extension.rs b/merkle/src/nodes/extension.rs
@@ -13,6 +13,15 @@ pub struct ExtensionNode {
 }
 
 impl ExtensionNode {
+    pub fn get(&self, path: &Nibbles, db: &dyn Db<B256, Vec<u8>>) -> Option<Vec<u8>> {
+        let common_prefix = self.path.common_prefix(path);
+        if self.path.len() == common_prefix {
+            self.node.clone().get(&path.skip_head(common_prefix), db)
+        } else {
+            None
+        }
+    }
+
     pub fn update(&self, path: Nibbles, value: Vec<u8>, db: &mut dyn Db<B256, Vec<u8>>) -> Node {
         let common_prefix = self.path.common_prefix(&path);
 

diff --git a/merkle/src/nodes/leaf.rs b/merkle/src/nodes/leaf.rs
@@ -17,6 +17,14 @@ pub struct LeafNode {
 }
 
 impl LeafNode {
+    pub fn get(&self, path: &Nibbles, _db: &dyn Db<B256, Vec<u8>>) -> Option<Vec<u8>> {
+        if self.path[..] == path[..] {
+            Some(self.value.clone())
+        } else {
+            None
+        }
+    }
+
     pub fn update(&self, path: Nibbles, value: Vec<u8>, db: &mut dyn Db<B256, Vec<u8>>) -> Node {
         let common_prefix = self.path.common_prefix(&path);
         // Path is the same

diff --git a/merkle/src/nodes/node.rs b/merkle/src/nodes/node.rs
@@ -116,6 +116,20 @@ impl From<Node> for NodeRef {
 }
 
 impl NodeRef {
+    pub fn get(&mut self, path: &Nibbles, db: &dyn Db<B256, Vec<u8>>) -> Option<Vec<u8>> {
+        self.load(db);
+
+        match &self.node {
+            None => panic!("Node should be present"),
+            Some(node) => match node.as_ref() {
+                Node::Nil => None,
+                Node::Branch(branch) => branch.get(path, db),
+                Node::Extension(extension) => extension.get(path, db),
+                Node::Leaf(leaf) => leaf.get(path, db),
+            },
+        }
+    }
+
     pub fn update(
         &mut self,
         path: Nibbles,