feat(trie): add revert functionality to state tries (#374)

## Summary
- Add `revert_to` method to `ClassesTrie`, `ContractsTrie`, `StoragesTrie`, and `BonsaiTrie` for unwinding trie state
- Implement `get_by_prefix` for `TrieDb` to support trie prefix operations
- Update bonsai-trie dependency to rev `95de4c6` for revert support

---------

Co-authored-by: Claude <noreply@anthropic.com>

Author: kariy

Cargo.lock

@@ -173,9 +173,22 @@ where
 
     fn get_by_prefix(
         &self,
-        _: &DatabaseKey<'_>,
+        prefix: &DatabaseKey<'_>,
     ) -> Result<Vec<(ByteVec, ByteVec)>, Self::DatabaseError> {
-        todo!()
+        let mut results = Vec::new();
+
+        let mut cursor = self.tx.cursor::<Tb>()?;
+        let walker = cursor.walk(None)?;
+
+        for entry in walker {
+            let (TrieDatabaseKey { key, .. }, value) = entry?;
+
+            if key.starts_with(prefix.as_slice()) {
+                results.push((key.to_smallvec(), value));
+            }
+        }
+
+        Ok(results)
     }
 
     fn insert(
@@ -300,8 +313,7 @@ where
         &self,
         prefix: &DatabaseKey<'_>,
     ) -> Result<Vec<(ByteVec, ByteVec)>, Self::DatabaseError> {
-        let _ = prefix;
-        todo!()
+        TrieDb::<Tb, Tx>::new(self.tx.clone()).get_by_prefix(prefix)
     }
 
     fn insert(
@@ -509,4 +521,67 @@ mod tests {
             assert_eq!(vec![value0, value1], result);
         }
     }
+
+    #[test]
+    fn revert_to() {
+        let db = test_utils::create_test_db();
+        let db_tx = db.tx_mut().expect("failed to get tx");
+
+        let mut trie = ClassesTrie::new(TrieDbMut::<tables::ClassesTrie, _>::new(&db_tx));
+
+        // Insert values at block 0
+        trie.insert(felt!("0x1"), felt!("0x100"));
+        trie.insert(felt!("0x2"), felt!("0x200"));
+        trie.commit(0);
+        let root_at_block_0 = trie.root();
+
+        // Insert more values at block 1
+        trie.insert(felt!("0x3"), felt!("0x300"));
+        trie.insert(felt!("0x4"), felt!("0x400"));
+        trie.commit(1);
+        let root_at_block_1 = trie.root();
+
+        // Roots should be different
+        assert_ne!(root_at_block_0, root_at_block_1);
+
+        // Insert even more values at block 2
+        trie.insert(felt!("0x5"), felt!("0x500"));
+        trie.commit(2);
+        let root_at_block_2 = trie.root();
+
+        // Roots should be different
+        assert_ne!(root_at_block_1, root_at_block_2);
+        assert_ne!(root_at_block_0, root_at_block_2);
+
+        // Revert to block 1
+        trie.revert_to(1, 2);
+        let root_after_revert = trie.root();
+
+        // After revert, root should match block 1
+        assert_eq!(root_after_revert, root_at_block_1);
+
+        // Revert to block 0
+        trie.revert_to(0, 1);
+        let root_after_second_revert = trie.root();
+
+        // After revert, root should match block 0
+        assert_eq!(root_after_second_revert, root_at_block_0);
+
+        // Insert more values at block 1
+        trie.insert(felt!("0x3"), felt!("0x300"));
+        trie.insert(felt!("0x4"), felt!("0x400"));
+        trie.commit(1);
+        let root_at_block_1_after_insert = trie.root();
+
+        // After insertion, root should match block 1
+        assert_eq!(root_at_block_1_after_insert, root_at_block_1);
+
+        // Insert even more values at block 2
+        trie.insert(felt!("0x5"), felt!("0x500"));
+        trie.commit(2);
+        let root_at_block_2_after_insert = trie.root();
+
+        // After insertion, root should match block 2
+        assert_eq!(root_at_block_2_after_insert, root_at_block_2);
+    }
 }

crates/storage/db/src/trie/mod.rs

@@ -17,7 +17,7 @@ starknet-types-core.workspace = true
 thiserror.workspace = true
 
 [dependencies.bonsai-trie]
-rev = "8b509fd"
+rev = "95de4c6"
 default-features = false
 features = [ "std" ]
 git = "https://github.com/dojoengine/bonsai-trie/"

crates/trie/Cargo.toml

@@ -47,6 +47,10 @@ impl<DB: BonsaiDatabase> ClassesTrie<DB> {
     pub fn multiproof(&mut self, class_hashes: Vec<ClassHash>) -> MultiProof {
         self.trie.multiproof(Self::BONSAI_IDENTIFIER, class_hashes)
     }
+
+    pub fn revert_to(&mut self, block: BlockNumber, latest_block: BlockNumber) {
+        self.trie.revert_to(block, latest_block);
+    }
 }
 
 impl<DB> ClassesTrie<DB>

crates/trie/src/classes.rs

