From 95d7096919cff64d20407527c2dc171f1cc9a5d1 Mon Sep 17 00:00:00 2001
From: Roman Krasiuk <rokrassyuk@gmail.com>
Date: Tue, 30 Jul 2024 23:29:46 +0200
Subject: [PATCH 1/2] feat(trie): database trie witness

---
 crates/trie/db/src/lib.rs     |  2 ++
 crates/trie/db/src/witness.rs | 46 +++++++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+)
 create mode 100644 crates/trie/db/src/witness.rs

diff --git a/crates/trie/db/src/lib.rs b/crates/trie/db/src/lib.rs
index 21deac82d27..8369812e4b2 100644
--- a/crates/trie/db/src/lib.rs
+++ b/crates/trie/db/src/lib.rs
@@ -3,7 +3,9 @@
 mod proof;
 mod state;
 mod storage;
+mod witness;
 
 pub use proof::DatabaseProof;
 pub use state::DatabaseStateRoot;
 pub use storage::DatabaseStorageRoot;
+pub use witness::DatabaseTrieWitness;
diff --git a/crates/trie/db/src/witness.rs b/crates/trie/db/src/witness.rs
new file mode 100644
index 00000000000..5ba4b9d5571
--- /dev/null
+++ b/crates/trie/db/src/witness.rs
@@ -0,0 +1,46 @@
+use reth_db_api::transaction::DbTx;
+use reth_execution_errors::TrieWitnessError;
+use reth_primitives::{Bytes, B256};
+use reth_trie::{
+    hashed_cursor::{DatabaseHashedCursorFactory, HashedPostStateCursorFactory},
+    trie_cursor::DatabaseTrieCursorFactory,
+    witness::TrieWitness,
+    HashedPostState,
+};
+use std::collections::HashMap;
+
+/// Extends [`TrieWitness`] with operations specific for working with a database transaction.
+pub trait DatabaseTrieWitness<'a, TX> {
+    /// Create a new [TrieWitness] from database transaction.
+    fn from_tx(tx: &'a TX) -> Self;
+
+    /// Generates trie witness for target state on top of this [`HashedPostState`].
+    fn overlay_witness(
+        tx: &'a TX,
+        post_state: HashedPostState,
+        target: HashedPostState,
+    ) -> Result<HashMap<B256, Bytes>, TrieWitnessError>;
+}
+
+impl<'a, TX: DbTx> DatabaseTrieWitness<'a, TX>
+    for TrieWitness<DatabaseTrieCursorFactory<'a, TX>, DatabaseHashedCursorFactory<'a, TX>>
+{
+    fn from_tx(tx: &'a TX) -> Self {
+        Self::new(DatabaseTrieCursorFactory::new(tx), DatabaseHashedCursorFactory::new(tx))
+    }
+
+    fn overlay_witness(
+        tx: &'a TX,
+        post_state: HashedPostState,
+        target: HashedPostState,
+    ) -> Result<HashMap<B256, Bytes>, TrieWitnessError> {
+        let prefix_sets = post_state.construct_prefix_sets();
+        let sorted = post_state.into_sorted();
+        let hashed_cursor_factory =
+            HashedPostStateCursorFactory::new(DatabaseHashedCursorFactory::new(tx), &sorted);
+        Self::from_tx(tx)
+            .with_hashed_cursor_factory(hashed_cursor_factory)
+            .with_prefix_sets_mut(prefix_sets)
+            .compute(target)
+    }
+}

From 75d12b5401c3648b485c917f98e0b816b6b03385 Mon Sep 17 00:00:00 2001
From: Roman Krasiuk <rokrassyuk@gmail.com>
Date: Tue, 30 Jul 2024 14:35:41 -0700
Subject: [PATCH 2/2] Update crates/trie/db/src/witness.rs

---
 crates/trie/db/src/witness.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/crates/trie/db/src/witness.rs b/crates/trie/db/src/witness.rs
index 5ba4b9d5571..12fd3150940 100644
--- a/crates/trie/db/src/witness.rs
+++ b/crates/trie/db/src/witness.rs
@@ -11,7 +11,7 @@ use std::collections::HashMap;
 
 /// Extends [`TrieWitness`] with operations specific for working with a database transaction.
 pub trait DatabaseTrieWitness<'a, TX> {
-    /// Create a new [TrieWitness] from database transaction.
+    /// Create a new [`TrieWitness`] from database transaction.
     fn from_tx(tx: &'a TX) -> Self;
 
     /// Generates trie witness for target state on top of this [`HashedPostState`].