Skip to content

Commit

Permalink
feat: add children discovery to trie_traversal
Browse files Browse the repository at this point in the history
  • Loading branch information
@morph-dev
morph-dev committed Sep 9, 2024
1 parent 26444cb commit 76bdd35
Show file tree
Hide file tree
Showing 2 changed files with 171 additions and 133 deletions.
113 changes: 111 additions & 2 deletions ethportal-api/src/types/state_trie/trie_traversal.rs
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
use std::{ops::Deref, vec};

use alloy_primitives::{Bytes, B256};
use eth_trie::node::Node;
use thiserror::Error;
Expand Down Expand Up @@ -50,10 +52,57 @@ pub enum TraversalResult<'path> {
Error(TraversalError),
}

#[derive(Debug, Clone, Default)]
pub struct NodeChildren(Vec<NodeChild>);

impl Deref for NodeChildren {
type Target = Vec<NodeChild>;

fn deref(&self) -> &Self::Target {
&self.0
}
}

impl IntoIterator for NodeChildren {
type Item = NodeChild;
type IntoIter = <Vec<NodeChild> as IntoIterator>::IntoIter;

fn into_iter(self) -> Self::IntoIter {
self.0.into_iter()
}
}

impl NodeChildren {
pub fn prepend_path(&mut self, prefix: &[u8]) {
for child in &mut self.0 {
child.prepend_path(prefix.iter().cloned());
}
}
}

#[derive(Debug, Clone)]
pub enum NodeChild {
Value { path: Vec<u8>, value: Bytes },
Node { path: Vec<u8>, hash: B256 },
}

impl NodeChild {
pub fn prepend_path(&mut self, prefix: impl IntoIterator<Item = u8>) {
let path = match self {
NodeChild::Value { path, .. } => path,
NodeChild::Node { path, .. } => path,
};
path.splice(0..0, prefix);
}
}

/// The trait for traversing Merkle Patricia Trie node.
pub trait NodeTraversal {
/// Traverses the node and returns the result.
fn traverse<'path>(&self, path: &'path [u8]) -> TraversalResult<'path>;

/// Returns all children Hash Nodes or Values within the node (in-order).
fn children(&self) -> Result<NodeChildren, TraversalError>;
}

impl NodeTraversal for Node {
Expand Down Expand Up @@ -103,12 +152,72 @@ impl NodeTraversal for Node {
},
}
}
Node::Hash(hash) => TraversalResult::Node(NextTraversalNode {
hash: hash.hash,
Node::Hash(hash_node) => TraversalResult::Node(NextTraversalNode {
hash: hash_node.hash,
remaining_path: path,
}),
}
}

fn children(&self) -> Result<NodeChildren, TraversalError> {
match self {
Node::Empty => Ok(NodeChildren::default()),
Node::Leaf(leaf_node) => {
let Some((last, nibbles)) = leaf_node.key.get_data().split_last() else {
return Err(TraversalError::InvalidLeafKey);
};
if *last != 16 {
return Err(TraversalError::InvalidLeafKey);
}
if nibbles.is_empty() {
return Err(TraversalError::EmptyLeafKey);
}

Ok(NodeChildren(vec![NodeChild::Value {
path: nibbles.into(),
value: Bytes::from_iter(&leaf_node.value),
}]))
}
Node::Extension(extension_node) => {
let Ok(extension_node) = extension_node.read() else {
return Err(TraversalError::ReadLockPoisoned);
};
let prefix = extension_node.prefix.get_data();
if prefix.is_empty() {
return Err(TraversalError::EmptyExtensionPrefix);
}

let mut children = extension_node.node.children()?;
children.prepend_path(prefix);
Ok(children)
}
Node::Branch(branch_node) => {
let Ok(branch_node) = branch_node.read() else {
return Err(TraversalError::ReadLockPoisoned);
};

let mut children = vec![];
if let Some(value) = &branch_node.value {
children.push(NodeChild::Value {
path: vec![],
value: Bytes::from_iter(value),
});
}

for (child_index, child) in branch_node.children.iter().enumerate() {
let mut child_children = child.children()?;
child_children.prepend_path(&[child_index as u8]);
children.extend(child_children.0);
}

Ok(NodeChildren(children))
}
Node::Hash(hash_node) => Ok(NodeChildren(vec![NodeChild::Node {
path: vec![],
hash: hash_node.hash,
}])),
}
}
}

#[cfg(test)]
Expand Down
191 changes: 60 additions & 131 deletions trin-execution/src/subcommands/era2.rs
View file
Original file line number Diff line number Diff line change
@@ -1,18 +1,16 @@
use std::sync::Arc;

