Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

feat(mpt): TrieNode retrieval #173

Merged
merged 1 commit into from
May 8, 2024
Merged

feat(mpt): TrieNode retrieval #173

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
149 changes: 145 additions & 4 deletions crates/mpt/src/node.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
use alloc::{boxed::Box, vec, vec::Vec};
use alloy_primitives::{keccak256, Bytes, B256};
use alloy_rlp::{Buf, BufMut, Decodable, Encodable, Header, EMPTY_STRING_CODE};
use alloy_trie::Nibbles;
use anyhow::{anyhow, Result};

/// The length of the branch list when RLP encoded
Expand All @@ -12,6 +13,9 @@ const BRANCH_LIST_LENGTH: usize = 17;
/// The length of a leaf or extension node's RLP encoded list
const LEAF_OR_EXTENSION_LIST_LENGTH: usize = 2;

/// The number of nibbles traversed in a branch node.
const BRANCH_NODE_NIBBLES: usize = 1;

/// Prefix for even-nibbled extension node paths.
const PREFIX_EXTENSION_EVEN: u8 = 0;

Expand All @@ -24,6 +28,9 @@ const PREFIX_LEAF_EVEN: u8 = 2;
/// Prefix for odd-nibbled leaf node paths.
const PREFIX_LEAF_ODD: u8 = 3;

/// Nibble bit width.
const NIBBLE_WIDTH: usize = 4;

/// A [TrieNode] is a node within a standard Ethereum Merkle Patricia Trie.
///
/// The [TrieNode] has several variants:
Expand Down Expand Up @@ -85,7 +92,7 @@ impl TrieNode {
let path = Bytes::decode(buf).map_err(|e| anyhow!("Failed to decode: {e}"))?;

// Check the high-order nibble of the path to determine the type of node.
match path[0] >> 4 {
match path[0] >> NIBBLE_WIDTH {
PREFIX_EXTENSION_EVEN | PREFIX_EXTENSION_ODD => {
// extension node
let extension_node_value =
Expand Down Expand Up @@ -114,6 +121,103 @@ impl TrieNode {
self
}
}

/// Walks down the trie to a leaf value with the given key, if it exists. Preimages for blinded
/// nodes along the path are fetched using the `fetcher` function, and persisted in the inner
/// [TrieNode] elements.
///
/// ## Takes
/// - `self` - The root trie node
/// - `path` - The nibbles representation of the path to the leaf node
/// - `nibble_offset` - The number of nibbles that have already been traversed in the `item_key`
/// - `fetcher` - The preimage fetcher for intermediate blinded nodes
///
/// ## Returns
/// - `Err(_)` - Could not retrieve the node with the given key from the trie.
/// - `Ok((_, _))` - The key and value of the node
pub fn open<'a>(
&'a mut self,
path: &Nibbles,
mut nibble_offset: usize,
fetcher: impl Fn(B256) -> Result<Bytes> + Copy,
) -> Result<&'a mut Bytes> {
match self {
TrieNode::Branch { ref mut stack } => {
let branch_nibble = path[nibble_offset] as usize;
nibble_offset += BRANCH_NODE_NIBBLES;

let branch_node = stack
.get_mut(branch_nibble)
.ok_or(anyhow!("Key does not exist in trie (branch element not found)"))?;
match branch_node {
TrieNode::Empty => {
anyhow::bail!("Key does not exist in trie (empty node in branch)")
}
TrieNode::Blinded { commitment } => {
// If the string is a hash, we need to grab the preimage for it and
// continue recursing.
let trie_node = TrieNode::decode(&mut fetcher(*commitment)?.as_ref())
.map_err(|e| anyhow!(e))?;
*branch_node = trie_node;

// If the value was found in the blinded node, return it.
branch_node.open(path, nibble_offset, fetcher)
}
node => {
// If the value was found in the blinded node, return it.
node.open(path, nibble_offset, fetcher)
}
}
}
TrieNode::Leaf { key, value } => {
let key_nibbles = Nibbles::unpack(key.clone());
let shared_nibbles = key_nibbles[1..].as_ref();

// If the key length is one, it only contains the prefix and no shared nibbles.
// Return the key and value.
if key.len() == 1 || nibble_offset + shared_nibbles.len() >= path.len() {
return Ok(value);
}

let item_key_nibbles =
path[nibble_offset..nibble_offset + shared_nibbles.len()].as_ref();

if item_key_nibbles == shared_nibbles {
Ok(value)
} else {
anyhow::bail!("Key does not exist in trie (leaf doesn't share nibbles)");
}
}
TrieNode::Extension { prefix, node } => {
let prefix_nibbles = Nibbles::unpack(prefix);
let shared_nibbles = prefix_nibbles[1..].as_ref();
let item_key_nibbles =
path[nibble_offset..nibble_offset + shared_nibbles.len()].as_ref();
if item_key_nibbles == shared_nibbles {
// Increase the offset within the key by the length of the shared nibbles
nibble_offset += shared_nibbles.len();

// Follow extension branch
if let TrieNode::Blinded { commitment } = node.as_ref() {
*node = Box::new(
TrieNode::decode(&mut fetcher(*commitment)?.as_ref())
.map_err(|e| anyhow!(e))?,
);
}
node.open(path, nibble_offset, fetcher)
} else {
anyhow::bail!("Key does not exist in trie (extension doesn't share nibbles)");
}
}
TrieNode::Blinded { commitment } => {
let trie_node = TrieNode::decode(&mut fetcher(*commitment)?.as_ref())
.map_err(|e| anyhow!(e))?;
*self = trie_node;
self.open(path, nibble_offset, fetcher)
}
_ => anyhow::bail!("Invalid trie node type encountered"),
}
}
}

