From 8d28c4c8f2b01d001ffef5e283eb1ec53dedd263 Mon Sep 17 00:00:00 2001 From: DaniPopes <57450786+DaniPopes@users.noreply.github.com> Date: Fri, 30 Jan 2026 23:42:57 +0100 Subject: [PATCH] chore(trie): add set_* methods alongside with_* builders (#21639) Co-authored-by: Amp --- crates/trie/sparse-parallel/src/trie.rs | 14 +++++------ crates/trie/sparse/src/state.rs | 22 ++++++++++++++--- crates/trie/sparse/src/traits.rs | 33 +++++++++++++++++++------ crates/trie/sparse/src/trie.rs | 16 ++++++------ 4 files changed, 57 insertions(+), 28 deletions(-) diff --git a/crates/trie/sparse-parallel/src/trie.rs b/crates/trie/sparse-parallel/src/trie.rs index 927916b24f..ab6f8ff951 100644 --- a/crates/trie/sparse-parallel/src/trie.rs +++ b/crates/trie/sparse-parallel/src/trie.rs @@ -152,27 +152,25 @@ impl Default for ParallelSparseTrie { } impl SparseTrie for ParallelSparseTrie { - fn with_root( - mut self, + fn set_root( + &mut self, root: TrieNode, masks: Option, retain_updates: bool, - ) -> SparseTrieResult { + ) -> SparseTrieResult<()> { // A fresh/cleared `ParallelSparseTrie` has a `SparseNode::Empty` at its root in the upper // subtrie. Delete that so we can reveal the new root node. let path = Nibbles::default(); let _removed_root = self.upper_subtrie.nodes.remove(&path).expect("root node should exist"); debug_assert_eq!(_removed_root, SparseNode::Empty); - self = self.with_updates(retain_updates); + self.set_updates(retain_updates); - self.reveal_upper_node(Nibbles::default(), &root, masks)?; - Ok(self) + self.reveal_upper_node(Nibbles::default(), &root, masks) } - fn with_updates(mut self, retain_updates: bool) -> Self { + fn set_updates(&mut self, retain_updates: bool) { self.updates = retain_updates.then(Default::default); - self } fn reveal_nodes(&mut self, mut nodes: Vec) -> SparseTrieResult<()> { diff --git a/crates/trie/sparse/src/state.rs b/crates/trie/sparse/src/state.rs index 36ba668830..fa3bc25027 100644 --- a/crates/trie/sparse/src/state.rs +++ b/crates/trie/sparse/src/state.rs @@ -114,21 +114,37 @@ impl SparseStateTrie { impl SparseStateTrie { /// Set the retention of branch node updates and deletions. - pub const fn with_updates(mut self, retain_updates: bool) -> Self { + pub const fn set_updates(&mut self, retain_updates: bool) { self.retain_updates = retain_updates; + } + + /// Set the retention of branch node updates and deletions. + pub const fn with_updates(mut self, retain_updates: bool) -> Self { + self.set_updates(retain_updates); self } /// Set the accounts trie to the given `RevealableSparseTrie`. - pub fn with_accounts_trie(mut self, trie: RevealableSparseTrie) -> Self { + pub fn set_accounts_trie(&mut self, trie: RevealableSparseTrie) { self.state = trie; + } + + /// Set the accounts trie to the given `RevealableSparseTrie`. + pub fn with_accounts_trie(mut self, trie: RevealableSparseTrie) -> Self { + self.set_accounts_trie(trie); self } /// Set the default trie which will be cloned when creating new storage /// [`RevealableSparseTrie`]s. - pub fn with_default_storage_trie(mut self, trie: RevealableSparseTrie) -> Self { + pub fn set_default_storage_trie(&mut self, trie: RevealableSparseTrie) { self.storage.default_trie = trie; + } + + /// Set the default trie which will be cloned when creating new storage + /// [`RevealableSparseTrie`]s. + pub fn with_default_storage_trie(mut self, trie: RevealableSparseTrie) -> Self { + self.set_default_storage_trie(trie); self } } diff --git a/crates/trie/sparse/src/traits.rs b/crates/trie/sparse/src/traits.rs index 0e7b7551f2..077e2fb9d5 100644 --- a/crates/trie/sparse/src/traits.rs +++ b/crates/trie/sparse/src/traits.rs @@ -40,17 +40,30 @@ pub trait SparseTrie: Sized + Debug + Send + Sync { /// /// # Returns /// - /// Self if successful, or an error if revealing fails. + /// `Ok(())` if successful, or an error if revealing fails. /// /// # Panics /// /// May panic if the trie is not new/cleared, and has already revealed nodes. - fn with_root( - self, + fn set_root( + &mut self, root: TrieNode, masks: Option, retain_updates: bool, - ) -> SparseTrieResult; + ) -> SparseTrieResult<()>; + + /// Configures the trie to have the given root node revealed. + /// + /// See [`Self::set_root`] for more details. + fn with_root( + mut self, + root: TrieNode, + masks: Option, + retain_updates: bool, + ) -> SparseTrieResult { + self.set_root(root, masks, retain_updates)?; + Ok(self) + } /// Configures the trie to retain information about updates. /// @@ -61,11 +74,15 @@ pub trait SparseTrie: Sized + Debug + Send + Sync { /// # Arguments /// /// * `retain_updates` - Whether to track updates + fn set_updates(&mut self, retain_updates: bool); + + /// Configures the trie to retain information about updates. /// - /// # Returns - /// - /// Self for method chaining. - fn with_updates(self, retain_updates: bool) -> Self; + /// See [`Self::set_updates`] for more details. + fn with_updates(mut self, retain_updates: bool) -> Self { + self.set_updates(retain_updates); + self + } /// Reserves capacity for additional trie nodes. /// diff --git a/crates/trie/sparse/src/trie.rs b/crates/trie/sparse/src/trie.rs index 025be1b393..91ab584791 100644 --- a/crates/trie/sparse/src/trie.rs +++ b/crates/trie/sparse/src/trie.rs @@ -110,7 +110,7 @@ impl RevealableSparseTrie { Box::default() }; - *revealed_trie = revealed_trie.with_root(root, masks, retain_updates)?; + revealed_trie.set_root(root, masks, retain_updates)?; *self = Self::Revealed(revealed_trie); } @@ -453,13 +453,13 @@ impl Default for SerialSparseTrie { } impl SparseTrieTrait for SerialSparseTrie { - fn with_root( - mut self, + fn set_root( + &mut self, root: TrieNode, masks: Option, retain_updates: bool, - ) -> SparseTrieResult { - self = self.with_updates(retain_updates); + ) -> SparseTrieResult<()> { + self.set_updates(retain_updates); // A fresh/cleared `SerialSparseTrie` has a `SparseNode::Empty` at its root. Delete that // so we can reveal the new root node. @@ -467,15 +467,13 @@ impl SparseTrieTrait for SerialSparseTrie { let _removed_root = self.nodes.remove(&path).expect("root node should exist"); debug_assert_eq!(_removed_root, SparseNode::Empty); - self.reveal_node(path, root, masks)?; - Ok(self) + self.reveal_node(path, root, masks) } - fn with_updates(mut self, retain_updates: bool) -> Self { + fn set_updates(&mut self, retain_updates: bool) { if retain_updates { self.updates = Some(SparseTrieUpdates::default()); } - self } fn reserve_nodes(&mut self, additional: usize) {