temp: dont record storage proofs

paradigmxyz · Sep 16, 2024 · 42064af · 42064af
1 parent 02f41f5
commit 42064af
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 61 deletions.
diff --git 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<B256, HashSet<B256>> = 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::<Vec<_>>()
-                        })
-                        .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
@@ -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 => {