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

chore(mpt): Do not expose recursion vars #197

Merged
merged 1 commit into from
Jun 4, 2024
+34 −18
Merged

chore(mpt): Do not expose recursion vars #197

Show file tree
Hide file tree
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
11 changes: 5 additions & 6 deletions crates/mpt/src/db/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -155,7 +155,7 @@ where
pub(crate) fn load_account_from_trie(&mut self, address: Address) -> Result<DbAccount> {
let hashed_address_nibbles = Nibbles::unpack(keccak256(address.as_slice()));
let trie_account_rlp =
self.root_node.open(&hashed_address_nibbles, 0, self.preimage_fetcher)?;
self.root_node.open(&hashed_address_nibbles, self.preimage_fetcher)?;
let trie_account = TrieAccount::decode(&mut trie_account_rlp.as_ref())
.map_err(|e| anyhow!("Error decoding trie account: {e}"))?;

Expand Down Expand Up @@ -231,7 +231,7 @@ where
value.encode(&mut rlp_buf);

if let Ok(storage_slot_rlp) =
storage_root.open(&Nibbles::unpack(hashed_slot_key), 0, self.preimage_fetcher)
storage_root.open(&Nibbles::unpack(hashed_slot_key), self.preimage_fetcher)
{
// If the storage slot already exists, update it.
*storage_slot_rlp = rlp_buf.into();
Expand All @@ -240,7 +240,6 @@ where
storage_root.insert(
&Nibbles::unpack(hashed_slot_key),
rlp_buf.into(),
0,
self.preimage_fetcher,
)?;
}
Expand Down Expand Up @@ -283,13 +282,13 @@ where
let mut account_buf = Vec::with_capacity(trie_account.length());
trie_account.encode(&mut account_buf);

if let Ok(account_rlp_ref) = self.root_node.open(&account_path, 0, preimage_fetcher) {
if let Ok(account_rlp_ref) = self.root_node.open(&account_path, preimage_fetcher) {
// Update the existing account in the trie.
*account_rlp_ref = account_buf.into();
} else {
// Insert the new account into the trie.
self.root_node
.insert(&account_path, account_buf.into(), 0, preimage_fetcher)
.insert(&account_path, account_buf.into(), preimage_fetcher)
.expect("Failed to insert account into trie");
}
}
Expand Down Expand Up @@ -352,7 +351,7 @@ where

let hashed_slot_key = keccak256(index.to_be_bytes::<32>().as_slice());
let slot_value =
storage_root.open(&Nibbles::unpack(hashed_slot_key), 0, fetcher)?;
storage_root.open(&Nibbles::unpack(hashed_slot_key), fetcher)?;

let int_slot = U256::decode(&mut slot_value.as_ref())
.map_err(|e| anyhow!("Failed to decode storage slot value: {e}"))?;
Expand Down
41 changes: 29 additions & 12 deletions crates/mpt/src/node.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,13 +99,21 @@ impl TrieNode {
/// ## 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,
fetcher: impl Fn(B256) -> Result<Bytes> + Copy,
) -> Result<&'a mut Bytes> {
self.open_inner(path, 0, fetcher)
}

/// Inner alias for `open` that keeps track of the nibble offset.
fn open_inner<'a>(
&'a mut self,
path: &Nibbles,
mut nibble_offset: usize,
Expand All @@ -131,11 +139,11 @@ impl TrieNode {
*branch_node = trie_node;

// If the value was found in the blinded node, return it.
branch_node.open(path, nibble_offset, fetcher)
branch_node.open_inner(path, nibble_offset, fetcher)
}
node => {
// If the value was found in the blinded node, return it.
node.open(path, nibble_offset, fetcher)
node.open_inner(path, nibble_offset, fetcher)
}
}
}
Expand Down Expand Up @@ -167,7 +175,7 @@ impl TrieNode {
.map_err(|e| anyhow!(e))?,
);
}
node.open(path, nibble_offset, fetcher)
node.open_inner(path, nibble_offset, fetcher)
} else {
anyhow::bail!("Key does not exist in trie (extension doesn't share nibbles) - {item_key_nibbles:?} {prefix:?} {path:?}");
}
Expand All @@ -176,7 +184,7 @@ impl TrieNode {
let trie_node = TrieNode::decode(&mut fetcher(*commitment)?.as_ref())
.map_err(|e| anyhow!(e))?;
*self = trie_node;
self.open(path, nibble_offset, fetcher)
self.open_inner(path, nibble_offset, fetcher)
}
_ => anyhow::bail!("Invalid trie node type encountered"),
}
Expand All @@ -188,13 +196,22 @@ impl TrieNode {
/// - `self` - The root trie node
/// - `path` - The nibbles representation of the path to the leaf node
/// - `node` - The node to insert at the given path
/// - `nibble_offset` - The number of nibbles that have already been traversed in the `path`
/// - `fetcher` - The preimage fetcher for intermediate blinded nodes
///
/// ## Returns
/// - `Err(_)` - Could not insert the node at the given path in the trie.
/// - `Ok(())` - The node was successfully inserted at the given path.
pub fn insert(
&mut self,
path: &Nibbles,
value: Bytes,
fetcher: impl Fn(B256) -> Result<Bytes> + Copy,
) -> Result<()> {
self.insert_inner(path, value, 0, fetcher)
}

/// Inner alias for `insert` that keeps track of the nibble offset.
fn insert_inner(
&mut self,
path: &Nibbles,
value: Bytes,
Expand Down Expand Up @@ -250,7 +267,7 @@ impl TrieNode {
// If the extension node shares some nibbles with the path, continue the
// insertion recursion.
nibble_offset += shared_extension_nibbles;
node.insert(path, value, nibble_offset, fetcher)?;
node.insert_inner(path, value, nibble_offset, fetcher)?;
return Ok(());
}

Expand Down Expand Up @@ -282,15 +299,15 @@ impl TrieNode {
// Follow the branch node to the next node in the path.
let branch_nibble = path[nibble_offset] as usize;
nibble_offset += BRANCH_NODE_NIBBLES;
stack[branch_nibble].insert(path, value, nibble_offset, fetcher)
stack[branch_nibble].insert_inner(path, value, nibble_offset, fetcher)
}
TrieNode::Blinded { commitment } => {
// If a blinded node is approached, reveal the node and continue the insertion
// recursion.
let trie_node = TrieNode::decode(&mut fetcher(*commitment)?.as_ref())
.map_err(|e| anyhow!(e))?;
*self = trie_node;
self.insert(path, value, nibble_offset, fetcher)
self.insert_inner(path, value, nibble_offset, fetcher)
}
}
}
Expand Down Expand Up @@ -648,7 +665,7 @@ mod test {
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 v = root_node.open(&path_nibbles, fetcher).unwrap();

let mut encoded_value = Vec::with_capacity(value.length());
value.encode(&mut encoded_value);
Expand All @@ -667,9 +684,9 @@ mod test {
fn test_insert_static() {
let mut node =
TrieNode::Leaf { prefix: Nibbles::unpack(hex!("01")), value: Default::default() };
node.insert(&Nibbles::unpack(hex!("012345")), bytes!("01"), 0, |_| Ok(Default::default()))
node.insert(&Nibbles::unpack(hex!("012345")), bytes!("01"), |_| Ok(Default::default()))
.unwrap();
node.insert(&Nibbles::unpack(hex!("012346")), bytes!("02"), 0, |_| Ok(Default::default()))
node.insert(&Nibbles::unpack(hex!("012346")), bytes!("02"), |_| Ok(Default::default()))
.unwrap();

let expected = TrieNode::Extension {
Expand Down
Loading