impl Encodable for TrieNode {
Expand Down Expand Up @@ -260,8 +364,12 @@ fn rlp_list_element_length(buf: &mut &[u8]) -> alloy_rlp::Result<usize> {
#[cfg(test)]
mod test {
use super::*;
use alloc::vec;
use alloy_primitives::{b256, bytes, hex};
use crate::{test_util::ordered_trie_with_encoder, TrieNode};
use alloc::{collections::BTreeMap, vec, vec::Vec};
use alloy_primitives::{b256, bytes, hex, keccak256, Bytes, B256};
use alloy_rlp::{Decodable, Encodable, EMPTY_STRING_CODE};
use alloy_trie::Nibbles;
use anyhow::{anyhow, Result};

#[test]
fn test_decode_branch() {
Expand Down Expand Up @@ -320,7 +428,7 @@ mod test {
hex!("e58300646fa0f3fe8b3c5b21d3e52860f1e4a5825a6100bb341069c1e88f4ebf6bd98de0c190");
let mut rlp_buf = Vec::new();

let opened = TrieNode::Leaf { key: bytes!("30"), value: bytes!("FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF") };
let opened = TrieNode::Leaf { key: bytes!("30"), value: [0xFF; 64].into() };
opened.encode(&mut rlp_buf);
let blinded = TrieNode::Blinded { commitment: keccak256(&rlp_buf) };

Expand All @@ -339,4 +447,37 @@ mod test {
let expected = TrieNode::Leaf { key: bytes!("20646f"), value: bytes!("76657262FF") };
assert_eq!(expected, TrieNode::decode(&mut LEAF_RLP.as_slice()).unwrap());
}

#[test]
fn test_retrieve_from_trie_simple() {
const VALUES: [&str; 5] = ["yeah", "dog", ", ", "laminar", "flow"];

let mut trie = ordered_trie_with_encoder(&VALUES, |v, buf| v.encode(buf));
let root = trie.root();

let preimages =
trie.take_proofs().into_iter().fold(BTreeMap::default(), |mut acc, (_, value)| {
acc.insert(keccak256(value.as_ref()), value);
acc
});
let fetcher = |h: B256| -> Result<Bytes> {
preimages.get(&h).cloned().ok_or(anyhow!("Failed to find preimage"))
};

let mut root_node = TrieNode::decode(&mut fetcher(root).unwrap().as_ref()).unwrap();
for (i, value) in VALUES.iter().enumerate() {
let path_nibbles = Nibbles::unpack([if i == 0 { EMPTY_STRING_CODE } else { i as u8 }]);
let v = root_node.open(&path_nibbles, 0, fetcher).unwrap();

let mut encoded_value = Vec::with_capacity(value.length());
value.encode(&mut encoded_value);

assert_eq!(v, encoded_value.as_slice());
}

let TrieNode::Blinded { commitment } = root_node.blind() else {
panic!("Expected blinded root node");
};
assert_eq!(commitment, root);
}
}
Loading