diff --git a/Cargo.lock b/Cargo.lock index cc1f9676ec..5930a70abe 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1943,7 +1943,7 @@ dependencies = [ "bitflags 2.6.0", "cexpr", "clang-sys", - "itertools 0.10.5", + "itertools 0.12.1", "lazy_static", "lazycell", "log", @@ -2113,6 +2113,22 @@ dependencies = [ "zeroize", ] +[[package]] +name = "bonsai-trie" +version = "0.1.0" +source = "git+https://github.com/madara-alliance/bonsai-trie/?rev=56d7d62#56d7d62232fd72419f1d50de8bc747b70a9db68f" +dependencies = [ + "bitvec", + "derive_more 0.99.18", + "hashbrown 0.14.5", + "log", + "parity-scale-codec", + "rayon", + "serde", + "smallvec", + "starknet-types-core", +] + [[package]] name = "borsh" version = "1.5.1" @@ -6726,6 +6742,7 @@ checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" dependencies = [ "ahash 0.8.11", "allocator-api2", + "rayon", "serde", ] @@ -7208,7 +7225,7 @@ dependencies = [ "iana-time-zone-haiku", "js-sys", "wasm-bindgen", - "windows-core 0.51.1", + "windows-core 0.52.0", ] [[package]] @@ -8191,9 +8208,11 @@ version = "1.0.0-rc.0" dependencies = [ "anyhow", "arbitrary", + "bitvec", "criterion", "dojo-metrics", "katana-primitives", + "katana-trie", "metrics", "page_size", "parking_lot 0.12.3", @@ -8202,7 +8221,9 @@ dependencies = [ "roaring", "serde", "serde_json", + "smallvec", "starknet 0.12.0", + "starknet-types-core", "tempfile", "thiserror", "tracing", @@ -8346,10 +8367,12 @@ dependencies = [ "alloy-primitives", "anyhow", "auto_impl", + "bitvec", "futures", "katana-db", "katana-primitives", "katana-runner", + "katana-trie", "lazy_static", "parking_lot 0.12.3", "rand 0.8.5", @@ -8357,6 +8380,7 @@ dependencies = [ "rstest_reuse", "serde_json", "starknet 0.12.0", + "starknet-types-core", "tempfile", "thiserror", "tokio", @@ -8509,6 +8533,21 @@ dependencies = [ "tracing", ] +[[package]] +name = "katana-trie" +version = "1.0.0-rc.0" +dependencies = [ + "anyhow", + "bitvec", + "bonsai-trie", + "katana-primitives", + "serde", + "slab", + "starknet 0.12.0", + "starknet-types-core", + "thiserror", +] + [[package]] name = "keccak" version = "0.1.5" @@ -8599,9 +8638,9 @@ dependencies = [ [[package]] name = "lambdaworks-crypto" -version = "0.7.0" +version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7fb5d4f22241504f7c7b8d2c3a7d7835d7c07117f10bff2a7d96a9ef6ef217c3" +checksum = "bbc2a4da0d9e52ccfe6306801a112e81a8fc0c76aa3e4449fefeda7fef72bb34" dependencies = [ "lambdaworks-math", "serde", @@ -8611,9 +8650,9 @@ dependencies = [ [[package]] name = "lambdaworks-math" -version = "0.7.0" +version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "358e172628e713b80a530a59654154bfc45783a6ed70ea284839800cebdf8f97" +checksum = "d1bd2632acbd9957afc5aeec07ad39f078ae38656654043bf16e046fa2730e23" dependencies = [ "serde", "serde_json", @@ -11032,7 +11071,7 @@ checksum = "22505a5c94da8e3b7c2996394d1c933236c4d743e81a410bcca4e6989fc066a4" dependencies = [ "bytes", "heck 0.5.0", - "itertools 0.10.5", + "itertools 0.12.1", "log", "multimap", "once_cell", @@ -11073,7 +11112,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "81bddcdb20abf9501610992b6759a4c888aef7d1a7247ef75e2404275ac24af1" dependencies = [ "anyhow", - "itertools 0.10.5", + "itertools 0.12.1", "proc-macro2", "quote", "syn 2.0.77", @@ -14064,9 +14103,9 @@ dependencies = [ [[package]] name = "starknet-types-core" -version = "0.1.6" +version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9b889ee5734db8b3c8a6551135c16764bf4ce1ab4955fffbb2ac5b6706542b64" +checksum = "fa1b9e01ccb217ab6d475c5cda05dbb22c30029f7bb52b192a010a00d77a3d74" dependencies = [ "arbitrary", "lambdaworks-crypto", @@ -14075,6 +14114,7 @@ dependencies = [ "num-bigint", "num-integer", "num-traits 0.2.19", + "parity-scale-codec", "serde", ] @@ -16525,6 +16565,15 @@ dependencies = [ "windows-targets 0.48.5", ] +[[package]] +name = "windows-core" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33ab640c8d7e35bf8ba19b884ba838ceb4fba93a4e8c65a9059d08afcfc683d9" +dependencies = [ + "windows-targets 0.52.6", +] + [[package]] name = "windows-core" version = "0.57.0" diff --git a/Cargo.toml b/Cargo.toml index b5a19cedad..d9d0542ba0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -35,6 +35,7 @@ members = [ "crates/katana/storage/db", "crates/katana/storage/provider", "crates/katana/tasks", + "crates/katana/trie", "crates/metrics", "crates/saya/core", "crates/saya/provider", @@ -100,6 +101,7 @@ katana-rpc-types-builder = { path = "crates/katana/rpc/rpc-types-builder" } katana-runner = { path = "crates/katana/runner" } katana-slot-controller = { path = "crates/katana/controller" } katana-tasks = { path = "crates/katana/tasks" } +katana-trie = { path = "crates/katana/trie" } # torii torii-client = { path = "crates/torii/client" } @@ -254,4 +256,6 @@ alloy-transport = { version = "0.3", default-features = false } starknet = "0.12.0" starknet-crypto = "0.7.1" -starknet-types-core = { version = "0.1.6", features = [ "arbitrary" ] } +starknet-types-core = { version = "0.1.7", features = [ "arbitrary", "hash" ] } + +bitvec = "1.0.1" diff --git a/crates/katana/storage/db/Cargo.toml b/crates/katana/storage/db/Cargo.toml index 36b86bd266..b92f75b51a 100644 --- a/crates/katana/storage/db/Cargo.toml +++ b/crates/katana/storage/db/Cargo.toml @@ -8,6 +8,7 @@ version.workspace = true [dependencies] katana-primitives = { workspace = true, features = [ "arbitrary" ] } +katana-trie.workspace = true anyhow.workspace = true dojo-metrics.workspace = true @@ -17,12 +18,16 @@ parking_lot.workspace = true roaring = { version = "0.10.3", features = [ "serde" ] } serde.workspace = true serde_json.workspace = true +starknet.workspace = true +starknet-types-core.workspace = true tempfile.workspace = true thiserror.workspace = true tracing.workspace = true # codecs +bitvec.workspace = true postcard = { workspace = true, optional = true } +smallvec = "1.13.2" [dependencies.libmdbx] git = "https://github.com/paradigmxyz/reth.git" @@ -32,7 +37,6 @@ rev = "b34b0d3" [dev-dependencies] arbitrary.workspace = true criterion.workspace = true -starknet.workspace = true [features] default = [ "postcard" ] diff --git a/crates/katana/storage/db/src/codecs/postcard.rs b/crates/katana/storage/db/src/codecs/postcard.rs index 7acf83bead..2a154fb756 100644 --- a/crates/katana/storage/db/src/codecs/postcard.rs +++ b/crates/katana/storage/db/src/codecs/postcard.rs @@ -11,6 +11,7 @@ use crate::error::CodecError; use crate::models::block::StoredBlockBodyIndices; use crate::models::contract::ContractInfoChangeList; use crate::models::list::BlockList; +use crate::models::trie::TrieDatabaseValue; macro_rules! impl_compress_and_decompress_for_table_values { ($($name:ty),*) => { @@ -38,6 +39,7 @@ impl_compress_and_decompress_for_table_values!( Header, Receipt, Felt, + TrieDatabaseValue, ContractAddress, BlockList, GenericContractInfo, diff --git a/crates/katana/storage/db/src/lib.rs b/crates/katana/storage/db/src/lib.rs index 8f9179d40a..7b72db2669 100644 --- a/crates/katana/storage/db/src/lib.rs +++ b/crates/katana/storage/db/src/lib.rs @@ -13,6 +13,7 @@ pub mod error; pub mod mdbx; pub mod models; pub mod tables; +pub mod trie; pub mod utils; pub mod version; diff --git a/crates/katana/storage/db/src/models/mod.rs b/crates/katana/storage/db/src/models/mod.rs index 4279b48986..09fe7d1e0c 100644 --- a/crates/katana/storage/db/src/models/mod.rs +++ b/crates/katana/storage/db/src/models/mod.rs @@ -3,3 +3,4 @@ pub mod class; pub mod contract; pub mod list; pub mod storage; +pub mod trie; diff --git a/crates/katana/storage/db/src/models/trie.rs b/crates/katana/storage/db/src/models/trie.rs new file mode 100644 index 0000000000..b10456af57 --- /dev/null +++ b/crates/katana/storage/db/src/models/trie.rs @@ -0,0 +1,81 @@ +use katana_trie::bonsai::ByteVec; +use serde::{Deserialize, Serialize}; + +use crate::codecs::{Decode, Encode}; +use crate::error::CodecError; + +#[repr(u8)] +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)] +pub enum TrieDatabaseKeyType { + Trie = 0, + Flat, + TrieLog, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct TrieDatabaseKey { + pub r#type: TrieDatabaseKeyType, + pub key: Vec, +} + +pub type TrieDatabaseValue = ByteVec; + +impl Encode for TrieDatabaseKey { + type Encoded = Vec; + + fn encode(self) -> Self::Encoded { + let mut encoded = Vec::new(); + encoded.push(self.r#type as u8); + encoded.extend(self.key); + encoded + } +} + +impl Decode for TrieDatabaseKey { + fn decode>(bytes: B) -> Result { + let bytes = bytes.as_ref(); + if bytes.is_empty() { + panic!("emptyy buffer") + } + + let r#type = match bytes[0] { + 0 => TrieDatabaseKeyType::Trie, + 1 => TrieDatabaseKeyType::Flat, + 2 => TrieDatabaseKeyType::TrieLog, + _ => panic!("Invalid trie database key type"), + }; + + let key = bytes[1..].to_vec(); + + Ok(TrieDatabaseKey { r#type, key }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_trie_key_roundtrip() { + let key = TrieDatabaseKey { r#type: TrieDatabaseKeyType::Trie, key: vec![1, 2, 3] }; + let encoded = key.clone().encode(); + let decoded = TrieDatabaseKey::decode(encoded).unwrap(); + assert_eq!(key, decoded); + } + + #[test] + fn test_flat_key_roundtrip() { + let key = TrieDatabaseKey { r#type: TrieDatabaseKeyType::Flat, key: vec![4, 5, 6] }; + let encoded = key.clone().encode(); + let decoded = TrieDatabaseKey::decode(encoded).unwrap(); + assert_eq!(key, decoded); + } + + #[test] + fn test_trielog_key_roundtrip() { + let key = TrieDatabaseKey { r#type: TrieDatabaseKeyType::TrieLog, key: vec![7, 8, 9] }; + let encoded = key.clone().encode(); + let decoded = TrieDatabaseKey::decode(encoded).unwrap(); + assert_eq!(key, decoded); + } +} diff --git a/crates/katana/storage/db/src/tables.rs b/crates/katana/storage/db/src/tables.rs index ccb02bd147..3ebfc050d0 100644 --- a/crates/katana/storage/db/src/tables.rs +++ b/crates/katana/storage/db/src/tables.rs @@ -10,6 +10,7 @@ use crate::models::block::StoredBlockBodyIndices; use crate::models::contract::{ContractClassChange, ContractInfoChangeList, ContractNonceChange}; use crate::models::list::BlockList; use crate::models::storage::{ContractStorageEntry, ContractStorageKey, StorageEntry}; +use crate::models::trie::{TrieDatabaseKey, TrieDatabaseValue}; pub trait Key: Encode + Decode + Clone + std::fmt::Debug {} pub trait Value: Compress + Decompress + std::fmt::Debug {} @@ -35,6 +36,8 @@ pub trait DupSort: Table { type SubKey: Key; } +pub trait Trie: Table {} + /// Enum for the types of tables present in libmdbx. #[derive(Debug, PartialEq, Copy, Clone)] pub enum TableType { @@ -44,7 +47,7 @@ pub enum TableType { DupSort, } -pub const NUM_TABLES: usize = 23; +pub const NUM_TABLES: usize = 26; /// Macro to declare `libmdbx` tables. #[macro_export] @@ -167,7 +170,10 @@ define_tables_enum! {[ (NonceChangeHistory, TableType::DupSort), (ClassChangeHistory, TableType::DupSort), (StorageChangeHistory, TableType::DupSort), - (StorageChangeSet, TableType::Table) + (StorageChangeSet, TableType::Table), + (ClassTrie, TableType::Table), + (ContractTrie, TableType::Table), + (ContractStorageTrie, TableType::Table) ]} tables! { @@ -219,14 +225,23 @@ tables! { NonceChangeHistory: (BlockNumber, ContractAddress) => ContractNonceChange, /// Contract class hash changes by block. ClassChangeHistory: (BlockNumber, ContractAddress) => ContractClassChange, - /// storage change set StorageChangeSet: (ContractStorageKey) => BlockList, /// Account storage change set - StorageChangeHistory: (BlockNumber, ContractStorageKey) => ContractStorageEntry - + StorageChangeHistory: (BlockNumber, ContractStorageKey) => ContractStorageEntry, + + /// Class trie + ClassTrie: (TrieDatabaseKey) => TrieDatabaseValue, + /// Contract trie + ContractTrie: (TrieDatabaseKey) => TrieDatabaseValue, + /// Contract storage trie + ContractStorageTrie: (TrieDatabaseKey) => TrieDatabaseValue } +impl Trie for ClassTrie {} +impl Trie for ContractTrie {} +impl Trie for ContractStorageTrie {} + #[cfg(test)] mod tests { @@ -258,6 +273,8 @@ mod tests { assert_eq!(Tables::ALL[20].name(), ClassChangeHistory::NAME); assert_eq!(Tables::ALL[21].name(), StorageChangeHistory::NAME); assert_eq!(Tables::ALL[22].name(), StorageChangeSet::NAME); + assert_eq!(Tables::ALL[23].name(), ClassTrie::NAME); + assert_eq!(Tables::ALL[24].name(), ContractTrie::NAME); assert_eq!(Tables::Headers.table_type(), TableType::Table); assert_eq!(Tables::BlockHashes.table_type(), TableType::Table); @@ -282,6 +299,8 @@ mod tests { assert_eq!(Tables::ClassChangeHistory.table_type(), TableType::DupSort); assert_eq!(Tables::StorageChangeHistory.table_type(), TableType::DupSort); assert_eq!(Tables::StorageChangeSet.table_type(), TableType::Table); + assert_eq!(Tables::ClassTrie.table_type(), TableType::Table); + assert_eq!(Tables::ContractTrie.table_type(), TableType::Table); } use katana_primitives::address; diff --git a/crates/katana/storage/db/src/trie/class.rs b/crates/katana/storage/db/src/trie/class.rs new file mode 100644 index 0000000000..4e853829e9 --- /dev/null +++ b/crates/katana/storage/db/src/trie/class.rs @@ -0,0 +1,55 @@ +use bitvec::order::Msb0; +use bitvec::vec::BitVec; +use bitvec::view::AsBits; +use katana_primitives::block::BlockNumber; +use katana_primitives::class::{ClassHash, CompiledClassHash}; +use katana_primitives::Felt; +use katana_trie::bonsai::id::BasicId; +use katana_trie::bonsai::{BonsaiStorage, BonsaiStorageConfig}; +use starknet::macros::short_string; +use starknet_types_core::hash::{Poseidon, StarkHash}; + +use crate::abstraction::DbTxMut; +use crate::tables; +use crate::trie::TrieDb; + +// https://docs.starknet.io/architecture-and-concepts/network-architecture/starknet-state/#classes_trie +const CONTRACT_CLASS_LEAF_V0: Felt = short_string!("CONTRACT_CLASS_LEAF_V0"); + +#[derive(Debug)] +pub struct ClassTrie { + inner: BonsaiStorage, Poseidon>, +} + +impl ClassTrie { + pub fn new(tx: Tx) -> Self { + let config = BonsaiStorageConfig { + max_saved_trie_logs: Some(0), + max_saved_snapshots: Some(0), + snapshot_interval: u64::MAX, + }; + + let db = TrieDb::::new(tx); + let inner = BonsaiStorage::new(db, config).unwrap(); + + Self { inner } + } + + pub fn insert(&mut self, hash: ClassHash, compiled_hash: CompiledClassHash) { + let value = Poseidon::hash(&CONTRACT_CLASS_LEAF_V0, &compiled_hash); + let key: BitVec = hash.to_bytes_be().as_bits()[5..].to_owned(); + self.inner.insert(self.bonsai_identifier(), &key, &value).unwrap(); + } + + pub fn commit(&mut self, block_number: BlockNumber) { + self.inner.commit(BasicId::new(block_number)).unwrap(); + } + + pub fn root(&self) -> Felt { + self.inner.root_hash(self.bonsai_identifier()).unwrap() + } + + fn bonsai_identifier(&self) -> &'static [u8] { + b"1" + } +} diff --git a/crates/katana/storage/db/src/trie/contract.rs b/crates/katana/storage/db/src/trie/contract.rs new file mode 100644 index 0000000000..5e460739aa --- /dev/null +++ b/crates/katana/storage/db/src/trie/contract.rs @@ -0,0 +1,83 @@ +use bitvec::order::Msb0; +use bitvec::vec::BitVec; +use bitvec::view::AsBits; +use katana_primitives::block::BlockNumber; +use katana_primitives::contract::{StorageKey, StorageValue}; +use katana_primitives::{ContractAddress, Felt}; +use katana_trie::bonsai::id::BasicId; +use katana_trie::bonsai::{BonsaiStorage, BonsaiStorageConfig}; +use starknet_types_core::hash::Poseidon; + +use crate::abstraction::DbTxMut; +use crate::tables; +use crate::trie::TrieDb; + +#[derive(Debug)] +pub struct StorageTrie { + inner: BonsaiStorage, Poseidon>, +} + +impl StorageTrie { + pub fn new(tx: Tx) -> Self { + let config = BonsaiStorageConfig { + max_saved_trie_logs: Some(0), + max_saved_snapshots: Some(0), + snapshot_interval: u64::MAX, + }; + + let db = TrieDb::::new(tx); + let inner = BonsaiStorage::new(db, config).unwrap(); + + Self { inner } + } + + pub fn insert(&mut self, address: ContractAddress, key: StorageKey, value: StorageValue) { + let key: BitVec = key.to_bytes_be().as_bits()[5..].to_owned(); + self.inner.insert(&address.to_bytes_be(), &key, &value).unwrap(); + } + + pub fn commit(&mut self, block_number: BlockNumber) { + self.inner.commit(BasicId::new(block_number)).unwrap(); + } + + pub fn root(&self, address: &ContractAddress) -> Felt { + self.inner.root_hash(&address.to_bytes_be()).unwrap() + } +} + +#[derive(Debug)] +pub struct ContractTrie { + inner: BonsaiStorage, Poseidon>, +} + +impl ContractTrie { + pub fn new(tx: Tx) -> Self { + let config = BonsaiStorageConfig { + max_saved_trie_logs: Some(0), + max_saved_snapshots: Some(0), + snapshot_interval: u64::MAX, + }; + + let db = TrieDb::::new(tx); + let inner = BonsaiStorage::new(db, config).unwrap(); + + Self { inner } + } + + pub fn insert(&mut self, address: ContractAddress, state_hash: Felt) { + let key: BitVec = address.to_bytes_be().as_bits()[5..].to_owned(); + self.inner.insert(self.bonsai_identifier(), &key, &state_hash).unwrap(); + } + + pub fn commit(&mut self, block_number: BlockNumber) { + self.inner.commit(BasicId::new(block_number)).unwrap(); + } + + pub fn root(&self) -> Felt { + self.inner.root_hash(self.bonsai_identifier()).unwrap() + } + + fn bonsai_identifier(&self) -> &'static [u8] { + b"1" + } +} diff --git a/crates/katana/storage/db/src/trie/mod.rs b/crates/katana/storage/db/src/trie/mod.rs new file mode 100644 index 0000000000..823c9f28ba --- /dev/null +++ b/crates/katana/storage/db/src/trie/mod.rs @@ -0,0 +1,150 @@ +use std::marker::PhantomData; + +use anyhow::Result; +use katana_trie::bonsai::id::BasicId; +use katana_trie::bonsai::{self, ByteVec, DatabaseKey}; +use smallvec::ToSmallVec; + +use crate::abstraction::{DbCursor, DbTxMut}; +use crate::models::trie::{TrieDatabaseKey, TrieDatabaseKeyType}; +use crate::models::{self}; +use crate::tables; + +mod class; +mod contract; + +pub use class::ClassTrie; +pub use contract::{ContractTrie, StorageTrie}; + +#[derive(Debug, thiserror::Error)] +#[error(transparent)] +pub struct Error(#[from] crate::error::DatabaseError); + +impl katana_trie::bonsai::DBError for Error {} + +#[derive(Debug)] +pub struct TrieDb { + tx: Tx, + _table: PhantomData, +} + +impl TrieDb +where + Tb: tables::Trie, + Tx: DbTxMut, +{ + pub fn new(tx: Tx) -> Self { + Self { tx, _table: PhantomData } + } +} + +impl bonsai::BonsaiDatabase for TrieDb +where + Tb: tables::Trie, + Tx: DbTxMut, +{ + type Batch = (); + type DatabaseError = Error; + + fn create_batch(&self) -> Self::Batch {} + + fn remove_by_prefix(&mut self, prefix: &DatabaseKey<'_>) -> Result<(), Self::DatabaseError> { + let mut cursor = self.tx.cursor_mut::()?; + let walker = cursor.walk(None)?; + + let mut keys_to_remove = Vec::new(); + // iterate over all entries in the table + for entry in walker { + let (key, _) = entry?; + if key.key.starts_with(prefix.as_slice()) { + keys_to_remove.push(key); + } + } + + for key in keys_to_remove { + let _ = self.tx.delete::(key, None)?; + } + + Ok(()) + } + + fn get(&self, key: &DatabaseKey<'_>) -> Result, Self::DatabaseError> { + let value = self.tx.get::(to_db_key(key))?; + Ok(value) + } + + fn get_by_prefix( + &self, + prefix: &DatabaseKey<'_>, + ) -> Result, Self::DatabaseError> { + let _ = prefix; + todo!() + } + + fn insert( + &mut self, + key: &DatabaseKey<'_>, + value: &[u8], + _batch: Option<&mut Self::Batch>, + ) -> Result, Self::DatabaseError> { + let key = to_db_key(key); + let value: ByteVec = value.to_smallvec(); + let old_value = self.tx.get::(key.clone())?; + self.tx.put::(key, value)?; + Ok(old_value) + } + + fn remove( + &mut self, + key: &DatabaseKey<'_>, + _batch: Option<&mut Self::Batch>, + ) -> Result, Self::DatabaseError> { + let key = to_db_key(key); + let old_value = self.tx.get::(key.clone())?; + self.tx.delete::(key, None)?; + Ok(old_value) + } + + fn contains(&self, key: &DatabaseKey<'_>) -> Result { + let key = to_db_key(key); + let value = self.tx.get::(key)?; + Ok(value.is_some()) + } + + fn write_batch(&mut self, _batch: Self::Batch) -> Result<(), Self::DatabaseError> { + Ok(()) + } +} + +impl bonsai::BonsaiPersistentDatabase for TrieDb +where + Tb: tables::Trie, + Tx: DbTxMut, +{ + type DatabaseError = Error; + type Transaction = TrieDb; + + fn snapshot(&mut self, _: BasicId) {} + + fn merge(&mut self, _: Self::Transaction) -> Result<(), Self::DatabaseError> { + todo!(); + } + + fn transaction(&self, _: BasicId) -> Option { + todo!(); + } +} + +fn to_db_key(key: &DatabaseKey<'_>) -> models::trie::TrieDatabaseKey { + match key { + DatabaseKey::Flat(bytes) => { + TrieDatabaseKey { key: bytes.to_vec(), r#type: TrieDatabaseKeyType::Flat } + } + DatabaseKey::Trie(bytes) => { + TrieDatabaseKey { key: bytes.to_vec(), r#type: TrieDatabaseKeyType::Trie } + } + DatabaseKey::TrieLog(bytes) => { + TrieDatabaseKey { key: bytes.to_vec(), r#type: TrieDatabaseKeyType::TrieLog } + } + } +} diff --git a/crates/katana/storage/provider/Cargo.toml b/crates/katana/storage/provider/Cargo.toml index 4cf24e81b8..fa7860cee3 100644 --- a/crates/katana/storage/provider/Cargo.toml +++ b/crates/katana/storage/provider/Cargo.toml @@ -9,6 +9,7 @@ version.workspace = true [dependencies] katana-db = { workspace = true, features = [ "test-utils" ] } katana-primitives = { workspace = true, features = [ "rpc" ] } +katana-trie.workspace = true anyhow.workspace = true auto_impl.workspace = true @@ -16,6 +17,9 @@ parking_lot.workspace = true thiserror.workspace = true tracing.workspace = true +bitvec.workspace = true +starknet-types-core.workspace = true + # fork provider deps futures = { workspace = true, optional = true } starknet = { workspace = true, optional = true } diff --git a/crates/katana/storage/provider/src/lib.rs b/crates/katana/storage/provider/src/lib.rs index b7c1f7b430..cc37ce473e 100644 --- a/crates/katana/storage/provider/src/lib.rs +++ b/crates/katana/storage/provider/src/lib.rs @@ -1,3 +1,4 @@ +use std::collections::BTreeMap; use std::ops::{Range, RangeInclusive}; use katana_db::models::block::StoredBlockBodyIndices; @@ -18,6 +19,7 @@ use traits::contract::{ContractClassProvider, ContractClassWriter}; use traits::env::BlockEnvProvider; use traits::state::{StateRootProvider, StateWriter}; use traits::transaction::{TransactionStatusProvider, TransactionTraceProvider}; +use traits::trie::{ClassTrieWriter, ContractTrieWriter}; pub mod error; pub mod providers; @@ -380,3 +382,29 @@ where self.provider.block_env_at(id) } } + +impl ClassTrieWriter for BlockchainProvider +where + Db: ClassTrieWriter, +{ + fn insert_updates( + &self, + block_number: BlockNumber, + updates: &BTreeMap, + ) -> ProviderResult { + self.provider.insert_updates(block_number, updates) + } +} + +impl ContractTrieWriter for BlockchainProvider +where + Db: ContractTrieWriter, +{ + fn insert_updates( + &self, + block_number: BlockNumber, + state_updates: &StateUpdates, + ) -> ProviderResult { + self.provider.insert_updates(block_number, state_updates) + } +} diff --git a/crates/katana/storage/provider/src/providers/db/mod.rs b/crates/katana/storage/provider/src/providers/db/mod.rs index 92f6ea1a0f..4d06be8dda 100644 --- a/crates/katana/storage/provider/src/providers/db/mod.rs +++ b/crates/katana/storage/provider/src/providers/db/mod.rs @@ -1,4 +1,5 @@ pub mod state; +pub mod trie; use std::collections::BTreeMap; use std::fmt::Debug; diff --git a/crates/katana/storage/provider/src/providers/db/trie.rs b/crates/katana/storage/provider/src/providers/db/trie.rs new file mode 100644 index 0000000000..65b5b96984 --- /dev/null +++ b/crates/katana/storage/provider/src/providers/db/trie.rs @@ -0,0 +1,117 @@ +use std::collections::{BTreeMap, HashMap}; +use std::fmt::Debug; + +use katana_db::abstraction::Database; +use katana_db::trie; +use katana_db::trie::{ContractTrie, StorageTrie}; +use katana_primitives::block::BlockNumber; +use katana_primitives::class::{ClassHash, CompiledClassHash}; +use katana_primitives::state::StateUpdates; +use katana_primitives::{ContractAddress, Felt}; +use katana_trie::compute_contract_state_hash; + +use crate::providers::db::DbProvider; +use crate::traits::state::{StateFactoryProvider, StateProvider}; +use crate::traits::trie::{ClassTrieWriter, ContractTrieWriter}; + +#[derive(Debug, Default)] +struct ContractLeaf { + pub class_hash: Option, + pub storage_root: Option, + pub nonce: Option, +} + +impl ClassTrieWriter for DbProvider { + fn insert_updates( + &self, + block_number: BlockNumber, + updates: &BTreeMap, + ) -> crate::ProviderResult { + let mut trie = trie::ClassTrie::new(self.0.tx_mut()?); + + for (class_hash, compiled_hash) in updates { + trie.insert(*class_hash, *compiled_hash); + } + + trie.commit(block_number); + Ok(trie.root()) + } +} + +impl ContractTrieWriter for DbProvider { + fn insert_updates( + &self, + block_number: BlockNumber, + state_updates: &StateUpdates, + ) -> crate::ProviderResult { + let mut contract_leafs: HashMap = HashMap::new(); + + let leaf_hashes: Vec<_> = { + let mut storage_trie_db = StorageTrie::new(self.0.tx_mut()?); + + // First we insert the contract storage changes + for (address, storage_entries) in &state_updates.storage_updates { + for (key, value) in storage_entries { + storage_trie_db.insert(*address, *key, *value); + } + // insert the contract address in the contract_leafs to put the storage root later + contract_leafs.insert(*address, Default::default()); + } + + // Then we commit them + storage_trie_db.commit(block_number); + + for (address, nonce) in &state_updates.nonce_updates { + contract_leafs.entry(*address).or_default().nonce = Some(*nonce); + } + + for (address, class_hash) in &state_updates.deployed_contracts { + contract_leafs.entry(*address).or_default().class_hash = Some(*class_hash); + } + + for (address, class_hash) in &state_updates.replaced_classes { + contract_leafs.entry(*address).or_default().class_hash = Some(*class_hash); + } + + contract_leafs + .into_iter() + .map(|(address, mut leaf)| { + let storage_root = storage_trie_db.root(&address); + leaf.storage_root = Some(storage_root); + + let latest_state = self.latest().unwrap(); + let leaf_hash = contract_state_leaf_hash(latest_state, &address, &leaf); + + (address, leaf_hash) + }) + .collect::>() + }; + + let mut contract_trie_db = ContractTrie::new(self.0.tx_mut()?); + + for (k, v) in leaf_hashes { + contract_trie_db.insert(k, v); + } + + contract_trie_db.commit(block_number); + Ok(contract_trie_db.root()) + } +} + +// computes the contract state leaf hash +fn contract_state_leaf_hash( + provider: impl StateProvider, + address: &ContractAddress, + contract_leaf: &ContractLeaf, +) -> Felt { + let nonce = + contract_leaf.nonce.unwrap_or(provider.nonce(*address).unwrap().unwrap_or_default()); + + let class_hash = contract_leaf + .class_hash + .unwrap_or(provider.class_hash_of_contract(*address).unwrap().unwrap_or_default()); + + let storage_root = contract_leaf.storage_root.expect("root need to set"); + + compute_contract_state_hash(&class_hash, &storage_root, &nonce) +} diff --git a/crates/katana/storage/provider/src/providers/fork/mod.rs b/crates/katana/storage/provider/src/providers/fork/mod.rs index 5383d3c960..115c6dd45a 100644 --- a/crates/katana/storage/provider/src/providers/fork/mod.rs +++ b/crates/katana/storage/provider/src/providers/fork/mod.rs @@ -1,6 +1,7 @@ pub mod backend; pub mod state; +use std::collections::BTreeMap; use std::ops::{Range, RangeInclusive}; use std::sync::Arc; @@ -16,6 +17,7 @@ use katana_primitives::receipt::Receipt; use katana_primitives::state::{StateUpdates, StateUpdatesWithDeclaredClasses}; use katana_primitives::trace::TxExecInfo; use katana_primitives::transaction::{Tx, TxHash, TxNumber, TxWithHash}; +use katana_primitives::Felt; use parking_lot::RwLock; use starknet::providers::jsonrpc::HttpTransport; use starknet::providers::JsonRpcClient; @@ -36,6 +38,7 @@ use crate::traits::transaction::{ ReceiptProvider, TransactionProvider, TransactionStatusProvider, TransactionTraceProvider, TransactionsProviderExt, }; +use crate::traits::trie::{ClassTrieWriter, ContractTrieWriter}; use crate::ProviderResult; #[derive(Debug)] @@ -582,3 +585,27 @@ impl BlockEnvProvider for ForkedProvider { })) } } + +impl ClassTrieWriter for ForkedProvider { + fn insert_updates( + &self, + block_number: BlockNumber, + updates: &BTreeMap, + ) -> ProviderResult { + let _ = block_number; + let _ = updates; + Ok(Felt::ZERO) + } +} + +impl ContractTrieWriter for ForkedProvider { + fn insert_updates( + &self, + block_number: BlockNumber, + state_updates: &StateUpdates, + ) -> ProviderResult { + let _ = block_number; + let _ = state_updates; + Ok(Felt::ZERO) + } +} diff --git a/crates/katana/storage/provider/src/traits/mod.rs b/crates/katana/storage/provider/src/traits/mod.rs index 762de20465..1fdbcf3de2 100644 --- a/crates/katana/storage/provider/src/traits/mod.rs +++ b/crates/katana/storage/provider/src/traits/mod.rs @@ -4,3 +4,4 @@ pub mod env; pub mod state; pub mod state_update; pub mod transaction; +pub mod trie; diff --git a/crates/katana/storage/provider/src/traits/trie.rs b/crates/katana/storage/provider/src/traits/trie.rs new file mode 100644 index 0000000000..8570a30b88 --- /dev/null +++ b/crates/katana/storage/provider/src/traits/trie.rs @@ -0,0 +1,26 @@ +use std::collections::BTreeMap; + +use katana_primitives::block::BlockNumber; +use katana_primitives::class::{ClassHash, CompiledClassHash}; +use katana_primitives::state::StateUpdates; +use katana_primitives::Felt; + +use crate::ProviderResult; + +#[auto_impl::auto_impl(&, Box, Arc)] +pub trait ClassTrieWriter: Send + Sync { + fn insert_updates( + &self, + block_number: BlockNumber, + updates: &BTreeMap, + ) -> ProviderResult; +} + +#[auto_impl::auto_impl(&, Box, Arc)] +pub trait ContractTrieWriter: Send + Sync { + fn insert_updates( + &self, + block_number: BlockNumber, + state_updates: &StateUpdates, + ) -> ProviderResult; +} diff --git a/crates/katana/trie/Cargo.toml b/crates/katana/trie/Cargo.toml new file mode 100644 index 0000000000..759dab5347 --- /dev/null +++ b/crates/katana/trie/Cargo.toml @@ -0,0 +1,23 @@ +[package] +edition.workspace = true +license.workspace = true +name = "katana-trie" +repository.workspace = true +version.workspace = true + +[dependencies] +katana-primitives.workspace = true + +anyhow.workspace = true +bitvec.workspace = true +serde.workspace = true +slab = "0.4.9" +starknet.workspace = true +starknet-types-core.workspace = true +thiserror.workspace = true + +[dependencies.bonsai-trie] +default-features = false +features = [ "std" ] +git = "https://github.com/madara-alliance/bonsai-trie/" +rev = "56d7d62" diff --git a/crates/katana/trie/src/lib.rs b/crates/katana/trie/src/lib.rs new file mode 100644 index 0000000000..1a3e858bd6 --- /dev/null +++ b/crates/katana/trie/src/lib.rs @@ -0,0 +1,88 @@ +use anyhow::Result; +use bitvec::vec::BitVec; +pub use bonsai_trie as bonsai; +use bonsai_trie::id::BasicId; +use bonsai_trie::{BonsaiDatabase, BonsaiPersistentDatabase}; +use katana_primitives::class::ClassHash; +use katana_primitives::Felt; +use starknet_types_core::hash::{Pedersen, StarkHash}; + +/// A helper trait to define a database that can be used as a Bonsai Trie. +/// +/// Basically a short hand for `BonsaiDatabase + BonsaiPersistentDatabase`. +pub trait BonsaiTrieDb: BonsaiDatabase + BonsaiPersistentDatabase {} +impl BonsaiTrieDb for T where T: BonsaiDatabase + BonsaiPersistentDatabase {} + +pub fn compute_merkle_root(values: &[Felt]) -> Result +where + H: StarkHash + Send + Sync, +{ + use bonsai_trie::id::BasicId; + use bonsai_trie::{databases, BonsaiStorage, BonsaiStorageConfig}; + + // the value is irrelevant + const IDENTIFIER: &[u8] = b"1"; + + let config = BonsaiStorageConfig::default(); + let bonsai_db = databases::HashMapDb::::default(); + let mut bs = BonsaiStorage::<_, _, H>::new(bonsai_db, config).unwrap(); + + for (id, value) in values.iter().enumerate() { + let key = BitVec::from_iter(id.to_be_bytes()); + bs.insert(IDENTIFIER, key.as_bitslice(), value).unwrap(); + } + + let id = bonsai_trie::id::BasicIdBuilder::new().new_id(); + bs.commit(id).unwrap(); + + Ok(bs.root_hash(IDENTIFIER).unwrap()) +} + +// H(H(H(class_hash, storage_root), nonce), 0), where H is the pedersen hash +pub fn compute_contract_state_hash( + class_hash: &ClassHash, + storage_root: &Felt, + nonce: &Felt, +) -> Felt { + const CONTRACT_STATE_HASH_VERSION: Felt = Felt::ZERO; + let hash = Pedersen::hash(class_hash, storage_root); + let hash = Pedersen::hash(&hash, nonce); + Pedersen::hash(&hash, &CONTRACT_STATE_HASH_VERSION) +} + +#[cfg(test)] +mod tests { + + use katana_primitives::contract::Nonce; + use katana_primitives::felt; + use starknet_types_core::hash; + + use super::*; + + // Taken from Pathfinder: https://github.com/eqlabs/pathfinder/blob/29f93d0d6ad8758fdcf5ae3a8bd2faad2a3bc92b/crates/merkle-tree/src/transaction.rs#L70-L88 + #[test] + fn test_commitment_merkle_tree() { + let hashes = vec![Felt::from(1), Felt::from(2), Felt::from(3), Felt::from(4)]; + + // Produced by the cairo-lang Python implementation: + // `hex(asyncio.run(calculate_patricia_root([1, 2, 3, 4], height=64, ffc=ffc))))` + let expected_root_hash = + felt!("0x1a0e579b6b444769e4626331230b5ae39bd880f47e703b73fa56bf77e52e461"); + let computed_root_hash = compute_merkle_root::(&hashes).unwrap(); + + assert_eq!(expected_root_hash, computed_root_hash); + } + + // Taken from Pathfinder: https://github.com/eqlabs/pathfinder/blob/29f93d0d6ad8758fdcf5ae3a8bd2faad2a3bc92b/crates/merkle-tree/src/contract_state.rs#L236C5-L252C6 + #[test] + fn test_compute_contract_state_hash() { + let root = felt!("0x4fb440e8ca9b74fc12a22ebffe0bc0658206337897226117b985434c239c028"); + let class_hash = felt!("0x2ff4903e17f87b298ded00c44bfeb22874c5f73be2ced8f1d9d9556fb509779"); + let nonce = Nonce::ZERO; + + let result = compute_contract_state_hash(&class_hash, &root, &nonce); + let expected = felt!("0x7161b591c893836263a64f2a7e0d829c92f6956148a60ce5e99a3f55c7973f3"); + + assert_eq!(result, expected); + } +}