diff --git a/crates/blockchain/blockchain.rs b/crates/blockchain/blockchain.rs index ffefff77f9d..81dd8e62197 100644 --- a/crates/blockchain/blockchain.rs +++ b/crates/blockchain/blockchain.rs @@ -34,6 +34,7 @@ use ethrex_rlp::encode::RLPEncode; use ethrex_storage::{ AccountUpdatesList, Store, UpdateBatch, error::StoreError, hash_address, hash_key, }; +use ethrex_trie::flattrie::FlatTrie; use ethrex_trie::node::{BranchNode, ExtensionNode}; use ethrex_trie::{Nibbles, Node, NodeRef, Trie}; use ethrex_vm::backends::levm::db::DatabaseLogger; @@ -988,7 +989,7 @@ impl Blockchain { } else { Trie::new_temp() }; - let mut storage_trie_roots = BTreeMap::new(); + let mut storage_tries = BTreeMap::new(); for key in &keys { if key.len() != 20 { continue; // not an address @@ -1011,16 +1012,18 @@ impl Blockchain { "execution witness does not contain non-empty storage trie".to_string(), )); }; - storage_trie_roots.insert(address, (*node).clone()); + storage_tries.insert(address, FlatTrie::from(&(*node))); } + let state_trie = state_trie_root.map(|n| FlatTrie::from(&n)); + Ok(ExecutionWitness { codes, block_headers_bytes, first_block_number: first_block_header.number, chain_config: self.storage.get_chain_config(), - state_trie_root, - storage_trie_roots, + state_trie, + storage_tries, keys, }) } diff --git a/crates/common/trie/flattrie.rs b/crates/common/trie/flattrie.rs new file mode 100644 index 00000000000..6b4fb10528c --- /dev/null +++ b/crates/common/trie/flattrie.rs @@ -0,0 +1,873 @@ +use ethereum_types::H256; +use ethrex_crypto::keccak::keccak_hash; +use ethrex_rlp::{ + constants::RLP_NULL, + decode::{RLPDecode, decode_bytes}, + encode::{RLPEncode, encode_length}, + error::RLPDecodeError, + structs::{Decoder, Encoder}, +}; +use rkyv::with::Skip; + +use crate::{EMPTY_TRIE_HASH, Nibbles, Node, NodeHash, NodeRef, rlp::decode_child}; + +/// A trie implementation that is non recursive, POD and avoids copying and encoding +/// nodes by providing views into a RLP flat buffer +#[derive( + Default, + serde::Serialize, + serde::Deserialize, + rkyv::Serialize, + rkyv::Deserialize, + rkyv::Archive, + Clone, +)] +pub struct FlatTrie { + /// Stores a contiguous byte buffer with each RLP encoded node + pub data: Vec, + /// Contains the structural information of the MPT + pub views: Vec, + /// The index of the view for the root of this trie + pub root_index: Option, + /// Root hash that gets initialized when calling `Self::authenticate` + #[serde(skip)] + #[rkyv(with = Skip)] + pub root_hash: Option, +} + +/// A view into a particular node +#[derive( + Clone, + Copy, + serde::Serialize, + serde::Deserialize, + rkyv::Serialize, + rkyv::Deserialize, + rkyv::Archive, +)] +pub struct NodeView { + /// Indices to the RLP code of this node over the flat data buffer + pub data_range: (usize, usize), + /// Handle into this node's childs + pub node_type: NodeType, +} + +/// Contains indices to the `handles` list for each child of a node +#[derive( + Clone, + Copy, + serde::Serialize, + serde::Deserialize, + rkyv::Serialize, + rkyv::Deserialize, + rkyv::Archive, +)] +pub enum NodeType { + /// A leaf node doesn't have any childs + Leaf, + /// An extension node always has a branch as a child + Extension { child: Option }, + /// A branch node can have any node type as any of its 16 childs + /// TODO: This can be optimized to a bitmap if the data vec is ordered (contiguous childs) + Branch { children: [Option; 16] }, +} + +pub enum NodeData { + Leaf { partial: Nibbles, value: Vec }, + Extension { path: Nibbles, child: NodeHash }, + Branch { children: [Option; 16] }, +} + +impl From<&Node> for FlatTrie { + fn from(root: &Node) -> Self { + let mut trie = FlatTrie::default(); + + fn recursive(value: &Node, trie: &mut FlatTrie) { + let childs = match value { + Node::Branch(node) => { + let mut childs = [None; 16]; + for (i, choice) in node.choices.iter().enumerate() { + if let NodeRef::Node(choice, _) = choice { + recursive(&(*choice), trie); + childs[i] = Some(trie.views.len() - 1); + } + } + NodeType::Branch { children: childs } + } + Node::Extension(node) => { + let mut child = None; + if let NodeRef::Node(child_node, _) = &node.child { + recursive(child_node, trie); + child = Some(trie.views.len() - 1); + } + NodeType::Extension { child } + } + Node::Leaf(_) => NodeType::Leaf, + }; + + let offset = trie.data.len(); + trie.data.extend(value.encode_to_vec()); + trie.views.push(NodeView { + data_range: (offset, trie.data.len()), + node_type: childs, + }); + } + + recursive(root, &mut trie); + trie.root_index = Some(trie.views.len() - 1); // last stored node is the root + trie + } +} + +impl FlatTrie { + // TODO: authentication should be done just once, when creating the trie. + pub fn root_hash(&mut self) -> Result, RLPDecodeError> { + if self.root_hash.is_none() && !self.authenticate()? { + return Ok(None); + } + Ok(self.root_hash) + } + + pub fn is_empty(&self) -> bool { + self.data.is_empty() + } + + /// Recursively traverses the trie, hashes each node's data and checks that their parents reference + /// those same hashes. This way the data gets authenticated. + /// + /// Initializes `Self::root_hash` if successful, and returns a boolean indicating success. + pub fn authenticate(&mut self) -> Result { + fn recursive<'a>( + trie: &'a FlatTrie, + view: &NodeView, + ) -> Result, RLPDecodeError> { + match view.node_type { + NodeType::Leaf => {} + NodeType::Extension { child } => { + if let Some(child) = child { + let Some(child_hash) = recursive(trie, &trie.views[child])? else { + return Ok(None); + }; + let Some(items) = trie.get_encoded_items(view)? else { + panic!(); // TODO: err + }; + let child_data_hash = decode_child(items[1]); + if child_hash != child_data_hash { + return Ok(None); + } + } + } + NodeType::Branch { children: childs } => { + // TODO: we can decode just the Some(_) childs + let Some(items) = trie.get_encoded_items(view)? else { + panic!(); // TODO: err + }; + for (i, child) in childs.iter().enumerate() { + if let Some(child) = child { + let child_view = &trie.views[*child]; + let Some(child_hash) = recursive(trie, child_view)? else { + return Ok(None); + }; + let child_data_hash = decode_child(items[i]); + if child_hash != child_data_hash { + return Ok(None); + } + } + } + } + } + Ok(Some(trie.get_hash_view(view))) + } + + let Some(root_view) = self.get_root_view() else { + self.root_hash = Some((*EMPTY_TRIE_HASH).into()); + return Ok(true); + }; + + let Some(root_hash) = recursive(&self, root_view)? else { + return Ok(false); + }; + self.root_hash = Some(root_hash); + Ok(true) + } + + pub fn get(&self, path: &[u8]) -> Result, RLPDecodeError> { + let mut path = Nibbles::from_bytes(path); + fn recursive<'a>( + trie: &'a FlatTrie, + path: &mut Nibbles, + view: &NodeView, + ) -> Result, RLPDecodeError> { + match view.node_type { + NodeType::Leaf => { + let Some(items) = trie.get_encoded_items(view)? else { + panic!(); // TODO: err + }; + + let (partial, _) = decode_bytes(items[0])?; + let partial = Nibbles::decode_compact(partial); + debug_assert!(partial.is_leaf()); + + if partial == *path { + let (value, _) = decode_bytes(items[1])?; + return Ok(Some(value)); + } else { + return Ok(None); + } + } + NodeType::Extension { child } => { + let Some(items) = trie.get_encoded_items(view)? else { + panic!(); // TODO: err + }; + + let (prefix, _) = decode_bytes(items[0])?; + let prefix = Nibbles::decode_compact(prefix); + debug_assert!(!prefix.is_leaf()); + + if path.skip_prefix(prefix) { + recursive(trie, path, &trie.views[child.unwrap()]) + } else { + Ok(None) + } + } + NodeType::Branch { children: childs } => { + if let Some(choice) = path.next_choice() { + let Some(child_view_index) = childs[choice] else { + return Ok(None); + }; + let child_view = trie.views[child_view_index]; + recursive(trie, path, &child_view) + } else { + let Some(items) = trie.get_encoded_items(view)? else { + panic!(); // TODO: err + }; + let (value, _) = decode_bytes(items[16])?; + Ok((!value.is_empty()).then_some(value)) + } + } + } + } + + let Some(root_view) = self.get_root_view() else { + panic!(); // TODO: err + }; + recursive(&self, &mut path, root_view) + } + + pub fn put(&mut self, node_type: NodeType, data: NodeData) -> usize { + // TODO: check that both match + + let start = self.data.len(); + let mut encoder = Encoder::new(&mut self.data); + match data { + NodeData::Leaf { partial, value } => { + encoder = encoder.encode_bytes(&partial.encode_compact()); + encoder = encoder.encode_bytes(&value); + } + NodeData::Extension { path, child } => { + encoder = encoder.encode_bytes(&path.encode_compact()); + encoder = child.encode(encoder); + } + NodeData::Branch { children } => { + for child in children { + if let Some(child) = child { + encoder = child.encode(encoder); + } else { + encoder = encoder.encode_raw(&[RLP_NULL]) + } + } + encoder = encoder.encode_raw(&[RLP_NULL]) + } + } + encoder.finish(); + let end = self.data.len(); + + let data_range = (start, end); + + let view = NodeView { + data_range, + node_type, + }; + self.views.push(view); + self.views.len() - 1 + } + + pub fn insert(&mut self, path: Vec, value: Vec) -> Result<(), RLPDecodeError> { + let path = Nibbles::from_bytes(&path); + if let Some(root_index) = self.root_index { + self.root_index = Some(self.insert_inner(root_index, path, value)?); + } else { + self.root_index = Some(self.put_leaf(path, value)); + } + self.root_hash = None; + Ok(()) + } + + fn insert_inner( + &mut self, + self_view_index: usize, + mut path: Nibbles, + value: Vec, + ) -> Result { + let self_view = self.views[self_view_index]; + match self_view.node_type { + NodeType::Leaf => { + let Some(items) = self.get_encoded_items(&self_view)? else { + panic!(); + }; + + let (partial, _) = decode_bytes(items[0])?; + let partial = Nibbles::decode_compact(partial); + debug_assert!(partial.is_leaf()); + + if partial == path { + Ok(self.put_leaf(partial, value)) + } else { + // Current node will be replaced with a branch or extension node + let match_index = path.count_prefix(&partial); + let self_choice_idx = partial.at(match_index); + let new_leaf_choice_idx = path.at(match_index); + + // Modify the partial of self + let (self_value, _) = decode_bytes(items[1])?; + let new_self_view_index = + self.put_leaf(partial.offset(match_index + 1), self_value.to_vec()); + + let branch_view_index = if self_choice_idx == 16 { + // Yields a new leaf node with the new value, and a parent branch node + // with the old value. We disallow branches with values. + unreachable!("leaf insertion yielded branch with old value"); + } else if new_leaf_choice_idx == 16 { + // Yields a new branch node with the current leaf as child, and the new + // value in the branch. We disallow branches with values. + unreachable!("leaf insertion yielded branch with new value") + } else { + // Yields a new leaf with the path and value in it, and a new branch + // with the new and old leaf as children. + let new_leaf_view_index = + self.put_leaf(path.offset(match_index + 1), value); + self.put_branch(vec![ + ( + new_leaf_choice_idx, + ( + Some(new_leaf_view_index), + self.get_hash(new_leaf_view_index), + ), + ), + ( + self_choice_idx, + ( + Some(new_self_view_index), + self.get_hash(new_self_view_index), + ), + ), + ]) + }; + + if match_index == 0 { + Ok(branch_view_index) + } else { + // Yields an extension node with the branch as child + Ok(self.put_extension( + path.slice(0, match_index), + self.get_hash(branch_view_index), + Some(branch_view_index), + )) + } + } + } + NodeType::Extension { child } => { + let Some(items) = self.get_encoded_items(&self_view)? else { + panic!(); + }; + + let (prefix, _) = decode_bytes(items[0])?; + let prefix = Nibbles::decode_compact(prefix); + debug_assert!(!prefix.is_leaf()); + + let match_index = path.count_prefix(&prefix); + if match_index == prefix.len() { + let path = path.offset(match_index); + let new_child_view_index = self.insert_inner( + child.expect("missing child of extension node at match_index == prefix"), + path, + value, + )?; + Ok(self.put_extension( + prefix, + self.get_hash(new_child_view_index), + Some(new_child_view_index), + )) + } else if match_index == 0 { + debug_assert!( + prefix.at(0) != 16, + "insertion into extension yielded branch with value" + ); + let branch_view_index = if prefix.len() == 1 { + self.put_branch(vec![( + prefix.at(0), + (child, self.get_extension_child(&self_view)?), + )]) + } else { + // New extension with self_node as a child + let new_node_view_index = self.put_extension( + prefix.offset(1), + self.get_extension_child(&self_view)?, + child, + ); + self.put_branch(vec![( + prefix.at(0), + ( + Some(new_node_view_index), + self.get_hash(new_node_view_index), + ), + )]) + }; + self.insert_inner(branch_view_index, path, value) + } else { + let new_extension_view_index = self.put_extension( + prefix.offset(match_index), + self.get_extension_child(&self_view)?, + child, + ); + let new_node_view_index = self.insert_inner( + new_extension_view_index, + path.offset(match_index), + value, + )?; + + Ok(self.put_extension( + prefix.slice(0, match_index), + self.get_hash(new_node_view_index), + Some(new_node_view_index), + )) + } + } + NodeType::Branch { mut children } => { + // let Some(items) = self.get_encoded_items(&self_view)? else { + // panic!(); + // }; + + // let children: [_; 16] = std::array::from_fn(|i| { + // let (child, _) = decode_bytes(items[i]).unwrap(); + // if child.is_empty() { + // None + // } else { + // Some(child.to_vec()) + // } + // }); + + if let Some(choice) = path.next_choice() { + let new_child_view_index = match children[choice] { + // TODO: we are not differentiating between a child which is not included in the data + // vs an empty choice. We should decode the RLP and get the information from there, + // or it could be added to the view. + None => self.put_leaf(path, value), + Some(view_index) => self.insert_inner(view_index, path, value)?, + }; + children[choice] = Some(new_child_view_index); + let children = children + .into_iter() + .enumerate() + .filter_map(|(i, child)| Some((i, (Some(child?), self.get_hash(child?))))) + .collect(); + Ok(self.put_branch(children)) + } else { + // We disallow values in branchs + unreachable!("wanted to insert value in a branch"); + } + } + } + } + + pub fn remove(&mut self, path: &[u8]) -> Result<(), RLPDecodeError> { + let path = Nibbles::from_bytes(path); + if let Some(root_index) = self.root_index { + self.root_index = self.remove_inner(root_index, path)?; + } + self.root_hash = None; + Ok(()) + } + + fn remove_inner( + &mut self, + self_view_index: usize, + mut path: Nibbles, + ) -> Result, RLPDecodeError> { + let self_view = self.views[self_view_index]; + match self_view.node_type { + NodeType::Leaf => { + let Some(items) = self.get_encoded_items(&self_view)? else { + panic!(); + }; + + let (partial, _) = decode_bytes(items[0])?; + let partial = Nibbles::decode_compact(partial); + debug_assert!(partial.is_leaf()); + + if partial == path { + Ok(None) + } else { + Ok(Some(self_view_index)) + } + } + NodeType::Extension { child } => { + let mut prefix = self.get_extension_prefix(&self_view)?; + + if path.skip_prefix(&prefix) { + let new_child_view_index = self.remove_inner( + child.expect("missing child of extension node at remove"), + path, + )?; + let Some(new_child_view_index) = new_child_view_index else { + return Ok(None); + }; + + let new_child_view = self.get_view(new_child_view_index).unwrap(); + let new_view_index = match new_child_view.node_type { + NodeType::Branch { .. } => self.put_extension( + prefix, + self.get_hash(new_child_view_index), + Some(new_child_view_index), + ), + NodeType::Extension { + child: new_extension_child, + } => { + let new_child_prefix = self.get_extension_prefix(new_child_view)?; + prefix.extend(&new_child_prefix); + let new_extension_child = new_extension_child + .expect("missing child of new extension at remove"); + self.put_extension( + prefix, + self.get_hash(new_extension_child), + Some(new_extension_child), + ) + } + NodeType::Leaf => { + let (partial, value) = + self.get_leaf_partial_and_value(new_child_view)?; + prefix.extend(&partial); + self.put_leaf(prefix, value.to_vec()) + } + }; + Ok(Some(new_view_index)) + } else { + Ok(Some(self_view_index)) + } + } + NodeType::Branch { mut children } => { + let Some(choice) = path.next_choice() else { + // We disallow values on branches + unreachable!(); + }; + let Some(child_view_index) = children[choice] else { + return Ok(Some(self_view_index)); + }; + + let new_child_index = self.remove_inner(child_view_index, path)?; + children[choice] = new_child_index; + + let children: Vec<(_, _)> = children + .into_iter() + .enumerate() + .filter_map(|(i, child)| Some((i, (Some(child?), self.get_hash(child?))))) + .collect(); + + match children.len() { + 0 => Ok(None), + 1 => { + let (choice_idx, (child_idx, _)) = children[0]; + let child_idx = child_idx.expect("missing child of branch at remove"); + let child_view = self + .get_view(child_idx) + .expect("missing child view of branch choice at remove"); + + match child_view.node_type { + NodeType::Leaf => { + let (mut partial, value) = + self.get_leaf_partial_and_value(child_view)?; + partial.prepend(choice_idx as u8); + Ok(Some(self.put_leaf(partial, value.to_vec()))) + } + NodeType::Extension { child } => { + let mut prefix = self.get_extension_prefix(child_view)?; + prefix.prepend(choice_idx as u8); + let child = child + .expect("missing child of extension at remove for branch case"); + Ok(Some(self.put_extension( + prefix, + self.get_hash(child), + Some(child), + ))) + } + NodeType::Branch { .. } => { + let prefix = Nibbles::from_hex(vec![choice_idx as u8]); + Ok(Some(self.put_extension( + prefix, + self.get_hash(child_idx), + Some(child_idx), + ))) + } + } + } + _ => Ok(Some(self.put_branch(children))), + } + } + } + } + + /// Puts a new leaf node from a prefix and a value. + /// + /// Returns the new node's view index. + pub fn put_leaf(&mut self, partial: Nibbles, value: Vec) -> usize { + let data = NodeData::Leaf { partial, value }; + let children = NodeType::Leaf; + self.put(children, data) + } + + /// Puts a new extension node from a path, child hash and child view. + /// + /// Returns the new node's view index. + pub fn put_extension( + &mut self, + path: Nibbles, + child_hash: NodeHash, + child_view_index: Option, + ) -> usize { + let data = NodeData::Extension { + path, + child: child_hash, + }; + let children = NodeType::Extension { + child: child_view_index, + }; + self.put(children, data) + } + + /// Puts a new branch node from a list of children in the form of (choice, view_index). Childrens + /// need to exist in the trie's data. + /// + /// Returns the new node's view index. + pub fn put_branch(&mut self, children_views: Vec<(usize, (Option, NodeHash))>) -> usize { + let data = { + let mut children: [_; 16] = std::array::from_fn(|_| None); + for (choice, child) in &children_views { + children[*choice] = Some(child.1); + } + NodeData::Branch { children } + }; + let children = { + let mut children = [None; 16]; + for (choice, child) in children_views { + children[choice] = child.0; + } + NodeType::Branch { children } + }; + self.put(children, data) + } + + pub fn get_view(&self, index: usize) -> Option<&NodeView> { + self.views.get(index) + } + + pub fn get_root_view(&self) -> Option<&NodeView> { + self.root_index + .map(|i| self.get_view(i).expect("missing root view")) + } + + /// Calculates the hash of a node from a view index of its data. + /// TODO: cache? we should cache a range into the hash, because its already stored in its parent + pub fn get_hash(&self, view_index: usize) -> NodeHash { + NodeHash::from_encoded(self.get_data(view_index)) + } + + /// Calculates the hash of a node from a view of its data. + /// TODO: cache? we should cache a range into the hash, because its already stored in its parent + /// TODO: consider changing into a NodeHash or similar + pub fn get_hash_view(&self, view: &NodeView) -> NodeHash { + NodeHash::from_encoded(self.get_data_view(view)) + } + + // TODO: cache decoded view? + pub fn get_encoded_items(&self, view: &NodeView) -> Result>, RLPDecodeError> { + let data = self.get_data_view(view); + let mut decoder = Decoder::new(data)?; + + let mut rlp_items = Vec::with_capacity(17); + while !decoder.is_done() && rlp_items.len() < 17 { + let (item, new_decoder) = decoder.get_encoded_item_ref()?; + decoder = new_decoder; + rlp_items.push(item); + } + + Ok(Some(rlp_items)) + } + + pub fn get_data(&self, view_index: usize) -> &[u8] { + self.get_data_view(&self.views[view_index]) + } + + pub fn get_data_view(&self, view: &NodeView) -> &[u8] { + &self.data[view.data_range.0..view.data_range.1] + } + + pub fn get_extension_prefix(&self, view: &NodeView) -> Result { + let Some(items) = self.get_encoded_items(view)? else { + panic!(); + }; + + let (prefix, _) = decode_bytes(items[0])?; + let prefix = Nibbles::decode_compact(prefix); + debug_assert!(!prefix.is_leaf()); + Ok(prefix) + } + + pub fn get_extension_child(&self, view: &NodeView) -> Result { + let Some(items) = self.get_encoded_items(view)? else { + panic!(); + }; + Ok(decode_child(items[1])) + } + + pub fn get_leaf_partial(&self, view: &NodeView) -> Result { + let Some(items) = self.get_encoded_items(view)? else { + panic!(); + }; + + let (partial, _) = decode_bytes(items[0])?; + let partial = Nibbles::decode_compact(partial); + debug_assert!(partial.is_leaf()); + Ok(partial) + } + + pub fn get_leaf_partial_and_value( + &self, + view: &NodeView, + ) -> Result<(Nibbles, &[u8]), RLPDecodeError> { + let Some(items) = self.get_encoded_items(view)? else { + panic!(); + }; + + let (partial, _) = decode_bytes(items[0])?; + let partial = Nibbles::decode_compact(partial); + debug_assert!(partial.is_leaf()); + + let (value, _) = decode_bytes(items[1])?; + + Ok((partial, value)) + } +} + +// fn encode_leaf(partial: &[u8], value: &[u8]) -> Vec { +// // TODO: fix headroom +// let mut buf = Vec::with_capacity(partial.len() + value.len() + 3); // 3 byte headroom +// Encoder::new(&mut buf) +// .encode_bytes(partial) +// .encode_bytes(value) +// .finish(); +// buf +// } +// +// fn encode_leaf_raw(raw_partial: &[u8], value: &[u8]) -> Vec { +// // TODO: fix headroom +// let mut buf = Vec::with_capacity(raw_partial.len() + value.len() + 3); // 3 byte headroom +// Encoder::new(&mut buf) +// .encode_raw(raw_partial) +// .encode_bytes(value) +// .finish(); +// buf +// } +// +// fn encode_branch(choices: [Option<[u8; 32]>; 16]) -> Vec { +// let value_len = 1; +// let choices_len = choices.iter().fold(0, |acc, choice| { +// acc + choice.map(|h| RLPEncode::length(&h)).unwrap_or(0) +// }); +// let payload_len = choices_len + value_len; +// +// let mut buf: Vec = Vec::with_capacity(choices_len + 3); // 3 byte prefix headroom +// +// encode_length(payload_len, &mut buf); +// for choice in choices { +// if let Some(choice) = choice { +// choice.encode(&mut buf); +// } else { +// buf.push(RLP_NULL); +// } +// } +// buf +// } + +#[cfg(test)] +mod test { + use ethrex_rlp::encode::RLPEncode; + use proptest::{ + collection::{btree_set, vec}, + prelude::*, + }; + + use crate::{Nibbles, Trie, flattrie::FlatTrie}; + + fn kv_pairs_strategy() -> impl Strategy, Vec)>, Vec)> { + // create random key-values, with keys all the same size, and a random permutation of indices + (1usize..32).prop_flat_map(|key_len| { + prop::collection::vec((vec(any::(), key_len), vec(any::(), 0..256)), 1..2) + .prop_flat_map(|kvs| { + let len = kvs.len(); + let shuffle = vec(..len, ..len).prop_shuffle(); + (Just(kvs), shuffle) + }) + }) + } + + proptest! { + #[test] + fn proptest_insert_compare_hash((kv, _) in kv_pairs_strategy()) { + let mut trie = Trie::new_temp(); + let mut flat_trie = FlatTrie::default(); + + for (key, value) in kv.iter(){ + trie.insert(key.clone(), value.clone()).unwrap(); + flat_trie.insert(key.clone(), value.clone()).unwrap(); + + let hash = trie.hash_no_commit(); + + prop_assert!(flat_trie.authenticate().unwrap()); + let flat_trie_hash = flat_trie.root_hash().unwrap().unwrap(); + + prop_assert_eq!(hash, flat_trie_hash.finalize()); + } + } + + #[test] + fn proptest_insert_remove_compare_hash((kv, shuffle) in kv_pairs_strategy()) { + let mut trie = Trie::new_temp(); + let mut flat_trie = FlatTrie::default(); + + for (key, value) in kv.iter() { + trie.insert(key.clone(), value.clone()).unwrap(); + flat_trie.insert(key.clone(), value.clone()).unwrap(); + + let hash = trie.hash_no_commit(); + + prop_assert!(flat_trie.authenticate().unwrap()); + let flat_trie_hash = flat_trie.root_hash().unwrap().unwrap(); + + prop_assert_eq!(hash, flat_trie_hash.finalize()); + } + + for i in shuffle.iter() { + let key = &kv[*i].0; + trie.remove(key).unwrap(); + flat_trie.remove(key).unwrap(); + + let hash = trie.hash_no_commit(); + + prop_assert!(flat_trie.authenticate().unwrap()); + let flat_trie_hash = flat_trie.root_hash().unwrap().unwrap(); + + prop_assert_eq!(hash, flat_trie_hash.finalize()); + } + } + } +} diff --git a/crates/common/trie/nibbles.rs b/crates/common/trie/nibbles.rs index b3c7d28fcb7..923d37778b8 100644 --- a/crates/common/trie/nibbles.rs +++ b/crates/common/trie/nibbles.rs @@ -101,10 +101,11 @@ impl Nibbles { /// If `prefix` is a prefix of self, move the offset after /// the prefix and return true, otherwise return false. - pub fn skip_prefix(&mut self, prefix: &Nibbles) -> bool { - if self.len() >= prefix.len() && &self.data[..prefix.len()] == prefix.as_ref() { + pub fn skip_prefix(&mut self, prefix: impl AsRef<[u8]>) -> bool { + let prefix = prefix.as_ref(); + if self.len() >= prefix.len() && &self.data[..prefix.len()] == prefix { self.data = self.data[prefix.len()..].to_vec(); - self.already_consumed.extend(&prefix.data); + self.already_consumed.extend(prefix); true } else { false diff --git a/crates/common/trie/rlp.rs b/crates/common/trie/rlp.rs index 80abea7eec0..eaf954f170e 100644 --- a/crates/common/trie/rlp.rs +++ b/crates/common/trie/rlp.rs @@ -67,10 +67,10 @@ impl RLPEncode for ExtensionNode { impl RLPEncode for LeafNode { fn encode(&self, buf: &mut dyn bytes::BufMut) { - Encoder::new(buf) - .encode_bytes(&self.partial.encode_compact()) - .encode_bytes(&self.value) - .finish() + let mut encoder = Encoder::new(buf); + encoder = encoder.encode_bytes(&self.partial.encode_compact()); + encoder = encoder.encode_bytes(&self.value); + encoder.finish() } } @@ -151,7 +151,7 @@ impl RLPDecode for Node { } } -fn decode_child(rlp: &[u8]) -> NodeHash { +pub fn decode_child(rlp: &[u8]) -> NodeHash { match decode_bytes(rlp) { Ok((hash, &[])) if hash.len() == 32 => NodeHash::from_slice(hash), Ok((&[], &[])) => NodeHash::default(), diff --git a/crates/common/trie/trie.rs b/crates/common/trie/trie.rs index 32b30235a7b..ca0f9a59f7e 100644 --- a/crates/common/trie/trie.rs +++ b/crates/common/trie/trie.rs @@ -1,5 +1,6 @@ pub mod db; pub mod error; +pub mod flattrie; pub mod logger; mod nibbles; pub mod node; diff --git a/crates/common/types/block_execution_witness.rs b/crates/common/types/block_execution_witness.rs index a970d96238b..cd76d259e52 100644 --- a/crates/common/types/block_execution_witness.rs +++ b/crates/common/types/block_execution_witness.rs @@ -10,7 +10,8 @@ use ethereum_types::{Address, H256, U256}; use ethrex_crypto::keccak::keccak_hash; use ethrex_rlp::error::RLPDecodeError; use ethrex_rlp::{decode::RLPDecode, encode::RLPEncode}; -use ethrex_trie::{EMPTY_TRIE_HASH, Node, Trie, TrieError}; +use ethrex_trie::flattrie::FlatTrie; +use ethrex_trie::{EMPTY_TRIE_HASH, Nibbles, Node, Trie, TrieError}; use rkyv::with::{Identity, MapKV}; use serde::{Deserialize, Serialize}; @@ -32,7 +33,7 @@ pub struct GuestProgramState { pub block_headers: BTreeMap, /// The accounts state trie containing the necessary state for the guest /// program execution. - pub state_trie: Trie, + pub state_trie: FlatTrie, /// The parent block header of the first block in the batch. pub parent_block_header: BlockHeader, /// The block number of the first block in the batch. @@ -40,7 +41,7 @@ pub struct GuestProgramState { /// The chain configuration. pub chain_config: ChainConfig, /// Map of storage root hashes to their corresponding storage tries. - pub storage_tries: BTreeMap, + pub storage_tries: BTreeMap, /// Map of account addresses to their corresponding hashed addresses. /// This is a convenience map to avoid recomputing the hashed address /// multiple times during guest program execution. @@ -72,10 +73,10 @@ pub struct ExecutionWitness { // The chain config. pub chain_config: ChainConfig, /// Root node embedded with the rest of the trie's nodes - pub state_trie_root: Option, + pub state_trie: Option, /// Root nodes per account storage embedded with the rest of the trie's nodes #[rkyv(with = MapKV)] - pub storage_trie_roots: BTreeMap, + pub storage_tries: BTreeMap, /// Flattened map of account addresses and storage keys whose values /// are needed for stateless execution. #[rkyv(with = crate::rkyv_utils::VecVecWrapper)] @@ -110,8 +111,15 @@ impl TryFrom for GuestProgramState { type Error = GuestProgramStateError; fn try_from(value: ExecutionWitness) -> Result { - let block_headers: BTreeMap = value - .block_headers_bytes + let ExecutionWitness { + codes, + block_headers_bytes, + first_block_number, + state_trie, + storage_tries: original_storage_tries, + .. + } = value; + let block_headers: BTreeMap = block_headers_bytes .into_iter() .map(|bytes| BlockHeader::decode(bytes.as_ref())) .collect::, _>>() @@ -123,8 +131,7 @@ impl TryFrom for GuestProgramState { .collect(); let parent_number = - value - .first_block_number + first_block_number .checked_sub(1) .ok_or(GuestProgramStateError::Custom( "First block number cannot be zero".to_string(), @@ -135,25 +142,27 @@ impl TryFrom for GuestProgramState { )?; // hash state trie nodes - let state_trie = if let Some(state_trie_root) = value.state_trie_root { - Trie::new_temp_with_root(state_trie_root.into()) + let state_trie = if let Some(mut state_trie) = state_trie { + if !state_trie.authenticate()? { + panic!(); + } + state_trie } else { - Trie::new_temp() + FlatTrie::default() }; - state_trie.hash_no_commit(); let mut storage_tries = BTreeMap::new(); - for (address, storage_trie_root) in value.storage_trie_roots { + for (address, mut storage_trie) in original_storage_tries { // hash storage trie nodes - let storage_trie = Trie::new_temp_with_root(storage_trie_root.into()); - storage_trie.hash_no_commit(); + if !storage_trie.authenticate()? { + panic!(); + } storage_tries.insert(address, storage_trie); } // hash codes // TODO: codes here probably needs to be Vec, rather than recomputing here. This requires rkyv implementation. - let codes_hashed = value - .codes + let codes_hashed = codes .into_iter() .map(|code| { let code = Code::from_bytecode(code.into()); @@ -246,7 +255,8 @@ impl GuestProgramState { .expect("failed to remove key"); } - let storage_root = storage_trie.hash_no_commit(); + storage_trie.authenticate()?; + let storage_root = storage_trie.root_hash()?.unwrap().finalize(); account_state.storage_root = storage_root; } @@ -260,8 +270,9 @@ impl GuestProgramState { /// Returns the root hash of the state trie /// Returns an error if the state trie is not built yet - pub fn state_trie_root(&self) -> Result { - Ok(self.state_trie.hash_no_commit()) + pub fn state_trie_root(&mut self) -> Result { + self.state_trie.authenticate()?; + Ok(self.state_trie.root_hash()?.unwrap().finalize()) // TOOD: unwrap } /// Returns Some(block_number) if the hash for block_number is not the parent @@ -429,7 +440,7 @@ impl GuestProgramState { pub fn get_valid_storage_trie( &mut self, address: Address, - ) -> Result, GuestProgramStateError> { + ) -> Result, GuestProgramStateError> { let is_storage_verified = *self.verified_storage_roots.get(&address).unwrap_or(&false); if is_storage_verified { Ok(self.storage_tries.get(&address)) @@ -439,9 +450,15 @@ impl GuestProgramState { // empty account return Ok(None); }; - let storage_trie = match self.storage_tries.get(&address) { + let storage_trie = match self.storage_tries.get_mut(&address) { None if storage_root == *EMPTY_TRIE_HASH => return Ok(None), - Some(trie) if trie.hash_no_commit() == storage_root => trie, + Some(trie) => { + if trie.root_hash()?.unwrap().finalize() == storage_root { + trie + } else { + panic!() + } + } _ => { return Err(GuestProgramStateError::Custom(format!( "invalid storage trie for account {address}" diff --git a/crates/networking/rpc/debug/execution_witness.rs b/crates/networking/rpc/debug/execution_witness.rs index d936df7db49..6b093e0107c 100644 --- a/crates/networking/rpc/debug/execution_witness.rs +++ b/crates/networking/rpc/debug/execution_witness.rs @@ -11,7 +11,7 @@ use ethrex_common::{ }; use ethrex_rlp::{decode::RLPDecode, error::RLPDecodeError}; use ethrex_storage::hash_address; -use ethrex_trie::{EMPTY_TRIE_HASH, Node, NodeRef, Trie, TrieError}; +use ethrex_trie::{EMPTY_TRIE_HASH, Node, NodeRef, Trie, TrieError, flattrie::FlatTrie}; use serde::{Deserialize, Serialize}; use serde_json::Value; use tracing::debug; @@ -45,23 +45,24 @@ pub struct RpcExecutionWitness { impl TryFrom for RpcExecutionWitness { type Error = TrieError; fn try_from(value: ExecutionWitness) -> Result { - let mut nodes = Vec::new(); - if let Some(state_trie_root) = value.state_trie_root { - state_trie_root.encode_subtrie(&mut nodes)?; - } - for node in value.storage_trie_roots.values() { - node.encode_subtrie(&mut nodes)?; - } - Ok(Self { - state: nodes.into_iter().map(Bytes::from).collect(), - keys: value.keys.into_iter().map(Bytes::from).collect(), - codes: value.codes.into_iter().map(Bytes::from).collect(), - headers: value - .block_headers_bytes - .into_iter() - .map(Bytes::from) - .collect(), - }) + panic!(); + // let mut nodes = Vec::new(); + // if let Some(state_trie_root) = value.state_trie_root { + // state_trie_root.encode_subtrie(&mut nodes)?; + // } + // for node in value.storage_trie_roots.values() { + // node.encode_subtrie(&mut nodes)?; + // } + // Ok(Self { + // state: nodes.into_iter().map(Bytes::from).collect(), + // keys: value.keys.into_iter().map(Bytes::from).collect(), + // codes: value.codes.into_iter().map(Bytes::from).collect(), + // headers: value + // .block_headers_bytes + // .into_iter() + // .map(Bytes::from) + // .collect(), + // }) } } @@ -118,7 +119,7 @@ pub fn execution_witness_from_rpc_chain_config( } else { Trie::new_temp() }; - let mut storage_trie_roots = BTreeMap::new(); + let mut storage_tries = BTreeMap::new(); for key in &rpc_witness.keys { if key.len() != 20 { continue; // not an address @@ -141,9 +142,12 @@ pub fn execution_witness_from_rpc_chain_config( "execution witness does not contain non-empty storage trie".to_string(), )); }; - storage_trie_roots.insert(address, (*node).clone()); + let trie = FlatTrie::from(&(*node)); + storage_tries.insert(address, trie.clone()); } + let state_trie = state_trie_root.map(|n| FlatTrie::from(&n)); + let witness = ExecutionWitness { codes: rpc_witness.codes.into_iter().map(|b| b.to_vec()).collect(), chain_config, @@ -153,8 +157,8 @@ pub fn execution_witness_from_rpc_chain_config( .into_iter() .map(|b| b.to_vec()) .collect(), - state_trie_root, - storage_trie_roots, + state_trie, + storage_tries, keys: rpc_witness.keys.into_iter().map(|b| b.to_vec()).collect(), }; diff --git a/crates/vm/levm/src/precompiles.rs b/crates/vm/levm/src/precompiles.rs index 3bfaa7c70e0..4ac2eab96dc 100644 --- a/crates/vm/levm/src/precompiles.rs +++ b/crates/vm/levm/src/precompiles.rs @@ -834,9 +834,11 @@ pub fn bn254_g1_mul(g1: G1, scalar: U256) -> Result { Fq::from_slice(&g1.1.to_big_endian()).map_err(|_| PrecompileError::ParsingInputError)?, ); - let g1: G1 = AffineG1::new(g1_x, g1_y) - .map_err(|_| PrecompileError::InvalidPoint)? - .into(); + let g1 = AffineG1::new(g1_x, g1_y).map_err(|_| PrecompileError::InvalidPoint)?; + + // Small difference between the patched versions of substrate-bn + #[cfg(feature = "zisk")] + let g1: G1 = g1.into(); let scalar = Fr::from_slice(&scalar.to_big_endian()).map_err(|_| PrecompileError::ParsingInputError)?;