feat: backfill address map

crates/optimism/trie/src/api.rs

@@ -2,7 +2,7 @@
 
 use crate::OpProofsStorageResult;
 use alloy_eips::eip1898::BlockWithParent;
-use alloy_primitives::{map::HashMap, B256, U256};
+use alloy_primitives::{map::HashMap, Address, B256, U256};
 use auto_impl::auto_impl;
 use derive_more::{AddAssign, Constructor};
 use reth_primitives_traits::Account;
@@ -110,6 +110,12 @@ pub trait OpProofsStore: Send + Sync + Debug {
         storages: Vec<(B256, U256)>,
     ) -> impl Future<Output = OpProofsStorageResult<()>> + Send;
 
+    /// Store a batch of address mappings from hashed to original addresses.
+    fn store_address_mappings(
+        &self,
+        mappings: Vec<(B256, Address)>,
+    ) -> impl Future<Output = OpProofsStorageResult<()>> + Send;
+
     /// Get the earliest block number and hash that has been stored
     ///
     /// This is used to determine the block number of trie nodes with block number 0.

crates/optimism/trie/src/backfill.rs

@@ -1,7 +1,7 @@
 //! Backfill job for proofs storage. Handles storing the existing state into the proofs storage.
 
 use crate::OpProofsStore;
