-
Notifications
You must be signed in to change notification settings - Fork 29
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
8 changed files
with
260 additions
and
325 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,3 +1,3 @@ | ||
| # `kona-mpt` | ||
|
|
||
| Utilities for interacting with and iterating through a merkle patricia trie | ||
| Utilities for interacting with a merkle patricia trie in the client program. |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,48 @@ | ||
| //! This module contains the [TrieAccount] struct. | ||
|
|
||
| use alloy_consensus::constants::KECCAK_EMPTY; | ||
| use alloy_primitives::{B256, U256}; | ||
| use alloy_rlp::{RlpDecodable, RlpEncodable}; | ||
| use revm_primitives::{Account, AccountInfo}; | ||
|
|
||
| /// An Ethereum account as represented in the trie. | ||
| #[derive(Clone, Copy, Debug, PartialEq, Eq, Default, RlpEncodable, RlpDecodable)] | ||
| pub struct TrieAccount { | ||
| /// Account nonce. | ||
| nonce: u64, | ||
| /// Account balance. | ||
| balance: U256, | ||
| /// Account's storage root. | ||
| storage_root: B256, | ||
| /// Hash of the account's bytecode. | ||
| code_hash: B256, | ||
| } | ||
|
|
||
| impl From<(Account, B256)> for TrieAccount { | ||
| fn from((account, storage_root): (Account, B256)) -> Self { | ||
| Self { | ||
| nonce: account.info.nonce, | ||
| balance: account.info.balance, | ||
| storage_root, | ||
| code_hash: account.info.code_hash, | ||
| } | ||
| } | ||
| } | ||
|
|
||
| impl From<(AccountInfo, B256)> for TrieAccount { | ||
| fn from((account, storage_root): (AccountInfo, B256)) -> Self { | ||
| Self { | ||
| nonce: account.nonce, | ||
| balance: account.balance, | ||
| storage_root, | ||
| code_hash: account.code_hash, | ||
| } | ||
| } | ||
| } | ||
|
|
||
| impl TrieAccount { | ||
| /// Get account's storage root. | ||
| pub fn storage_root(&self) -> B256 { | ||
| self.storage_root | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,138 @@ | ||
| //! Contains the [retrieve] function, allowing for retrieving values from leaves within a Merkle | ||
| //! Patricia Trie by key. | ||
|
|
||
| use crate::{NodeElement, TrieNode}; | ||
| use alloy_primitives::{Bytes, B256}; | ||
| use alloy_rlp::Decodable; | ||
| use alloy_trie::Nibbles; | ||
| use anyhow::{anyhow, Result}; | ||
|
|
||
| /// 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. | ||
| /// | ||
| /// ## Takes | ||
| /// - `item_key` - The nibbles representation of the key being retrieved | ||
| /// - `trie_node` - The root trie node | ||
| /// - `pos` - 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 retrieve( | ||
| item_key: &Nibbles, | ||
| trie_node: TrieNode, | ||
| mut pos: usize, | ||
| fetcher: impl Fn(B256) -> Result<Bytes> + Copy, | ||
| ) -> Result<Bytes> { | ||
| match trie_node { | ||
| TrieNode::Branch { mut stack } => { | ||
| let branch_nibble = item_key[pos]; | ||
| pos += 1; | ||
|
|
||
| match stack | ||
| .remove(branch_nibble as usize) | ||
| .ok_or(anyhow!("Key does not exist in trie"))? | ||
| { | ||
| NodeElement::String(s) => { | ||
| // If the string is a hash, we need to grab the preimage for it and | ||
| // continue recursing. | ||
| let hash: B256 = | ||
| s.as_ref().try_into().map_err(|e| anyhow!("Conversion error: {e}"))?; | ||
| let trie_node = | ||
| TrieNode::decode(&mut fetcher(hash)?.as_ref()).map_err(|e| anyhow!(e))?; | ||
|
|
||
| // If the value was found in the blinded node, return it. | ||
| if let Ok(value) = retrieve(item_key, trie_node, pos, fetcher) { | ||
| return Ok(value); | ||
| } | ||
| } | ||
| list @ NodeElement::List(_) => { | ||
| let trie_node = list.try_list_into_node()?; | ||
|
|
||
| // If the value was found in the blinded node, return it. | ||
| if let Ok(value) = retrieve(item_key, trie_node, pos, fetcher) { | ||
| return Ok(value); | ||
| } | ||
| } | ||
| _ => { /* Skip over empty lists and strings; We're looking for leaves */ } | ||
| }; | ||
|
|
||
| anyhow::bail!("Key does not exist in trie"); | ||
| } | ||
| TrieNode::Leaf { key, value } => { | ||
| // If the key length is one, it only contains the prefix and no shared nibbles. Return | ||
| // the key and value. | ||
| if key.len() == 1 { | ||
| return Ok(value); | ||
| } | ||
|
|
||
| let key_nibbles = Nibbles::unpack(key.clone()); | ||
| let shared_nibbles = key_nibbles[1..].as_ref(); | ||
| let item_key_nibbles = item_key[pos..pos + shared_nibbles.len()].as_ref(); | ||
|
|
||
| if item_key_nibbles == shared_nibbles { | ||
| Ok(value) | ||
| } else { | ||
| anyhow::bail!("Key does not exist in trie"); | ||
| } | ||
| } | ||
| TrieNode::Extension { prefix, node } => { | ||
| let prefix_nibbles = Nibbles::unpack(prefix); | ||
| let shared_nibbles = prefix_nibbles[1..].as_ref(); | ||
| let item_key_nibbles = item_key[pos..pos + shared_nibbles.len()].as_ref(); | ||
| if item_key_nibbles == shared_nibbles { | ||
| // Increase the offset within the key by the length of the shared nibbles | ||
| pos += shared_nibbles.len(); | ||
|
|
||
| // Follow extension branch | ||
| let hash = B256::from_slice(node.as_ref()); | ||
| let extension_link = | ||
| TrieNode::decode(&mut fetcher(hash)?.as_ref()).map_err(|e| anyhow!(e))?; | ||
| retrieve(item_key, extension_link, pos, fetcher) | ||
| } else { | ||
| anyhow::bail!("Key does not exist in trie"); | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
||
| #[cfg(test)] | ||
| mod test { | ||
| use alloc::{collections::BTreeMap, vec::Vec}; | ||
| use alloy_primitives::{keccak256, Bytes, B256}; | ||
| use alloy_rlp::{Decodable, Encodable, EMPTY_STRING_CODE}; | ||
| use alloy_trie::Nibbles; | ||
| use anyhow::{anyhow, Result}; | ||
|
|
||
| use crate::{retrieve, test_util::ordered_trie_with_encoder, TrieNode}; | ||
|
|
||
| #[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 root = TrieNode::decode(&mut fetcher(root).unwrap().as_ref()).unwrap(); | ||
|
|
||
| for (i, value) in VALUES.iter().enumerate() { | ||
| let key_nibbles = Nibbles::unpack([if i == 0 { EMPTY_STRING_CODE } else { i as u8 }]); | ||
| let v = retrieve(&key_nibbles, root.clone(), 0, fetcher).unwrap(); | ||
|
|
||
| let mut encoded_value = Vec::with_capacity(value.length()); | ||
| value.encode(&mut encoded_value); | ||
|
|
||
| assert_eq!(v, encoded_value); | ||
| } | ||
| } | ||
| } |
Oops, something went wrong.