diff --git a/crates/trie/trie/src/prefix_set.rs b/crates/trie/trie/src/prefix_set.rs
index 4997228050..af0fb173d9 100644
--- a/crates/trie/trie/src/prefix_set.rs
+++ b/crates/trie/trie/src/prefix_set.rs
@@ -20,9 +20,9 @@ pub struct TriePrefixSetsMut {
 impl TriePrefixSetsMut {
     /// Extends prefix sets with contents of another prefix set.
     pub fn extend(&mut self, other: Self) {
-        self.account_prefix_set.extend(other.account_prefix_set.keys);
+        self.account_prefix_set.extend(other.account_prefix_set);
         for (hashed_address, prefix_set) in other.storage_prefix_sets {
-            self.storage_prefix_sets.entry(hashed_address).or_default().extend(prefix_set.keys);
+            self.storage_prefix_sets.entry(hashed_address).or_default().extend(prefix_set);
         }
         self.destroyed_accounts.extend(other.destroyed_accounts);
     }
@@ -115,12 +115,18 @@ impl PrefixSetMut {
         self.keys.push(nibbles);
     }
 
+    /// Extend prefix set with contents of another prefix set.
+    pub fn extend(&mut self, other: Self) {
+        self.all |= other.all;
+        self.keys.extend(other.keys);
+    }
+
     /// Extend prefix set keys with contents of provided iterator.
-    pub fn extend<I>(&mut self, nibbles_iter: I)
+    pub fn extend_keys<I>(&mut self, keys: I)
     where
         I: IntoIterator<Item = Nibbles>,
     {
-        self.keys.extend(nibbles_iter);
+        self.keys.extend(keys);
     }
 
     /// Returns the number of elements in the set.
@@ -270,4 +276,11 @@ mod tests {
         assert_eq!(prefix_set.keys.len(), 3); // Length should be 3 (excluding duplicate)
         assert_eq!(prefix_set.keys.capacity(), 3); // Capacity should be 3 after shrinking
     }
+
+    #[test]
+    fn test_prefix_set_all_extend() {
+        let mut prefix_set_mut = PrefixSetMut::default();
+        prefix_set_mut.extend(PrefixSetMut::all());
+        assert!(prefix_set_mut.all);
+    }
 }
diff --git a/crates/trie/trie/src/proof.rs b/crates/trie/trie/src/proof.rs
index bff681a349..d31d63fd9a 100644
--- a/crates/trie/trie/src/proof.rs
+++ b/crates/trie/trie/src/proof.rs
@@ -96,7 +96,7 @@ where
 
         // Create the walker.
         let mut prefix_set = self.prefix_sets.account_prefix_set.clone();
-        prefix_set.extend(targets.keys().map(Nibbles::unpack));
+        prefix_set.extend_keys(targets.keys().map(Nibbles::unpack));
         let walker = TrieWalker::new(trie_cursor, prefix_set.freeze());
 
         // Create a hash builder to rebuild the root node since it is not available in the database.
@@ -225,7 +225,7 @@ where
         }
 
         let target_nibbles = targets.into_iter().map(Nibbles::unpack).collect::<Vec<_>>();
-        self.prefix_set.extend(target_nibbles.clone());
+        self.prefix_set.extend_keys(target_nibbles.clone());
 
         let trie_cursor = self.trie_cursor_factory.storage_trie_cursor(self.hashed_address)?;
         let walker = TrieWalker::new(trie_cursor, self.prefix_set.freeze());