test(katana): trie snapshot db (#2845)

dojoengine · kariy · Dec 31, 2024 · Dec 16, 2024 · Dec 16, 2024 · Dec 18, 2024
commit 7b2e4a7d1f5057a88b90bc7a8808b42c65434b5c
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/katana/storage/db/Cargo.toml b/crates/katana/storage/db/Cargo.toml
@@ -34,6 +34,8 @@ rev = "b34b0d3"
 [dev-dependencies]
 arbitrary.workspace = true
 criterion.workspace = true
+proptest = "1.6.0"
+rstest.workspace = true
 starknet.workspace = true
 
 [features]

diff --git a/crates/katana/storage/db/src/models/trie.rs b/crates/katana/storage/db/src/models/trie.rs
@@ -42,6 +42,23 @@ pub enum TrieDatabaseKeyType {
     TrieLog,
 }
 
+#[derive(Debug, thiserror::Error)]
+#[error("unknown trie key type: {0}")]
+pub struct TrieDatabaseKeyTypeTryFromError(u8);
+
+impl TryFrom<u8> for TrieDatabaseKeyType {
+    type Error = TrieDatabaseKeyTypeTryFromError;
+
+    fn try_from(value: u8) -> Result<Self, Self::Error> {
+        match value {
+            0 => Ok(Self::Trie),
+            1 => Ok(Self::Flat),
+            2 => Ok(Self::TrieLog),
+            invalid => Err(TrieDatabaseKeyTypeTryFromError(invalid)),
+        }
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
 pub struct TrieDatabaseKey {
     pub r#type: TrieDatabaseKeyType,
@@ -68,16 +85,11 @@ impl Decode for TrieDatabaseKey {
 
         if bytes.len() < 2 {
             // Need at least type and length bytes
-            panic!("emptyy buffer")
+            panic!("empty buffer")
         }
 
-        let r#type = match bytes[0] {
-            0 => TrieDatabaseKeyType::Trie,
-            1 => TrieDatabaseKeyType::Flat,
-            2 => TrieDatabaseKeyType::TrieLog,
-            _ => panic!("Invalid trie database key type"),
-        };
-
+        let r#type =
+            TrieDatabaseKeyType::try_from(bytes[0]).expect("Invalid trie database key type");
         let key_len = bytes[1] as usize;
 
         if bytes.len() < 2 + key_len {

diff --git a/crates/katana/storage/db/src/trie/mod.rs b/crates/katana/storage/db/src/trie/mod.rs
@@ -361,7 +361,6 @@ where
 
     // TODO: check if the snapshot exist
     fn transaction(&self, id: CommitId) -> Option<(CommitId, Self::Transaction<'_>)> {
-        dbg!("getting snapshot", id);
         Some((id, SnapshotTrieDb::new(self.tx.clone(), id)))
     }
 }
@@ -379,3 +378,104 @@ fn to_db_key(key: &DatabaseKey<'_>) -> models::trie::TrieDatabaseKey {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use katana_primitives::hash::{Pedersen, StarkHash};
+    use katana_primitives::{felt, hash};
+    use katana_trie::{verify_proof, ClassesTrie, CommitId};
+    use starknet::macros::short_string;
+
+    use super::TrieDbMut;
+    use crate::abstraction::Database;
+    use crate::mdbx::test_utils;
+    use crate::tables;
+    use crate::trie::SnapshotTrieDb;
+
+    #[test]
+    fn snapshot() {
+        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));
+
+        let root0 = {
+            let entries = [
+                (felt!("0x9999"), felt!("0xdead")),
+                (felt!("0x5555"), felt!("0xbeef")),
+                (felt!("0x1337"), felt!("0xdeadbeef")),
+            ];
+
+            for (key, value) in entries {
+                trie.insert(key, value);
+            }
+
+            trie.commit(0);
+            trie.root()
+        };
+
+        let root1 = {
+            let entries = [
+                (felt!("0x6969"), felt!("0x80085")),
+                (felt!("0x3333"), felt!("0x420")),
+                (felt!("0x2222"), felt!("0x7171")),
+            ];
+
+            for (key, value) in entries {
+                trie.insert(key, value);
+            }
+
+            trie.commit(1);
+            trie.root()
+        };
+
+        assert_ne!(root0, root1);
+
+        {
+            let db = SnapshotTrieDb::<tables::ClassesTrie, _>::new(&db_tx, CommitId::new(0));
+            let mut snapshot0 = ClassesTrie::new(db);
+
+            let snapshot_root0 = snapshot0.root();
+            assert_eq!(snapshot_root0, root0);
+
+            let proofs0 = snapshot0.multiproof(vec![felt!("0x9999")]);
+            let verify_result0 =
+                verify_proof::<Pedersen>(&proofs0, snapshot_root0, vec![felt!("0x9999")]);
+
+            let value =
+                hash::Poseidon::hash(&short_string!("CONTRACT_CLASS_LEAF_V0"), &felt!("0xdead"));
+            assert_eq!(vec![value], verify_result0);
+        }
+
+        {
+            let commit = CommitId::new(1);
+            let mut snapshot1 =
+                ClassesTrie::new(SnapshotTrieDb::<tables::ClassesTrie, _>::new(&db_tx, commit));
+
+            let snapshot_root1 = snapshot1.root();
+            assert_eq!(snapshot_root1, root1);
+
+            let proofs1 = snapshot1.multiproof(vec![felt!("0x6969")]);
+            let verify_result1 =
+                verify_proof::<Pedersen>(&proofs1, snapshot_root1, vec![felt!("0x6969")]);
+
+            let value =
+                hash::Poseidon::hash(&short_string!("CONTRACT_CLASS_LEAF_V0"), &felt!("0x80085"));
+            assert_eq!(vec![value], verify_result1);
+        }
+
+        {
+            let root = trie.root();
+            let proofs = trie.multiproof(vec![felt!("0x6969"), felt!("0x9999")]);
+            let result =
+                verify_proof::<Pedersen>(&proofs, root, vec![felt!("0x6969"), felt!("0x9999")]);
+
+            let value0 =
+                hash::Poseidon::hash(&short_string!("CONTRACT_CLASS_LEAF_V0"), &felt!("0x80085"));
+            let value1 =
+                hash::Poseidon::hash(&short_string!("CONTRACT_CLASS_LEAF_V0"), &felt!("0xdead"));
+
+            assert_eq!(vec![value0, value1], result);
+        }
+    }
+}
diff --git a/crates/katana/storage/db/src/trie/snapshot.rs b/crates/katana/storage/db/src/trie/snapshot.rs
@@ -121,3 +121,124 @@ where
         unimplemented!("modifying trie snapshot is not supported")
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashMap;
+
+    use katana_primitives::felt;
+    use katana_trie::bonsai::DatabaseKey;
+    use katana_trie::{BonsaiPersistentDatabase, CommitId};
+    use proptest::prelude::*;
+    use proptest::strategy;
+
+    use super::*;
+    use crate::abstraction::{Database, DbTx};
+    use crate::mdbx::test_utils;
+    use crate::models::trie::TrieDatabaseKeyType;
+    use crate::tables;
+    use crate::trie::{SnapshotTrieDb, TrieDbMut};
+
+    #[allow(unused)]
+    fn arb_db_key_type() -> BoxedStrategy<TrieDatabaseKeyType> {
+        prop_oneof![
+            Just(TrieDatabaseKeyType::Trie),
+            Just(TrieDatabaseKeyType::Flat),
+            Just(TrieDatabaseKeyType::TrieLog),
+        ]
+        .boxed()
+    }
+
+    #[derive(Debug)]
+    struct Case {
+        number: BlockNumber,
+        keyvalues: HashMap<(TrieDatabaseKeyType, [u8; 32]), [u8; 32]>,
+    }
+
+    prop_compose! {
+        // This create a strategy that generates a random values but always a hardcoded key
+        fn arb_keyvalues_with_fixed_key() (
+            value in any::<[u8;32]>()
+        ) -> HashMap<(TrieDatabaseKeyType, [u8; 32]), [u8; 32]> {
+            let key = (TrieDatabaseKeyType::Trie, felt!("0x112345678921541231").to_bytes_be());
+            HashMap::from_iter([(key, value)])
+        }
+    }
+
+    prop_compose! {
+        fn arb_keyvalues() (
+            keyvalues in prop::collection::hash_map(
+                (arb_db_key_type(), any::<[u8;32]>()),
+                any::<[u8;32]>(),
+                1..100
+            )
+        ) -> HashMap<(TrieDatabaseKeyType, [u8; 32]), [u8; 32]> {
+            keyvalues
+        }
+    }
+
+    prop_compose! {
+        fn arb_block(count: u64, step: u64) (
+            number in (count * step)..((count * step) + step),
+            keyvalues in arb_keyvalues_with_fixed_key()
+        ) -> Case {
+            Case { number, keyvalues }
+        }
+    }
+
+    /// Strategy for generating a list of blocks with `count` size where each block is within a
+    /// range of `step` size. See [`arb_block`].
+    fn arb_blocklist(step: u64, count: usize) -> impl strategy::Strategy<Value = Vec<Case>> {
+        let mut strats = Vec::with_capacity(count);
+        for i in 0..count {
+            strats.push(arb_block(i as u64, step));
+        }
+        strategy::Strategy::prop_map(strats, move |strats| strats)
+    }
+
+    proptest! {
+        #[test]
+        fn test_get_insert(blocks in arb_blocklist(10, 1000)) {
+            let db = test_utils::create_test_db();
+            let tx = db.tx_mut().expect("failed to create rw tx");
+
+            for block in &blocks {
+                let mut trie = TrieDbMut::<tables::ClassesTrie, _>::new(&tx);
+
+                // Insert key/value pairs
+                for ((r#type, key), value) in &block.keyvalues {
+                    let db_key = match r#type {
+                        TrieDatabaseKeyType::Trie => DatabaseKey::Trie(key.as_ref()),
+                        TrieDatabaseKeyType::Flat => DatabaseKey::Flat(key.as_ref()),
+                        TrieDatabaseKeyType::TrieLog => DatabaseKey::TrieLog(key.as_ref()),
+                     };
+
+                    trie.insert(&db_key, value.as_ref(), None).expect("failed to insert");
+                }
+
+                let snapshot_id = CommitId::from(block.number);
+                trie.snapshot(snapshot_id);
+            }
+
+            tx.commit().expect("failed to commit tx");
+            let tx = db.tx().expect("failed to create ro tx");
+
+            for block in &blocks {
+                let snapshot_id = CommitId::from(block.number);
+                let snapshot_db = SnapshotTrieDb::<tables::ClassesTrie, _>::new(&tx, snapshot_id);
+
+                // Verify snapshots
+                for ((r#type, key), value) in &block.keyvalues {
+                    let db_key = match r#type {
+                        TrieDatabaseKeyType::Trie => DatabaseKey::Trie(key.as_ref()),
+                        TrieDatabaseKeyType::Flat => DatabaseKey::Flat(key.as_ref()),
+                        TrieDatabaseKeyType::TrieLog => DatabaseKey::TrieLog(key.as_ref()),
+                     };
+
+                    let result = snapshot_db.get(&db_key).unwrap();
+                    prop_assert_eq!(result.as_ref().map(|x| x.as_slice()), Some(value.as_slice()));
+                }
+            }
+        }
+    }
+}
diff --git a/crates/katana/trie/src/lib.rs b/crates/katana/trie/src/lib.rs
@@ -38,8 +38,14 @@ where
 {
     pub fn new(db: DB) -> Self {
         let config = BonsaiStorageConfig {
-            max_saved_trie_logs: Some(usize::MAX),
-            max_saved_snapshots: Some(usize::MAX),
+            // we have our own implementation of storing trie changes
+            max_saved_trie_logs: Some(0),
+            // 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()
+            max_saved_snapshots: Some(64usize),
             snapshot_interval: 1,
         };