@@ -31,6 +31,10 @@ impl<DB: BonsaiDatabase> ContractsTrie<DB> {
         let keys = addresses.into_iter().map(Felt::from).collect::<Vec<Felt>>();
         self.trie.multiproof(Self::BONSAI_IDENTIFIER, keys)
     }
+
+    pub fn revert_to(&mut self, block: BlockNumber, latest_block: BlockNumber) {
+        self.trie.revert_to(block, latest_block);
+    }
 }
 
 impl<DB> ContractsTrie<DB>

crates/trie/src/contracts.rs

@@ -3,6 +3,7 @@ pub use bonsai::{BitVec, MultiProof, Path, ProofNode};
 pub use bonsai_trie::databases::HashMapDb;
 use bonsai_trie::BonsaiStorage;
 pub use bonsai_trie::{BonsaiDatabase, BonsaiPersistentDatabase, BonsaiStorageConfig};
+use katana_primitives::block::BlockNumber;
 use katana_primitives::class::ClassHash;
 use katana_primitives::Felt;
 use starknet_types_core::hash::{Pedersen, StarkHash};
@@ -39,14 +40,23 @@ where
 {
     pub fn new(db: DB) -> Self {
         let config = BonsaiStorageConfig {
-            // we have our own implementation of storing trie changes
-            max_saved_trie_logs: Some(0),
+            // This field controls what's the oldest block we can revert to.
+            //
+            // The value 5 is chosen arbitrarily as a placeholder. This value should be
+            // configurable.
+            max_saved_trie_logs: Some(5),
+
             // in the bonsai-trie crate, this field seems to be only used in rocksdb impl.
             // i dont understand why would they add a config thats implementation specific ????
             //
             // this config should be used by our implementation of the
             // BonsaiPersistentDatabase::snapshot()
+            //
+            // note: currently, this value is not being used for anything. our trie will stores
+            // all created snapshots.
             max_saved_snapshots: Some(64usize),
+
+            // creates a snapshot for every block
             snapshot_interval: 1,
         };
 
@@ -61,6 +71,10 @@ where
         let keys = keys.into_iter().map(|key| key.to_bytes_be().as_bits()[5..].to_owned());
         self.storage.get_multi_proof(id, keys).expect("failed to get multiproof")
     }
+
+    pub fn revert_to(&mut self, block: BlockNumber, latest_block: BlockNumber) {
+        self.storage.revert_to(block.into(), latest_block.into()).expect("failed to revert trie");
+    }
 }
 
 impl<DB, Hash> BonsaiTrie<DB, Hash>
@@ -169,4 +183,54 @@ mod tests {
 
         assert_eq!(result, expected);
     }
+
+    #[test]
+    fn test_revert_to() {
+        use bonsai_trie::databases;
+
+        // the identifier for the trie
+        const IDENTIFIER: &[u8] = b"test_trie";
+
+        // Create a BonsaiStorage with in-memory database and trie logs enabled
+        let bonsai_db = databases::HashMapDb::<CommitId>::default();
+        let mut trie = BonsaiTrie::<_, hash::Pedersen>::new(bonsai_db);
+
+        // Insert values at block 0
+        trie.insert(IDENTIFIER, Felt::from(1), Felt::from(100));
+        trie.insert(IDENTIFIER, Felt::from(2), Felt::from(200));
+        trie.commit(0.into());
+        let root_at_block_0 = trie.root(IDENTIFIER);
+
+        // Insert more values at block 1
+        trie.insert(IDENTIFIER, Felt::from(3), Felt::from(300));
+        trie.insert(IDENTIFIER, Felt::from(4), Felt::from(400));
+        trie.commit(1.into());
+        let root_at_block_1 = trie.root(IDENTIFIER);
+
+        // Roots should be different
+        assert_ne!(root_at_block_0, root_at_block_1);
+
+        // Insert even more values at block 2
+        trie.insert(IDENTIFIER, Felt::from(5), Felt::from(500));
+        trie.commit(2.into());
+        let root_at_block_2 = trie.root(IDENTIFIER);
+
+        // Roots should be different
+        assert_ne!(root_at_block_1, root_at_block_2);
+        assert_ne!(root_at_block_0, root_at_block_2);
+
+        // Revert to block 1
+        trie.revert_to(1, 2);
+        let root_after_revert = trie.root(IDENTIFIER);
+
+        // After revert, root should match block 1
+        assert_eq!(root_after_revert, root_at_block_1);
+
+        // Revert to block 0
+        trie.revert_to(0, 1);
+        let root_after_second_revert = trie.root(IDENTIFIER);
+
+        // After revert, root should match block 0
+        assert_eq!(root_after_second_revert, root_at_block_0);
+    }
 }

crates/trie/src/lib.rs

@@ -25,6 +25,10 @@ impl<DB: BonsaiDatabase> StoragesTrie<DB> {
     pub fn multiproof(&mut self, storage_keys: Vec<StorageKey>) -> MultiProof {
         self.trie.multiproof(&self.address.to_bytes_be(), storage_keys)
     }
+
+    pub fn revert_to(&mut self, block: BlockNumber, latest_block: BlockNumber) {
+        self.trie.revert_to(block, latest_block);
+    }
 }
 
 impl<DB> StoragesTrie<DB>