diff --git a/crates/mpt/src/db/mod.rs b/crates/mpt/src/db/mod.rs index 58dfc5019..238dba83a 100644 --- a/crates/mpt/src/db/mod.rs +++ b/crates/mpt/src/db/mod.rs @@ -155,7 +155,7 @@ where pub(crate) fn load_account_from_trie(&mut self, address: Address) -> Result { 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}"))?; @@ -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(); @@ -240,7 +240,6 @@ where storage_root.insert( &Nibbles::unpack(hashed_slot_key), rlp_buf.into(), - 0, self.preimage_fetcher, )?; } @@ -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"); } } @@ -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}"))?; diff --git a/crates/mpt/src/node.rs b/crates/mpt/src/node.rs index 1804044fb..34797b271 100644 --- a/crates/mpt/src/node.rs +++ b/crates/mpt/src/node.rs @@ -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 + 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, @@ -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) } } } @@ -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:?}"); } @@ -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"), } @@ -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 + 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, @@ -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(()); } @@ -282,7 +299,7 @@ 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 @@ -290,7 +307,7 @@ impl TrieNode { 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) } } } @@ -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); @@ -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 {