feat: create index table to efficiently prune history tables

crates/optimism/trie/Cargo.toml

@@ -35,6 +35,7 @@ tokio = { workspace = true, features = ["sync"] }
 # codec
 bytes.workspace = true
 serde.workspace = true
+bincode.workspace = true
 
 # misc
 thiserror.workspace = true

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

@@ -0,0 +1,125 @@
+use crate::db::{HashedStorageKey, StorageTrieKey};
+use alloy_primitives::B256;
+use reth_db::{
+    table::{self, Decode, Encode},
+    DatabaseError,
+};
+use reth_trie::StoredNibbles;
+use serde::{Deserialize, Serialize};
+
+/// The keys of the entries in the history tables.
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
+pub struct ChangeSet {
+    /// Keys changed in [`AccountTrieHistory`](super::AccountTrieHistory) table.
+    pub account_trie_keys: Vec<StoredNibbles>,
+    /// Keys changed in [`StorageTrieHistory`](super::StorageTrieHistory) table.
+    pub storage_trie_keys: Vec<StorageTrieKey>,
+    /// Keys changed in [`HashedAccountHistory`](super::HashedAccountHistory) table.
+    pub hashed_account_keys: Vec<B256>,
+    /// Keys changed in [`HashedStorageHistory`](super::HashedStorageHistory) table.
+    pub hashed_storage_keys: Vec<HashedStorageKey>,
+}
+
+impl table::Encode for ChangeSet {
+    type Encoded = Vec<u8>;
+
+    fn encode(self) -> Self::Encoded {
+        bincode::serialize(&self).expect("ChangeSet serialization should not fail")
+    }
+}
+
+impl table::Decode for ChangeSet {
+    fn decode(value: &[u8]) -> Result<Self, DatabaseError> {
+        bincode::deserialize(value).map_err(|_| DatabaseError::Decode)
+    }
+}
+
+impl table::Compress for ChangeSet {
+    type Compressed = Vec<u8>;
+
+    fn compress_to_buf<B: bytes::BufMut + AsMut<[u8]>>(&self, buf: &mut B) {
+        let encoded = self.clone().encode();
+        buf.put_slice(&encoded);
+    }
+}
+
+impl table::Decompress for ChangeSet {
+    fn decompress(value: &[u8]) -> Result<Self, DatabaseError> {
+        Self::decode(value)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use alloy_primitives::B256;
+    use reth_db::table::{Compress, Decompress};
+
+    #[test]
+    fn test_encode_decode_empty_change_set() {
+        let change_set = ChangeSet {
+            account_trie_keys: vec![],
+            storage_trie_keys: vec![],
+            hashed_account_keys: vec![],
+            hashed_storage_keys: vec![],
+        };
+
+        let encoded = change_set.clone().encode();
+        let decoded = ChangeSet::decode(&encoded).expect("Failed to decode");
+        assert_eq!(change_set, decoded);
+    }
+
+    #[test]
+    fn test_encode_decode_populated_change_set() {
+        let account_key = StoredNibbles::from(vec![1, 2, 3, 4]);
+        let storage_key = StorageTrieKey {
+            hashed_address: B256::repeat_byte(0x11),
+            path: StoredNibbles::from(vec![5, 6, 7, 8]),
+        };
+        let hashed_storage_key = HashedStorageKey {
+            hashed_address: B256::repeat_byte(0x22),
+            hashed_storage_key: B256::repeat_byte(0x33),
+        };
+
+        let change_set = ChangeSet {
+            account_trie_keys: vec![account_key],
+            storage_trie_keys: vec![storage_key],
+            hashed_account_keys: vec![B256::repeat_byte(0x44)],
+            hashed_storage_keys: vec![hashed_storage_key],
+        };
+
+        let encoded = change_set.clone().encode();
+        let decoded = ChangeSet::decode(&encoded).expect("Failed to decode");
+        assert_eq!(change_set, decoded);
+    }
+
+    #[test]
+    fn test_decode_invalid_data() {
+        let invalid_data = vec![0xFF; 32];
+        let result = ChangeSet::decode(&invalid_data);
+        assert!(result.is_err());
+        assert!(matches!(result.unwrap_err(), DatabaseError::Decode));
+    }
+
+    #[test]
+    fn test_compress_decompress() {
+        let change_set = ChangeSet {
+            account_trie_keys: vec![StoredNibbles::from(vec![1, 2, 3])],
+            storage_trie_keys: vec![StorageTrieKey {
+                hashed_address: B256::ZERO,
+                path: StoredNibbles::from(vec![4, 5, 6]),
+            }],
+            hashed_account_keys: vec![B256::ZERO],
+            hashed_storage_keys: vec![HashedStorageKey {
+                hashed_address: B256::ZERO,
+                hashed_storage_key: B256::repeat_byte(0x42),
+            }],
+        };
+
+        let mut buf = Vec::new();
+        change_set.compress_to_buf(&mut buf);
+
+        let decompressed = ChangeSet::decompress(&buf).expect("Failed to decompress");
+        assert_eq!(change_set, decompressed);
+    }
+}

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

@@ -11,6 +11,8 @@ mod version;
 pub use version::*;
 mod storage;
 pub use storage::*;
+mod change_set;
+pub use change_set::*;
 
 use alloy_primitives::B256;
 use reth_db::{
@@ -70,4 +72,11 @@ tables! {
       type Key = ProofWindowKey;
       type Value = BlockNumberHash;
     }
+
+    /// A reverse mapping of block numbers to a keys of the tables.
+    /// This is used for efficiently locating data by block number.
+    table BlockChangeSet {
+        type Key = u64; // Block number
+        type Value = ChangeSet;
+    }
 }