feat(trie): TrieInput (#10918)

This commit is contained in:
Roman Krasiuk
2024-09-16 10:04:17 +02:00
committed by GitHub
parent f37bd47e2e
commit b0eb78ced6
14 changed files with 131 additions and 203 deletions

View File

@ -839,7 +839,7 @@ mod tests {
AccountReader, BlockHashReader, StateProofProvider, StateProvider, StateRootProvider,
StorageRootProvider,
};
use reth_trie::{prefix_set::TriePrefixSetsMut, AccountProof, HashedStorage};
use reth_trie::{AccountProof, HashedStorage, TrieInput};
fn create_mock_state(
test_block_builder: &mut TestBlockBuilder,
@ -913,12 +913,7 @@ mod tests {
Ok(B256::random())
}
fn state_root_from_nodes(
&self,
_nodes: TrieUpdates,
_post_state: HashedPostState,
_prefix_sets: TriePrefixSetsMut,
) -> ProviderResult<B256> {
fn state_root_from_nodes(&self, _input: TrieInput) -> ProviderResult<B256> {
Ok(B256::random())
}
@ -931,9 +926,7 @@ mod tests {
fn state_root_from_nodes_with_updates(
&self,
_nodes: TrieUpdates,
_post_state: HashedPostState,
_prefix_sets: TriePrefixSetsMut,
_input: TrieInput,
) -> ProviderResult<(B256, TrieUpdates)> {
Ok((B256::random(), TrieUpdates::default()))
}

View File

@ -7,10 +7,7 @@ use reth_storage_api::{
AccountReader, BlockHashReader, StateProofProvider, StateProvider, StateProviderBox,
StateRootProvider, StorageRootProvider,
};
use reth_trie::{
prefix_set::TriePrefixSetsMut, updates::TrieUpdates, AccountProof, HashedPostState,
HashedStorage,
};
use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, TrieInput};
use std::{collections::HashMap, sync::OnceLock};
/// A state provider that stores references to in-memory blocks along with their state as well as
@ -45,13 +42,12 @@ impl MemoryOverlayStateProvider {
/// Return lazy-loaded trie state aggregated from in-memory blocks.
fn trie_state(&self) -> &MemoryOverlayTrieState {
self.trie_state.get_or_init(|| {
let mut hashed_state = HashedPostState::default();
let mut trie_nodes = TrieUpdates::default();
let mut trie_state = MemoryOverlayTrieState::default();
for block in self.in_memory.iter().rev() {
hashed_state.extend_ref(block.hashed_state.as_ref());
trie_nodes.extend_ref(block.trie.as_ref());
trie_state.state.extend_ref(block.hashed_state.as_ref());
trie_state.nodes.extend_ref(block.trie.as_ref());
}
MemoryOverlayTrieState { trie_nodes, hashed_state }
trie_state
})
}
}
@ -102,54 +98,38 @@ impl AccountReader for MemoryOverlayStateProvider {
}
impl StateRootProvider for MemoryOverlayStateProvider {
fn state_root(&self, hashed_state: HashedPostState) -> ProviderResult<B256> {
let prefix_sets = hashed_state.construct_prefix_sets();
self.state_root_from_nodes(TrieUpdates::default(), hashed_state, prefix_sets)
fn state_root(&self, state: HashedPostState) -> ProviderResult<B256> {
self.state_root_from_nodes(TrieInput::from_state(state))
}
fn state_root_from_nodes(
&self,
nodes: TrieUpdates,
state: HashedPostState,
prefix_sets: TriePrefixSetsMut,
) -> ProviderResult<B256> {
let MemoryOverlayTrieState { mut trie_nodes, mut hashed_state } = self.trie_state().clone();
trie_nodes.extend(nodes);
hashed_state.extend(state);
self.historical.state_root_from_nodes(trie_nodes, hashed_state, prefix_sets)
fn state_root_from_nodes(&self, mut input: TrieInput) -> ProviderResult<B256> {
let MemoryOverlayTrieState { nodes, state } = self.trie_state().clone();
input.prepend_cached(nodes, state);
self.historical.state_root_from_nodes(input)
}
fn state_root_with_updates(
&self,
hashed_state: HashedPostState,
state: HashedPostState,
) -> ProviderResult<(B256, TrieUpdates)> {
let prefix_sets = hashed_state.construct_prefix_sets();
self.state_root_from_nodes_with_updates(TrieUpdates::default(), hashed_state, prefix_sets)
self.state_root_from_nodes_with_updates(TrieInput::from_state(state))
}
fn state_root_from_nodes_with_updates(
&self,
nodes: TrieUpdates,
state: HashedPostState,
prefix_sets: TriePrefixSetsMut,
mut input: TrieInput,
) -> ProviderResult<(B256, TrieUpdates)> {
let MemoryOverlayTrieState { mut trie_nodes, mut hashed_state } = self.trie_state().clone();
trie_nodes.extend(nodes);
hashed_state.extend(state);
self.historical.state_root_from_nodes_with_updates(trie_nodes, hashed_state, prefix_sets)
let MemoryOverlayTrieState { nodes, state } = self.trie_state().clone();
input.prepend_cached(nodes, state);
self.historical.state_root_from_nodes_with_updates(input)
}
}
impl StorageRootProvider for MemoryOverlayStateProvider {
// TODO: Currently this does not reuse available in-memory trie nodes.
fn storage_root(&self, address: Address, storage: HashedStorage) -> ProviderResult<B256> {
let mut hashed_storage = self
.trie_state()
.hashed_state
.storages
.get(&keccak256(address))
.cloned()
.unwrap_or_default();
let mut hashed_storage =
self.trie_state().state.storages.get(&keccak256(address)).cloned().unwrap_or_default();
hashed_storage.extend(&storage);
self.historical.storage_root(address, hashed_storage)
}
@ -163,7 +143,7 @@ impl StateProofProvider for MemoryOverlayStateProvider {
address: Address,
slots: &[B256],
) -> ProviderResult<AccountProof> {
let mut hashed_state = self.trie_state().hashed_state.clone();
let mut hashed_state = self.trie_state().state.clone();
hashed_state.extend(state);
self.historical.proof(hashed_state, address, slots)
}
@ -174,7 +154,7 @@ impl StateProofProvider for MemoryOverlayStateProvider {
overlay: HashedPostState,
target: HashedPostState,
) -> ProviderResult<HashMap<B256, Bytes>> {
let mut hashed_state = self.trie_state().hashed_state.clone();
let mut hashed_state = self.trie_state().state.clone();
hashed_state.extend(overlay);
self.historical.witness(hashed_state, target)
}
@ -207,10 +187,10 @@ impl StateProvider for MemoryOverlayStateProvider {
}
/// The collection of data necessary for trie-related operations for [`MemoryOverlayStateProvider`].
#[derive(Clone, Debug)]
#[derive(Clone, Default, Debug)]
pub(crate) struct MemoryOverlayTrieState {
/// The collection of aggregated in-memory trie updates.
pub(crate) trie_nodes: TrieUpdates,
pub(crate) nodes: TrieUpdates,
/// The collection of hashed state from in-memory blocks.
pub(crate) hashed_state: HashedPostState,
pub(crate) state: HashedPostState,
}