From 42064afafcbfde974f88eb11a08905e2e15b044c Mon Sep 17 00:00:00 2001 From: Roman Krasiuk Date: Mon, 16 Sep 2024 22:05:48 +0200 Subject: [PATCH] temp: dont record storage proofs --- crates/engine/tree/src/tree/mod.rs | 86 +++++++++------------------ crates/engine/tree/src/tree/proofs.rs | 6 +- 2 files changed, 31 insertions(+), 61 deletions(-) diff --git a/crates/engine/tree/src/tree/mod.rs b/crates/engine/tree/src/tree/mod.rs index be4eab14b4e62..23239d8edb7e0 100644 --- a/crates/engine/tree/src/tree/mod.rs +++ b/crates/engine/tree/src/tree/mod.rs @@ -46,11 +46,13 @@ use reth_stages_api::ControlFlow; use reth_tasks::TaskSpawner; use reth_trie::{ hashed_cursor::HashedPostStateCursorFactory, + metrics::TrieRootMetrics, proof::Proof, trie_cursor::InMemoryTrieCursorFactory, updates::TrieUpdates, witness::{next_root_from_proofs, target_nodes}, - HashedPostState, MultiProof, Nibbles, TrieAccount, TrieInput, TrieInputSorted, EMPTY_ROOT_HASH, + HashedPostState, MultiProof, Nibbles, StorageRoot, TrieAccount, TrieInput, TrieInputSorted, + EMPTY_ROOT_HASH, }; use reth_trie_db::{DatabaseHashedCursorFactory, DatabaseTrieCursorFactory}; use reth_trie_parallel::parallel_root::ParallelStateRoot; @@ -2358,6 +2360,10 @@ where let consistent_view = ConsistentDbView::new_with_latest_tip(self.provider.clone())?; let provider_ro = consistent_view.provider_ro()?; + let mut prefix_sets = input.prefix_sets.clone(); + prefix_sets.extend(state.construct_prefix_sets()); + + // TODO: clean up let proof_targets: HashMap> = HashMap::from_iter( state .accounts @@ -2371,64 +2377,28 @@ where let account_trie_nodes = proof_targets .into_par_iter() .map(|(hashed_address, hashed_slots)| { - // Gather and record storage trie nodes for this account. - let mut storage_trie_nodes = Vec::with_capacity(hashed_slots.len()); - let storage = state.storages.get(&hashed_address); - for hashed_slot in hashed_slots { - let slot_key = Nibbles::unpack(hashed_slot); - let slot_value = storage - .and_then(|s| s.storage.get(&hashed_slot)) - .filter(|v| !v.is_zero()) - .map(|v| alloy_rlp::encode_fixed_size(v).to_vec()); - let proof = multiproof - .storages + let (storage_root, _, _storage_trie_updates) = StorageRoot::new_hashed( + InMemoryTrieCursorFactory::new( + DatabaseTrieCursorFactory::new(provider_ro.tx_ref()), + &input.nodes, + ), + HashedPostStateCursorFactory::new( + DatabaseHashedCursorFactory::new(provider_ro.tx_ref()), + &input.state, + ), + hashed_address, + TrieRootMetrics::default(), + ) + .with_prefix_set( + prefix_sets + .storage_prefix_sets .get(&hashed_address) - .map(|proof| { - proof - .subtree - .iter() - .filter(|e| slot_key.starts_with(e.0)) - .collect::>() - }) - .unwrap_or_default(); - storage_trie_nodes.extend(target_nodes( - slot_key.clone(), - slot_value, - proof, - None, - )?); - } - - let (storage_root, _storage_trie_updates) = - next_root_from_proofs(storage_trie_nodes, true, |key: Nibbles| { - // Right pad the target with 0s. - let mut padded_key = key.pack(); - padded_key.resize(32, 0); - let mut proof = Proof::new( - InMemoryTrieCursorFactory::new( - DatabaseTrieCursorFactory::new(provider_ro.tx_ref()), - &input.nodes, - ), - HashedPostStateCursorFactory::new( - DatabaseHashedCursorFactory::new(provider_ro.tx_ref()), - &input.state, - ), - ) - .with_target(( - hashed_address, - HashSet::from([B256::from_slice(&padded_key)]), - )) - .multiproof() - .unwrap(); - - // The subtree only contains the proof for a single target. - let node = proof - .storages - .get_mut(&hashed_address) - .and_then(|storage_multiproof| storage_multiproof.subtree.remove(&key)) - .ok_or(TrieWitnessError::MissingTargetNode(key))?; - Ok(node) - })?; + .cloned() + .unwrap_or_default() + .freeze(), + ) + .calculate(true) + .unwrap(); // Gather and record account trie nodes. let account = state diff --git a/crates/engine/tree/src/tree/proofs.rs b/crates/engine/tree/src/tree/proofs.rs index c0b6806dc1ee4..866476cf431bd 100644 --- a/crates/engine/tree/src/tree/proofs.rs +++ b/crates/engine/tree/src/tree/proofs.rs @@ -29,7 +29,7 @@ pub(crate) async fn gather_proofs( let mut targets = HashMap::from([match next { StateAccess::Account(address) => (keccak256(address), HashSet::default()), StateAccess::StorageSlot(address, slot) => { - (keccak256(address), HashSet::from([keccak256(slot)])) + (keccak256(address), HashSet::default()) // HashSet::from([keccak256(slot)])) } }]); @@ -39,7 +39,7 @@ pub(crate) async fn gather_proofs( targets.entry(keccak256(address)).or_default(); } StateAccess::StorageSlot(address, slot) => { - targets.entry(keccak256(address)).or_default().insert(keccak256(slot)); + targets.entry(keccak256(address)).or_default(); // .insert(keccak256(slot)); } } } @@ -105,7 +105,7 @@ where targets.entry(keccak256(address)).or_default(); } StateAccess::StorageSlot(address, slot) => { - targets.entry(keccak256(address)).or_default().insert(keccak256(slot)); + targets.entry(keccak256(address)).or_default(); // .insert(keccak256(slot)); } }, None => {