diff --git a/Cargo.lock b/Cargo.lock index 4a066bb2fd..c34d601928 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2021,7 +2021,16 @@ version = "0.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0700ddab506f33b20a03b13996eccd309a48e5ff77d0d95926aa0210fb4e95f1" dependencies = [ - "bit-vec", + "bit-vec 0.6.3", +] + +[[package]] +name = "bit-set" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08807e080ed7f9d5433fa9b275196cfc35414f66a0c79d864dc51a0d825231a3" +dependencies = [ + "bit-vec 0.8.0", ] [[package]] @@ -2030,6 +2039,12 @@ version = "0.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "349f9b6a179ed607305526ca489b34ad0a41aed5f7980fa90eb03160b69598fb" +[[package]] +name = "bit-vec" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7" + [[package]] name = "bit_field" version = "0.10.2" @@ -8565,8 +8580,10 @@ dependencies = [ "page_size", "parking_lot 0.12.3", "postcard", + "proptest", "reth-libmdbx", "roaring", + "rstest 0.18.2", "serde", "serde_json", "smallvec", @@ -9023,7 +9040,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "55cb077ad656299f160924eb2912aa147d7339ea7d69e1b5517326fdcec3c1ca" dependencies = [ "ascii-canvas", - "bit-set", + "bit-set 0.5.3", "ena", "itertools 0.11.0", "lalrpop-util", @@ -11535,12 +11552,12 @@ dependencies = [ [[package]] name = "proptest" -version = "1.5.0" +version = "1.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b4c2511913b88df1637da85cc8d96ec8e43a3f8bb8ccb71ee1ac240d6f3df58d" +checksum = "14cae93065090804185d3b75f0bf93b8eeda30c7a9b4a33d3bdb3988d6229e50" dependencies = [ - "bit-set", - "bit-vec", + "bit-set 0.8.0", + "bit-vec 0.8.0", "bitflags 2.6.0", "lazy_static", "num-traits 0.2.19", diff --git a/crates/katana/storage/db/Cargo.toml b/crates/katana/storage/db/Cargo.toml index 8385b2dffe..47ff7f3464 100644 --- 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 index 3e3628919c..4f5cc49ec3 100644 --- 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 for TrieDatabaseKeyType { + type Error = TrieDatabaseKeyTypeTryFromError; + + fn try_from(value: u8) -> Result { + 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 index de5aaa12e2..c8f741c33f 100644 --- 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::::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::::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::(&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::::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::(&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::(&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 index b9adff8d72..a9012888f9 100644 --- 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 { + 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> { + 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::::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::::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 index 0753352c80..e0af40b4bc 100644 --- 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, };