-use alloy_primitives::B256;
+use alloy_primitives::{keccak256, Address, B256};
 use reth_db::{
     cursor::{DbCursorRO, DbDupCursorRO},
     tables,
@@ -82,6 +82,7 @@ macro_rules! define_dup_cursor_iter {
 
 // Generate iterators for all 4 table types
 define_simple_cursor_iter!(HashedAccountsIter, tables::HashedAccounts, B256, Account);
+define_simple_cursor_iter!(AddressLookupIter, tables::PlainAccountState, Address, Account);
 define_dup_cursor_iter!(HashedStoragesIter, tables::HashedStorages, B256, StorageEntry);
 define_simple_cursor_iter!(
     AccountsTrieIter,
@@ -263,6 +264,29 @@ impl<'a, Tx: DbTx, S: OpProofsStore + Send> BackfillJob<'a, Tx, S> {
         Ok(())
     }
 
+    /// Backfill address mappings data
+    async fn backfill_address_mappings(&self) -> eyre::Result<()> {
+        let start_cursor = self.tx.cursor_read::<tables::PlainAccountState>()?;
+
+        let source = AddressLookupIter::new(start_cursor)
+            .map(|res| res.map(|(addr, _)| (keccak256(addr), addr)));
+        let save_fn = async |entries: Vec<(B256, Address)>| -> eyre::Result<()> {
+            self.storage.store_address_mappings(entries).await?;
+            Ok(())
+        };
+
+        backfill(
+            "address mappings",
+            source,
+            BACKFILL_STORAGE_THRESHOLD,
+            BACKFILL_LOG_THRESHOLD,
+            save_fn,
+        )
+        .await?;
+
+        Ok(())
+    }
+
     /// Backfill accounts trie data
     async fn backfill_accounts_trie(&self) -> eyre::Result<()> {
         let start_cursor = self.tx.cursor_read::<tables::AccountsTrie>()?;
@@ -328,6 +352,7 @@ impl<'a, Tx: DbTx, S: OpProofsStore + Send> BackfillJob<'a, Tx, S> {
     async fn backfill_trie(&self) -> eyre::Result<()> {
         self.backfill_hashed_accounts().await?;
         self.backfill_hashed_storages().await?;
+        self.backfill_address_mappings().await?;
         self.backfill_storages_trie().await?;
         self.backfill_accounts_trie().await?;
 
@@ -520,6 +545,34 @@ mod tests {
         assert_eq!(count, 3);
     }
 
+    #[tokio::test]
+    async fn test_backfill_address_mappings() {
+        let db = create_test_rw_db();
+        let storage = InMemoryProofsStorage::new();
+
+        // Insert test address mappings into database
+        let tx = db.tx_mut().unwrap();
+        let mut cursor = tx.cursor_write::<tables::PlainAccountState>().unwrap();
+
+        let accounts = vec![
+            (Address::repeat_byte(0x01), Account { nonce: 1, ..Default::default() }),
+            (Address::repeat_byte(0x02), Account { nonce: 2, ..Default::default() }),
+        ];
+
+        for (addr, account) in &accounts {
+            cursor.append(*addr, account).unwrap();
+        }
+        drop(cursor);
+        tx.commit().unwrap();
+
+        // Run backfill
+        let tx = db.tx().unwrap();
+        let job = BackfillJob::new(storage.clone(), &tx);
+        job.backfill_address_mappings().await.unwrap();
+
+        // Todo: Verify data was stored
+    }
+
     #[tokio::test]
     async fn test_backfill_storages_trie() {
         let db = create_test_rw_db();

crates/optimism/trie/src/db/models/mod.rs

@@ -16,7 +16,7 @@ pub(crate) mod kv;
 pub use change_set::*;
 pub use kv::*;
 
-use alloy_primitives::B256;
+use alloy_primitives::{Address, B256};
 use reth_db::{
     table::{DupSort, TableInfo},
     tables, TableSet, TableType, TableViewer,
@@ -81,4 +81,10 @@ tables! {
         type Key = u64; // Block number
         type Value = ChangeSet;
     }
+
+    /// A mapping table from hashed addresses to their original addresses.
+    table AddressLookup {
+        type Key = B256;
+        type Value = Address;
+    }
 }

crates/optimism/trie/src/db/store.rs

@@ -4,16 +4,16 @@ use crate::{
     db::{
         cursor::Dup,
         models::{
-            kv::IntoKV, AccountTrieHistory, BlockChangeSet, ChangeSet, HashedAccountHistory,
-            HashedStorageHistory, HashedStorageKey, MaybeDeleted, StorageTrieHistory,
-            StorageTrieKey, StorageValue, VersionedValue,
+            kv::IntoKV, AccountTrieHistory, AddressLookup, BlockChangeSet, ChangeSet,
+            HashedAccountHistory, HashedStorageHistory, HashedStorageKey, MaybeDeleted,
+            StorageTrieHistory, StorageTrieKey, StorageValue, VersionedValue,
         },
         MdbxAccountCursor, MdbxStorageCursor, MdbxTrieCursor,
     },
     BlockStateDiff, OpProofsStorageError, OpProofsStorageResult, OpProofsStore,
 };
 use alloy_eips::{eip1898::BlockWithParent, NumHash};
-use alloy_primitives::{map::HashMap, B256, U256};
+use alloy_primitives::{map::HashMap, Address, B256, U256};
 #[cfg(feature = "metrics")]
 use metrics::{gauge, Label};
 use reth_db::{
@@ -514,6 +514,27 @@ impl OpProofsStore for MdbxProofsStorage {
         })?
     }
 
+    async fn store_address_mappings(
+        &self,
+        mappings: Vec<(B256, Address)>,
+    ) -> OpProofsStorageResult<()> {
+        let mut mappings = mappings;
+        if mappings.is_empty() {
+            return Ok(());
+        }
+
+        // sort the mappings by key to ensure insertion is efficient
+        mappings.sort_by_key(|(key, _)| *key);
+
+        self.env.update(|tx| {
+            let mut cur = tx.cursor_write::<AddressLookup>()?;
+            for (k, v) in mappings {
+                cur.append(k, &v)?;
+            }
+            Ok(())
+        })?
+    }
+
     async fn get_earliest_block_number(&self) -> OpProofsStorageResult<Option<(u64, B256)>> {
         self.env.view(|tx| self.inner_get_block_number_hash(tx, ProofWindowKey::EarliestBlock))?
     }
@@ -901,7 +922,7 @@ mod tests {
         StorageTrieKey,
     };
     use alloy_eips::NumHash;
-    use alloy_primitives::B256;
+    use alloy_primitives::{keccak256, B256};
     use reth_db::{
         cursor::DbDupCursorRO,
         transaction::{DbTx, DbTxMut},
@@ -1322,6 +1343,40 @@ mod tests {
         }
     }
 
+    #[tokio::test]
+    async fn test_store_address_mappings() {
+        let dir = TempDir::new().unwrap();
+        let store = MdbxProofsStorage::new(dir.path()).expect("env");
+
+        let a1 = Address::random();
+        let h1 = keccak256(a1);
+        let a2 = Address::random();
+        let h2 = keccak256(a2);
+        let a3 = Address::random();
+        let h3 = keccak256(a3);
+
+        // Input is unsorted to verify the method sorts them before appending
+        // (MDBX append requires sorted keys)
+        let mappings = vec![(h3, a3), (h1, a1), (h2, a2)];
+
+        store.store_address_mappings(mappings).await.expect("store");
+
+        let tx = store.env.tx().expect("ro tx");
+        let mut cur = tx.cursor_read::<AddressLookup>().expect("cursor");
+
+        // Verify h1
+        let v1 = cur.seek_exact(h1).expect("seek").expect("exists");
+        assert_eq!(v1, (h1, a1));
+
+        // Verify h2
+        let v2 = cur.seek_exact(h2).expect("seek").expect("exists");
+        assert_eq!(v2, (h2, a2));
+
+        // Verify h3
+        let v3 = cur.seek_exact(h3).expect("seek").expect("exists");
+        assert_eq!(v3, (h3, a3));
+    }
+
     #[tokio::test]
     async fn test_store_trie_updates_comprehensive() {
         let dir = TempDir::new().unwrap();

crates/optimism/trie/src/in_memory.rs

@@ -4,7 +4,7 @@ use crate::{
     api::WriteCounts, BlockStateDiff, OpProofsStorageError, OpProofsStorageResult, OpProofsStore,
 };
 use alloy_eips::eip1898::BlockWithParent;
-use alloy_primitives::{map::HashMap, B256, U256};
+use alloy_primitives::{map::HashMap, Address, B256, U256};
 use reth_db::DatabaseError;
 use reth_primitives_traits::Account;
 use reth_trie::{
@@ -37,6 +37,9 @@ struct InMemoryStorageInner {
     /// Hashed storages: (`block_number`, `hashed_address`, `hashed_slot`) -> value
     hashed_storages: BTreeMap<(u64, B256, B256), U256>,
 
+    /// Address mappings: (`block_number`, `hashed_address`) -> original address
+    address_mappings: BTreeMap<B256, Address>,
+
     /// Trie updates by block number
     trie_updates: BTreeMap<u64, TrieUpdatesSorted>,
 
@@ -538,6 +541,19 @@ impl OpProofsStore for InMemoryProofsStorage {
         Ok(())
     }
 
+    async fn store_address_mappings(
+        &self,
+        mappings: Vec<(B256, alloy_primitives::Address)>,
+    ) -> OpProofsStorageResult<()> {
+        let mut inner = self.inner.write().await;
+
+        for (hashed_address, original_address) in mappings {
+            inner.address_mappings.insert(hashed_address, original_address);
+        }
+
+        Ok(())
+    }
+
     async fn get_earliest_block_number(&self) -> OpProofsStorageResult<Option<(u64, B256)>> {
         let inner = self.inner.read().await;
         Ok(inner.earliest_block)

crates/optimism/trie/src/metrics.rs

@@ -51,6 +51,8 @@ pub enum StorageOperation {
     StoreHashedAccount,
     /// Store hashed storage
     StoreHashedStorage,
+    /// Store address mapping
+    StoreAddressMapping,
     /// Trie cursor seek exact operation
     TrieCursorSeekExact,
     /// Trie cursor seek
@@ -73,6 +75,7 @@ impl StorageOperation {
             Self::StoreStorageBranch => "store_storage_branch",
             Self::StoreHashedAccount => "store_hashed_account",
             Self::StoreHashedStorage => "store_hashed_storage",
+            Self::StoreAddressMapping => "store_address_mapping",
             Self::TrieCursorSeekExact => "trie_cursor_seek_exact",
             Self::TrieCursorSeek => "trie_cursor_seek",
             Self::TrieCursorNext => "trie_cursor_next",
@@ -469,6 +472,27 @@ where
         result
     }
 
+    async fn store_address_mappings(
+        &self,
+        mappings: Vec<(B256, alloy_primitives::Address)>,
+    ) -> OpProofsStorageResult<()> {
+        let count = mappings.len();
+        let start = Instant::now();
+        let result = self.storage.store_address_mappings(mappings).await;
+        let duration = start.elapsed();
+
+        // Record per-item duration
+        if count > 0 {
+            self.metrics.record_duration_per_item(
+                StorageOperation::StoreAddressMapping,
+                duration,
+                count,
+            );
+        }
+
+        result
+    }
+
     #[inline]
     async fn get_earliest_block_number(&self) -> OpProofsStorageResult<Option<(u64, B256)>> {
         self.storage.get_earliest_block_number().await