feat(trie): use Arc instead of references for nodes and state cursor …

…factories
paradigmxyz · Jan 20, 2025 · 63b94ab · 63b94ab
1 parent 6cc660c
commit 63b94ab
Show file tree

Hide file tree

Showing 17 changed files with 226 additions and 141 deletions.
diff --git a/crates/engine/tree/benches/state_root_task.rs b/crates/engine/tree/benches/state_root_task.rs
@@ -247,11 +247,11 @@ fn bench_state_root(c: &mut Criterion) {
                         let blinded_provider_factory = ProofBlindedProviderFactory::new(
                             InMemoryTrieCursorFactory::new(
                                 DatabaseTrieCursorFactory::new(provider.tx_ref()),
-                                &nodes_sorted,
+                                nodes_sorted,
                             ),
                             HashedPostStateCursorFactory::new(
                                 DatabaseHashedCursorFactory::new(provider.tx_ref()),
-                                &state_sorted,
+                                state_sorted,
                             ),
                             prefix_sets,
                         );

diff --git a/crates/engine/tree/src/tree/mod.rs b/crates/engine/tree/src/tree/mod.rs
@@ -2311,13 +2311,13 @@ where
 
                     let in_memory_trie_cursor = InMemoryTrieCursorFactory::new(
                         DatabaseTrieCursorFactory::new(context.provider_ro.tx_ref()),
-                        &context.nodes_sorted,
+                        context.nodes_sorted.clone(),
                     );
                     let blinded_provider_factory = ProofBlindedProviderFactory::new(
                         in_memory_trie_cursor.clone(),
                         HashedPostStateCursorFactory::new(
                             DatabaseHashedCursorFactory::new(context.provider_ro.tx_ref()),
-                            &context.state_sorted,
+                            context.state_sorted.clone(),
                         ),
                         context.prefix_sets.clone(),
                     );

diff --git a/crates/engine/tree/src/tree/root.rs b/crates/engine/tree/src/tree/root.rs
@@ -973,11 +973,11 @@ mod tests {
         let blinded_provider_factory = ProofBlindedProviderFactory::new(
             InMemoryTrieCursorFactory::new(
                 DatabaseTrieCursorFactory::new(provider.tx_ref()),
-                &nodes_sorted,
+                nodes_sorted,
             ),
             HashedPostStateCursorFactory::new(
                 DatabaseHashedCursorFactory::new(provider.tx_ref()),
-                &state_sorted,
+                state_sorted,
             ),
             config.prefix_sets.clone(),
         );

diff --git a/crates/trie/db/src/proof.rs b/crates/trie/db/src/proof.rs
@@ -10,6 +10,9 @@ use reth_trie::{
     StorageMultiProof, TrieInput,
 };
 
+extern crate alloc;
+use alloc::sync::Arc;
+
 /// Extends [`Proof`] with operations specific for working with a database transaction.
 pub trait DatabaseProof<'a, TX> {
     /// Create a new [Proof] from database transaction.
@@ -50,11 +53,11 @@ impl<'a, TX: DbTx> DatabaseProof<'a, TX>
         Self::from_tx(tx)
             .with_trie_cursor_factory(InMemoryTrieCursorFactory::new(
                 DatabaseTrieCursorFactory::new(tx),
-                &nodes_sorted,
+                Arc::new(nodes_sorted),
             ))
             .with_hashed_cursor_factory(HashedPostStateCursorFactory::new(
                 DatabaseHashedCursorFactory::new(tx),
-                &state_sorted,
+                Arc::new(state_sorted),
             ))
             .with_prefix_sets_mut(input.prefix_sets)
             .account_proof(address, slots)
@@ -70,11 +73,11 @@ impl<'a, TX: DbTx> DatabaseProof<'a, TX>
         Self::from_tx(tx)
             .with_trie_cursor_factory(InMemoryTrieCursorFactory::new(
                 DatabaseTrieCursorFactory::new(tx),
-                &nodes_sorted,
+                Arc::new(nodes_sorted),
             ))
             .with_hashed_cursor_factory(HashedPostStateCursorFactory::new(
                 DatabaseHashedCursorFactory::new(tx),
-                &state_sorted,
+                Arc::new(state_sorted),
             ))
             .with_prefix_sets_mut(input.prefix_sets)
             .multiproof(targets)
@@ -125,7 +128,7 @@ impl<'a, TX: DbTx> DatabaseStorageProof<'a, TX>
         Self::from_tx(tx, address)
             .with_hashed_cursor_factory(HashedPostStateCursorFactory::new(
                 DatabaseHashedCursorFactory::new(tx),
-                &state_sorted,
+                Arc::new(state_sorted),
             ))
             .with_prefix_set_mut(prefix_set)
             .storage_proof(slot)
@@ -147,7 +150,7 @@ impl<'a, TX: DbTx> DatabaseStorageProof<'a, TX>
         Self::from_tx(tx, address)
             .with_hashed_cursor_factory(HashedPostStateCursorFactory::new(
                 DatabaseHashedCursorFactory::new(tx),
-                &state_sorted,
+                Arc::new(state_sorted),
             ))
             .with_prefix_set_mut(prefix_set)
             .storage_multiproof(targets)

diff --git a/crates/trie/db/src/state.rs b/crates/trie/db/src/state.rs
@@ -19,6 +19,9 @@ use reth_trie::{
 use std::{collections::HashMap, ops::RangeInclusive};
 use tracing::debug;
 
+extern crate alloc;
+use alloc::sync::Arc;
+
 /// Extends [`StateRoot`] with operations specific for working with a database transaction.
 pub trait DatabaseStateRoot<'a, TX>: Sized {
     /// Create a new [`StateRoot`] instance.
@@ -173,7 +176,10 @@ impl<'a, TX: DbTx> DatabaseStateRoot<'a, TX>
         let state_sorted = post_state.into_sorted();
         StateRoot::new(
             DatabaseTrieCursorFactory::new(tx),
-            HashedPostStateCursorFactory::new(DatabaseHashedCursorFactory::new(tx), &state_sorted),
+            HashedPostStateCursorFactory::new(
+                DatabaseHashedCursorFactory::new(tx),
+                Arc::new(state_sorted),
+            ),
         )
         .with_prefix_sets(prefix_sets)
         .root()
@@ -187,7 +193,10 @@ impl<'a, TX: DbTx> DatabaseStateRoot<'a, TX>
         let state_sorted = post_state.into_sorted();
         StateRoot::new(
             DatabaseTrieCursorFactory::new(tx),
-            HashedPostStateCursorFactory::new(DatabaseHashedCursorFactory::new(tx), &state_sorted),
+            HashedPostStateCursorFactory::new(
+                DatabaseHashedCursorFactory::new(tx),
+                Arc::new(state_sorted),
+            ),
         )
         .with_prefix_sets(prefix_sets)
         .root_with_updates()
@@ -197,8 +206,14 @@ impl<'a, TX: DbTx> DatabaseStateRoot<'a, TX>
         let state_sorted = input.state.into_sorted();
         let nodes_sorted = input.nodes.into_sorted();
         StateRoot::new(
-            InMemoryTrieCursorFactory::new(DatabaseTrieCursorFactory::new(tx), &nodes_sorted),
-            HashedPostStateCursorFactory::new(DatabaseHashedCursorFactory::new(tx), &state_sorted),
+            InMemoryTrieCursorFactory::new(
+                DatabaseTrieCursorFactory::new(tx),
+                Arc::new(nodes_sorted),
+            ),
+            HashedPostStateCursorFactory::new(
+                DatabaseHashedCursorFactory::new(tx),
+                Arc::new(state_sorted),
+            ),
         )
         .with_prefix_sets(input.prefix_sets.freeze())
         .root()
@@ -211,8 +226,14 @@ impl<'a, TX: DbTx> DatabaseStateRoot<'a, TX>
         let state_sorted = input.state.into_sorted();
         let nodes_sorted = input.nodes.into_sorted();
         StateRoot::new(
-            InMemoryTrieCursorFactory::new(DatabaseTrieCursorFactory::new(tx), &nodes_sorted),
-            HashedPostStateCursorFactory::new(DatabaseHashedCursorFactory::new(tx), &state_sorted),
+            InMemoryTrieCursorFactory::new(
+                DatabaseTrieCursorFactory::new(tx),
+                Arc::new(nodes_sorted),
+            ),
+            HashedPostStateCursorFactory::new(
+                DatabaseHashedCursorFactory::new(tx),
+                Arc::new(state_sorted),
+            ),
         )
         .with_prefix_sets(input.prefix_sets.freeze())
         .root_with_updates()

diff --git a/crates/trie/db/src/storage.rs b/crates/trie/db/src/storage.rs
@@ -7,6 +7,9 @@ use reth_trie::{
     hashed_cursor::HashedPostStateCursorFactory, HashedPostState, HashedStorage, StorageRoot,
 };
 
+extern crate alloc;
+use alloc::sync::Arc;
+
 #[cfg(feature = "metrics")]
 use reth_trie::metrics::{TrieRootMetrics, TrieType};
 
@@ -68,7 +71,10 @@ impl<'a, TX: DbTx> DatabaseStorageRoot<'a, TX>
             HashedPostState::from_hashed_storage(keccak256(address), hashed_storage).into_sorted();
         StorageRoot::new(
             DatabaseTrieCursorFactory::new(tx),
-            HashedPostStateCursorFactory::new(DatabaseHashedCursorFactory::new(tx), &state_sorted),
+            HashedPostStateCursorFactory::new(
+                DatabaseHashedCursorFactory::new(tx),
+                Arc::new(state_sorted),
+            ),
             address,
             prefix_set,
             #[cfg(feature = "metrics")]

diff --git a/crates/trie/db/src/witness.rs b/crates/trie/db/src/witness.rs
@@ -7,6 +7,9 @@ use reth_trie::{
     witness::TrieWitness, HashedPostState, TrieInput,
 };
 
+extern crate alloc;
+use alloc::sync::Arc;
+
 /// Extends [`TrieWitness`] with operations specific for working with a database transaction.
 pub trait DatabaseTrieWitness<'a, TX> {
     /// Create a new [`TrieWitness`] from database transaction.
@@ -37,11 +40,11 @@ impl<'a, TX: DbTx> DatabaseTrieWitness<'a, TX>
         Self::from_tx(tx)
             .with_trie_cursor_factory(InMemoryTrieCursorFactory::new(
                 DatabaseTrieCursorFactory::new(tx),
-                &nodes_sorted,
+                Arc::new(nodes_sorted),
             ))
             .with_hashed_cursor_factory(HashedPostStateCursorFactory::new(
                 DatabaseHashedCursorFactory::new(tx),
-                &state_sorted,
+                Arc::new(state_sorted),
             ))
             .with_prefix_sets_mut(input.prefix_sets)
             .compute(target)

diff --git a/crates/trie/db/tests/fuzz_in_memory_nodes.rs b/crates/trie/db/tests/fuzz_in_memory_nodes.rs
@@ -16,13 +16,12 @@ use reth_trie::{
     HashedPostState, HashedStorage, StateRoot, StorageRoot,
 };
 use reth_trie_db::{DatabaseStateRoot, DatabaseStorageRoot, DatabaseTrieCursorFactory};
-use std::collections::BTreeMap;
+use std::{collections::BTreeMap, sync::Arc};
 
 proptest! {
     #![proptest_config(ProptestConfig {
         cases: 128, ..ProptestConfig::default()
     })]
-
     #[test]
     fn fuzz_in_memory_account_nodes(mut init_state: BTreeMap<B256, U256>, state_updates: [BTreeMap<B256, Option<U256>>; 10]) {
         let factory = create_test_provider_factory();
@@ -58,8 +57,8 @@ proptest! {
             // Compute root with in-memory trie nodes overlay
             let (state_root, trie_updates) = StateRoot::from_tx(provider.tx_ref())
                 .with_prefix_sets(hashed_state.construct_prefix_sets().freeze())
-                .with_trie_cursor_factory(InMemoryTrieCursorFactory::new(
-                    DatabaseTrieCursorFactory::new(provider.tx_ref()), &trie_nodes.clone().into_sorted())
+                .with_trie_cursor_factory(
+                    InMemoryTrieCursorFactory::new(DatabaseTrieCursorFactory::new(provider.tx_ref()), Arc::new(trie_nodes.clone().into_sorted()))
                 )
                 .root_with_updates()
                 .unwrap();
@@ -113,10 +112,9 @@ proptest! {
             let (storage_root, _, trie_updates) =
                 StorageRoot::from_tx_hashed(provider.tx_ref(), hashed_address)
                     .with_prefix_set(hashed_storage.construct_prefix_set().freeze())
-                    .with_trie_cursor_factory(InMemoryTrieCursorFactory::new(
-                        DatabaseTrieCursorFactory::new(provider.tx_ref()),
-                        &trie_nodes.into_sorted(),
-                    ))
+                .with_trie_cursor_factory(
+                    InMemoryTrieCursorFactory::new(DatabaseTrieCursorFactory::new(provider.tx_ref()), Arc::new(trie_nodes.clone().into_sorted()))
+                    )
                     .root_with_updates()
                     .unwrap();