feat: add construct_prefix_sets methods for sorted state types

Add construct_prefix_sets() to HashedPostStateSorted and construct_prefix_set()
to HashedStorageSorted to support the append_ref method on TrieInputSorted.

These methods efficiently iterate over already-sorted data to build prefix sets
needed for trie computation, without requiring any sorting operations.

Also fix BuiltPayloadExecutedBlock conversion to correctly use Either::Right for
sorted variants and convert unsorted to sorted when needed.

Author: yongkangc

crates/optimism/payload/src/builder.rs

@@ -386,8 +386,8 @@ impl<Txs> OpBuilder<'_, Txs> {
         let executed: BuiltPayloadExecutedBlock<N> = BuiltPayloadExecutedBlock {
             recovered_block: Arc::new(block),
             execution_output: Arc::new(execution_outcome),
-            hashed_state: either::Either::Left(Arc::new(hashed_state.into_sorted())),
-            trie_updates: either::Either::Left(Arc::new(trie_updates.into_sorted())),
+            hashed_state: either::Either::Right(Arc::new(hashed_state.into_sorted())),
+            trie_updates: either::Either::Right(Arc::new(trie_updates.into_sorted())),
         };
 
         let no_tx_pool = ctx.attributes().no_tx_pool();

crates/payload/primitives/src/traits.rs

@@ -42,17 +42,17 @@ pub struct BuiltPayloadExecutedBlock<N: NodePrimitives> {
 impl<N: NodePrimitives> BuiltPayloadExecutedBlock<N> {
     /// Converts this into an [`reth_chain_state::ExecutedBlock`].
     ///
-    /// If the hashed state or trie updates are in sorted form, they will be converted
-    /// back to their unsorted representations.
+    /// If the hashed state or trie updates are in unsorted form, they will be converted
+    /// to their sorted representations.
     pub fn into_executed_payload(self) -> reth_chain_state::ExecutedBlock<N> {
         let hashed_state = match self.hashed_state {
-            Either::Left(unsorted) => unsorted,
-            Either::Right(sorted) => Arc::new(Arc::unwrap_or_clone(sorted).into()),
+            Either::Left(unsorted) => Arc::new(Arc::unwrap_or_clone(unsorted).into_sorted()),
+            Either::Right(sorted) => sorted,
         };
 
         let trie_updates = match self.trie_updates {
-            Either::Left(unsorted) => unsorted,
-            Either::Right(sorted) => Arc::new(Arc::unwrap_or_clone(sorted).into()),
+            Either::Left(unsorted) => Arc::new(Arc::unwrap_or_clone(unsorted).into_sorted()),
+            Either::Right(sorted) => sorted,
         };
 
         reth_chain_state::ExecutedBlock {

crates/trie/common/src/hashed_state.rs

@@ -560,6 +560,38 @@ impl HashedPostStateSorted {
         self.accounts.destroyed_accounts.clear();
         self.storages.clear();
     }
+
+    /// Construct [`TriePrefixSetsMut`] from sorted hashed post state.
+    /// The prefix sets contain the hashed account and storage keys that have been changed in the
+    /// post state.
+    pub fn construct_prefix_sets(&self) -> TriePrefixSetsMut {
+        // Populate account prefix set from both live and destroyed accounts.
+        let mut account_prefix_set = PrefixSetMut::with_capacity(
+            self.accounts.accounts.len() + self.accounts.destroyed_accounts.len(),
+        );
+
+        for (hashed_address, _) in &self.accounts.accounts {
+            account_prefix_set.insert(Nibbles::unpack(hashed_address));
+        }
+
+        for hashed_address in &self.accounts.destroyed_accounts {
+            account_prefix_set.insert(Nibbles::unpack(hashed_address));
+        }
+
+        // Populate storage prefix sets.
+        let mut storage_prefix_sets =
+            HashMap::with_capacity_and_hasher(self.storages.len(), Default::default());
+        for (hashed_address, hashed_storage) in &self.storages {
+            account_prefix_set.insert(Nibbles::unpack(hashed_address));
+            storage_prefix_sets.insert(*hashed_address, hashed_storage.construct_prefix_set());
+        }
+
+        TriePrefixSetsMut {
+            account_prefix_set,
+            storage_prefix_sets,
+            destroyed_accounts: self.accounts.destroyed_accounts.clone(),
+        }
+    }
 }
 
 impl AsRef<Self> for HashedPostStateSorted {
@@ -659,6 +691,24 @@ impl HashedStorageSorted {
         // Merge zero valued slots sets
         self.zero_valued_slots.extend(&other.zero_valued_slots);
     }
+
+    /// Construct [`PrefixSetMut`] from sorted hashed storage.
+    pub fn construct_prefix_set(&self) -> PrefixSetMut {
+        if self.wiped {
+            PrefixSetMut::all()
+        } else {
+            let mut prefix_set = PrefixSetMut::with_capacity(
+                self.non_zero_valued_slots.len() + self.zero_valued_slots.len(),
+            );
+            for (hashed_slot, _) in &self.non_zero_valued_slots {
+                prefix_set.insert(Nibbles::unpack(hashed_slot));
+            }
+            for hashed_slot in &self.zero_valued_slots {
+                prefix_set.insert(Nibbles::unpack(hashed_slot));
+            }
+            prefix_set
+        }
+    }
 }
 
 impl From<HashedStorageSorted> for HashedStorage {