use alloy_rlp::{Decodable, EMPTY_STRING_CODE};
use anyhow::{ensure, Error};
use anyhow::{anyhow, ensure, Error};
use e2store::era2::{
AccountEntry, AccountOrStorageEntry, Era2, StorageEntry, StorageItem, MAX_STORAGE_ITEMS,
};
use eth_trie::{
decode_node,
node::{LeafNode, Node},
EthTrie, Trie, DB,
use eth_trie::{decode_node, EthTrie, Trie, DB};
use ethportal_api::{
types::state_trie::trie_traversal::{NodeChild, NodeTraversal},
Header,
};
use ethportal_api::Header;
use parking_lot::Mutex;
use revm_primitives::{keccak256, B256, KECCAK_EMPTY, U256};
use revm_primitives::{keccak256, Bytes, B256, KECCAK_EMPTY, U256};
use tracing::info;

use crate::{
Expand All @@ -25,118 +23,52 @@ use crate::{
#[derive(Debug)]
struct LeafNodeWithKeyHash {
key_hash: B256,
value: Vec<u8>,
}

#[derive(Debug)]
struct TrieNode {
node_hash: B256,
path: Vec<u8>,
}

impl TrieNode {
fn new(node_hash: B256, path: Vec<u8>) -> Self {
Self { node_hash, path }
}
value: Bytes,
}

struct TrieLeafIterator<TrieDB: DB> {
trie: Arc<Mutex<EthTrie<TrieDB>>>,
stack: Vec<TrieNode>,
leafs_to_process: Vec<LeafNodeWithKeyHash>,
db: Arc<TrieDB>,
stack: Vec<NodeChild>,
}

impl<TrieDB: DB> TrieLeafIterator<TrieDB> {
pub fn new(trie: Arc<Mutex<EthTrie<TrieDB>>>) -> anyhow::Result<Self> {
let stack = vec![TrieNode::new(trie.lock().root_hash()?, vec![])];
Ok(Self {
trie,
stack,
leafs_to_process: vec![],
})
}

fn process_leaf(&mut self, leaf: Arc<LeafNode>, path: Vec<u8>) {
// reconstruct the address hash from the path so that we can fetch the
// address from the database
let mut partial_key_path = leaf.key.get_data().to_vec();
partial_key_path.pop();
let full_key_path = [&path, partial_key_path.as_slice()].concat();
let address_hash = full_nibble_path_to_address_hash(&full_key_path);
self.leafs_to_process.push(LeafNodeWithKeyHash {
key_hash: address_hash,
value: leaf.value.clone(),
});
}

fn process_node(&mut self, node: Node, path: Vec<u8>) -> anyhow::Result<()> {
match node {
Node::Leaf(leaf) => self.process_leaf(leaf, path),
Node::Extension(extension) => {
let extension = extension.read().expect("Extension node must be readable");
let path_with_extension_prefix =
[path, extension.prefix.get_data().to_vec()].concat();
match &extension.node {
Node::Hash(hash) => {
self.stack
.push(TrieNode::new(hash.hash, path_with_extension_prefix));
}
Node::Leaf(leaf) => self.process_leaf(leaf.clone(), path_with_extension_prefix),
_ => {
panic!("Invalid extension node, must be either a leaf or a hash");
}
}
}
Node::Branch(branch) => {
let branch = branch.read().expect("Branch node must be readable");
for (i, child) in branch.children.iter().enumerate() {
let branch_path = [path.clone(), vec![i as u8]].concat();
match child {
Node::Leaf(leaf) => self.process_leaf(leaf.clone(), branch_path),
Node::Hash(hash) => self.stack.push(TrieNode::new(hash.hash, branch_path)),
Node::Empty => {} // Do nothing
_ => {
panic!("Invalid branch node, must be either a leaf or a hash")
}
}
}
if let Some(node) = &branch.value {
let decoded_node = decode_node(&mut node.as_slice())?;
if let Node::Leaf(leaf) = decoded_node {
self.process_leaf(leaf, path);
} else {
panic!("Invalid branch value, must be a leaf");
}
}
}
Node::Hash(_) => {}
Node::Empty => {}
pub fn new(db: Arc<TrieDB>, root_hash: B256) -> Self {
Self {
db,
stack: vec![NodeChild::Node {
path: vec![],
hash: root_hash,
}],
}
}

fn process_node(&mut self, hash: B256, path: Vec<u8>) -> anyhow::Result<()> {
let encoded_trie_node = self
.db
.get(hash.as_slice())
.map_err(|err| anyhow!("Unable to get node from trie db: {err}"))?
.ok_or_else(|| {
anyhow!("Node must exist, as we are walking a valid trie | node_hash: {hash} path: {path:?}")
})?;
let trie_node = decode_node(&mut encoded_trie_node.as_slice())?;

let mut children = trie_node.children()?;
children.prepend_path(&path);
// Add children to stack in reverse order, so that final iteration is in-order
self.stack.extend(children.into_iter().rev());
Ok(())
}

pub fn next(&mut self) -> Result<Option<LeafNodeWithKeyHash>, Error> {
// Well process the storage trie, since values can be small, they may be inlined into
// branches, so we will process them here
if let Some(leaf_node) = self.leafs_to_process.pop() {
return Ok(Some(leaf_node));
}

while let Some(TrieNode { node_hash, path }) = self.stack.pop() {
let node = self
.trie
.lock()
.db
.get(node_hash.as_slice())
.expect("Unable to get node from trie db")
.unwrap_or_else(|| panic!("Node must exist, as we are walking a valid trie | node_hash: {} path: {:?}",
node_hash, path));

self.process_node(decode_node(&mut node.as_slice())?, path)?;

if let Some(leaf_node) = self.leafs_to_process.pop() {
return Ok(Some(leaf_node));
while let Some(child) = self.stack.pop() {
match child {
NodeChild::Value { path, value } => {
return Ok(Some(LeafNodeWithKeyHash {
key_hash: full_nibble_path_to_address_hash(&path),
value,
}))
}
NodeChild::Node { path, hash } => self.process_node(hash, path)?,
}
}

Expand Down Expand Up @@ -168,15 +100,16 @@ impl StateExporter {
);
let mut era2 = Era2::create(self.exporter_config.path_to_era2.clone(), header)?;
info!("Era2 initiated");
let mut leaf_iterator = TrieLeafIterator::new(self.trin_execution.database.trie.clone())?;

let mut account_trie = self.trin_execution.database.trie.lock();
let mut leaf_iterator =
TrieLeafIterator::new(account_trie.db.clone(), account_trie.root_hash()?);
drop(account_trie);

info!("Trie leaf iterator initiated");
let mut accounts_processed = 0;
while let Ok(Some(LeafNodeWithKeyHash {
key_hash: account_hash,
value: account_state,
})) = leaf_iterator.next()
{
let account_state: Account = Decodable::decode(&mut account_state.as_slice())?;
while let Ok(Some(account_leaf)) = leaf_iterator.next() {
let account_state: Account = Decodable::decode(&mut account_leaf.value.as_ref())?;
let bytecode = if account_state.code_hash != KECCAK_EMPTY {
self.trin_execution
.database
Expand All @@ -189,22 +122,18 @@ impl StateExporter {

let mut storage: Vec<StorageItem> = vec![];
if account_state.storage_root != keccak256([EMPTY_STRING_CODE]) {
let account_db =
AccountDB::new(account_hash, self.trin_execution.database.db.clone());
let account_trie = Arc::new(Mutex::new(EthTrie::from(
Arc::new(account_db),
account_state.storage_root,
)?));
let mut storage_iterator = TrieLeafIterator::new(account_trie)?;
while let Ok(Some(LeafNodeWithKeyHash {
key_hash: storage_index_hash,
value: storage_value,
})) = storage_iterator.next()
{
let storage_slot_value: U256 = Decodable::decode(&mut storage_value.as_slice())
.expect("Failed to decode storage slot value");
let account_db = AccountDB::new(
account_leaf.key_hash,
self.trin_execution.database.db.clone(),
);
let mut storage_leaf_iterator =
TrieLeafIterator::new(Arc::new(account_db), account_state.storage_root);
while let Ok(Some(storage_leaf)) = storage_leaf_iterator.next() {
let storage_slot_value: U256 =
Decodable::decode(&mut storage_leaf.value.as_ref())
.expect("Failed to decode storage slot value");
storage.push(StorageItem {
storage_index_hash,
storage_index_hash: storage_leaf.key_hash,
value: storage_slot_value,
});
}
Expand All @@ -215,7 +144,7 @@ impl StateExporter {
+ (storage.len() % MAX_STORAGE_ITEMS != 0) as usize;

era2.append_entry(&AccountOrStorageEntry::Account(AccountEntry {
address_hash: account_hash,
address_hash: account_leaf.key_hash,
account_state: (&account_state).into(),
bytecode,
storage_count: storage_count as u32,
Expand Down

0 comments on commit 76bdd35

Please sign in to comment.