diff --git a/crates/katana/executor/src/abstraction/mod.rs b/crates/katana/executor/src/abstraction/mod.rs index 78483e8ddd..d9a0e99af4 100644 --- a/crates/katana/executor/src/abstraction/mod.rs +++ b/crates/katana/executor/src/abstraction/mod.rs @@ -233,6 +233,10 @@ impl<'a> StateRootProvider for StateProviderDb<'a> { self.0.contracts_root() } + fn storage_root(&self, contract: ContractAddress) -> ProviderResult> { + self.0.storage_root(contract) + } + fn state_root(&self) -> ProviderResult { self.0.state_root() } diff --git a/crates/katana/executor/src/implementation/blockifier/state.rs b/crates/katana/executor/src/implementation/blockifier/state.rs index cbda7a187f..86aba56190 100644 --- a/crates/katana/executor/src/implementation/blockifier/state.rs +++ b/crates/katana/executor/src/implementation/blockifier/state.rs @@ -274,6 +274,10 @@ impl StateRootProvider for CachedState { fn contracts_root(&self) -> ProviderResult { unimplemented!("not supported in executor's state") } + + fn storage_root(&self, _: katana_primitives::ContractAddress) -> ProviderResult> { + unimplemented!("not supported in executor's state") + } } #[cfg(test)] diff --git a/crates/katana/executor/src/implementation/noop.rs b/crates/katana/executor/src/implementation/noop.rs index 09643823dc..0392e633eb 100644 --- a/crates/katana/executor/src/implementation/noop.rs +++ b/crates/katana/executor/src/implementation/noop.rs @@ -204,4 +204,8 @@ impl StateRootProvider for NoopStateProvider { fn contracts_root(&self) -> ProviderResult { Ok(Felt::ZERO) } + + fn storage_root(&self, _: ContractAddress) -> ProviderResult> { + Ok(Some(Felt::ZERO)) + } } diff --git a/crates/katana/rpc/rpc-types/src/trie.rs b/crates/katana/rpc/rpc-types/src/trie.rs index a92f0b4426..e7345fb899 100644 --- a/crates/katana/rpc/rpc-types/src/trie.rs +++ b/crates/katana/rpc/rpc-types/src/trie.rs @@ -92,6 +92,10 @@ pub struct ContractsProof { #[derive(Debug, Serialize, Deserialize)] pub struct ContractLeafData { + // NOTE: This field is not specified in the RPC specs, but the contract storage root is + // required to compute the contract state hash (ie the value of the contracts trie). We + // include this in the response for now to ease the conversions over on SNOS side. + pub storage_root: Felt, pub nonce: Felt, pub class_hash: Felt, } diff --git a/crates/katana/rpc/rpc/src/starknet/mod.rs b/crates/katana/rpc/rpc/src/starknet/mod.rs index 90bd2c6d3d..520b3d507c 100644 --- a/crates/katana/rpc/rpc/src/starknet/mod.rs +++ b/crates/katana/rpc/rpc/src/starknet/mod.rs @@ -1180,7 +1180,8 @@ where for address in addresses { let nonce = state.nonce(address)?.unwrap_or_default(); let class_hash = state.class_hash_of_contract(address)?.unwrap_or_default(); - contract_leaves_data.push(ContractLeafData { class_hash, nonce }); + let storage_root = state.storage_root(address)?.unwrap_or_default(); + contract_leaves_data.push(ContractLeafData { storage_root, class_hash, nonce }); } ContractsProof { nodes: proofs.into(), contract_leaves_data } diff --git a/crates/katana/storage/db/src/trie/mod.rs b/crates/katana/storage/db/src/trie/mod.rs index c8f741c33f..fce7968a24 100644 --- a/crates/katana/storage/db/src/trie/mod.rs +++ b/crates/katana/storage/db/src/trie/mod.rs @@ -5,6 +5,7 @@ use std::marker::PhantomData; use anyhow::Result; use katana_primitives::block::BlockNumber; +use katana_primitives::ContractAddress; use katana_trie::bonsai::{BonsaiDatabase, BonsaiPersistentDatabase, ByteVec, DatabaseKey}; use katana_trie::CommitId; use smallvec::ToSmallVec; @@ -68,9 +69,13 @@ where katana_trie::ClassesTrie::new(TrieDb::new(self.tx.clone())) } + // TODO: makes this return an Option /// Returns the storages trie. - pub fn storages_trie(&self) -> katana_trie::StoragesTrie> { - katana_trie::StoragesTrie::new(TrieDb::new(self.tx.clone())) + pub fn storages_trie( + &self, + address: ContractAddress, + ) -> katana_trie::StoragesTrie> { + katana_trie::StoragesTrie::new(TrieDb::new(self.tx.clone()), address) } } @@ -104,12 +109,14 @@ where katana_trie::ClassesTrie::new(SnapshotTrieDb::new(self.tx.clone(), commit)) } + // TODO: makes this return an Option /// Returns the historical storages trie. pub fn storages_trie( &self, + address: ContractAddress, ) -> katana_trie::StoragesTrie> { let commit = CommitId::new(self.block); - katana_trie::StoragesTrie::new(SnapshotTrieDb::new(self.tx.clone(), commit)) + katana_trie::StoragesTrie::new(SnapshotTrieDb::new(self.tx.clone(), commit), address) } } diff --git a/crates/katana/storage/provider/src/providers/db/state.rs b/crates/katana/storage/provider/src/providers/db/state.rs index fbf2498898..74ccb03ea1 100644 --- a/crates/katana/storage/provider/src/providers/db/state.rs +++ b/crates/katana/storage/provider/src/providers/db/state.rs @@ -11,6 +11,7 @@ use katana_primitives::class::{ClassHash, CompiledClass, CompiledClassHash, Cont use katana_primitives::contract::{ ContractAddress, GenericContractInfo, Nonce, StorageKey, StorageValue, }; +use katana_primitives::Felt; use super::DbProvider; use crate::error::ProviderError; @@ -185,8 +186,8 @@ where address: ContractAddress, storage_keys: Vec, ) -> ProviderResult { - let mut trie = TrieDbFactory::new(&self.0).latest().storages_trie(); - let proofs = trie.multiproof(address, storage_keys); + let mut trie = TrieDbFactory::new(&self.0).latest().storages_trie(address); + let proofs = trie.multiproof(storage_keys); Ok(proofs) } } @@ -195,15 +196,20 @@ impl StateRootProvider for LatestStateProvider where Tx: DbTx + fmt::Debug + Send + Sync, { - fn classes_root(&self) -> ProviderResult { + fn classes_root(&self) -> ProviderResult { let trie = TrieDbFactory::new(&self.0).latest().classes_trie(); Ok(trie.root()) } - fn contracts_root(&self) -> ProviderResult { + fn contracts_root(&self) -> ProviderResult { let trie = TrieDbFactory::new(&self.0).latest().contracts_trie(); Ok(trie.root()) } + + fn storage_root(&self, contract: ContractAddress) -> ProviderResult> { + let trie = TrieDbFactory::new(&self.0).latest().storages_trie(contract); + Ok(Some(trie.root())) + } } /// A historical state provider. @@ -372,8 +378,8 @@ where let proofs = TrieDbFactory::new(&self.tx) .historical(self.block_number) .expect("should exist") - .storages_trie() - .multiproof(address, storage_keys); + .storages_trie(address) + .multiproof(storage_keys); Ok(proofs) } } @@ -400,6 +406,15 @@ where Ok(root) } + fn storage_root(&self, contract: ContractAddress) -> ProviderResult> { + let root = TrieDbFactory::new(&self.tx) + .historical(self.block_number) + .expect("should exist") + .storages_trie(contract) + .root(); + Ok(Some(root)) + } + fn state_root(&self) -> ProviderResult { let header = self.tx.get::(self.block_number)?.expect("should exist"); Ok(header.state_root) diff --git a/crates/katana/storage/provider/src/providers/db/trie.rs b/crates/katana/storage/provider/src/providers/db/trie.rs index cb98362d4c..539835e764 100644 --- a/crates/katana/storage/provider/src/providers/db/trie.rs +++ b/crates/katana/storage/provider/src/providers/db/trie.rs @@ -10,6 +10,7 @@ use katana_primitives::state::StateUpdates; use katana_primitives::{ContractAddress, Felt}; use katana_trie::{compute_contract_state_hash, ClassesTrie, ContractsTrie, StoragesTrie}; +use crate::error::ProviderError; use crate::providers::db::DbProvider; use crate::traits::state::{StateFactoryProvider, StateProvider}; use crate::traits::trie::TrieWriter; @@ -48,24 +49,25 @@ impl TrieWriter for DbProvider { self.0.update(|tx| { let mut contract_trie_db = ContractsTrie::new(TrieDbMut::::new(tx)); - let mut storage_trie_db = - StoragesTrie::new(TrieDbMut::::new(tx)); let mut contract_leafs: HashMap = HashMap::new(); let leaf_hashes: Vec<_> = { // First we insert the contract storage changes for (address, storage_entries) in &state_updates.storage_updates { + let mut storage_trie_db = + StoragesTrie::new(TrieDbMut::::new(tx), *address); + for (key, value) in storage_entries { - storage_trie_db.insert(*address, *key, *value); + storage_trie_db.insert(*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); + // 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); @@ -82,15 +84,19 @@ impl TrieWriter for DbProvider { contract_leafs .into_iter() .map(|(address, mut leaf)| { - let storage_root = storage_trie_db.root(address); + let storage_trie = StoragesTrie::new( + TrieDbMut::::new(tx), + address, + ); + let storage_root = storage_trie.root(); leaf.storage_root = Some(storage_root); - let latest_state = self.latest().unwrap(); + let latest_state = self.latest()?; let leaf_hash = contract_state_leaf_hash(latest_state, &address, &leaf); - (address, leaf_hash) + Ok((address, leaf_hash)) }) - .collect::>() + .collect::, ProviderError>>()? }; for (k, v) in leaf_hashes { diff --git a/crates/katana/storage/provider/src/providers/fork/backend.rs b/crates/katana/storage/provider/src/providers/fork/backend.rs index 6e0c588869..b6113532b5 100644 --- a/crates/katana/storage/provider/src/providers/fork/backend.rs +++ b/crates/katana/storage/provider/src/providers/fork/backend.rs @@ -695,6 +695,10 @@ impl StateRootProvider for SharedStateProvider { fn contracts_root(&self) -> ProviderResult { unimplemented!("not supported in forked mode") } + + fn storage_root(&self, _: ContractAddress) -> ProviderResult> { + unimplemented!("not supported in forked mode") + } } /// A helper function to convert a contract/class not found error returned by the RPC provider into diff --git a/crates/katana/storage/provider/src/providers/fork/state.rs b/crates/katana/storage/provider/src/providers/fork/state.rs index 3c6fe5debd..7ab28ca869 100644 --- a/crates/katana/storage/provider/src/providers/fork/state.rs +++ b/crates/katana/storage/provider/src/providers/fork/state.rs @@ -131,6 +131,10 @@ impl StateRootProvider for ForkedStateDb { fn contracts_root(&self) -> ProviderResult { Ok(katana_primitives::Felt::ZERO) } + + fn storage_root(&self, _: ContractAddress) -> ProviderResult> { + Ok(Some(katana_primitives::Felt::ZERO)) + } } #[derive(Debug)] @@ -203,6 +207,10 @@ impl StateRootProvider for LatestStateProvider { fn contracts_root(&self) -> ProviderResult { Ok(katana_primitives::Felt::ZERO) } + + fn storage_root(&self, _: ContractAddress) -> ProviderResult> { + Ok(Some(katana_primitives::Felt::ZERO)) + } } impl StateProvider for ForkedSnapshot { @@ -306,6 +314,10 @@ impl StateRootProvider for ForkedSnapshot { fn contracts_root(&self) -> ProviderResult { Ok(katana_primitives::Felt::ZERO) } + + fn storage_root(&self, _: ContractAddress) -> ProviderResult> { + Ok(Some(katana_primitives::Felt::ZERO)) + } } #[cfg(test)] diff --git a/crates/katana/storage/provider/src/traits/state.rs b/crates/katana/storage/provider/src/traits/state.rs index b99e865d42..36e81d62d0 100644 --- a/crates/katana/storage/provider/src/traits/state.rs +++ b/crates/katana/storage/provider/src/traits/state.rs @@ -26,6 +26,9 @@ pub trait StateRootProvider: Send + Sync { /// Retrieves the root of the contracts trie. fn contracts_root(&self) -> ProviderResult; + + /// Retrieves the root of a contract's storage trie. + fn storage_root(&self, contract: ContractAddress) -> ProviderResult>; } #[auto_impl::auto_impl(&, Box, Arc)] diff --git a/crates/katana/trie/src/classes.rs b/crates/katana/trie/src/classes.rs index c74f3dffc1..e2481c7dcd 100644 --- a/crates/katana/trie/src/classes.rs +++ b/crates/katana/trie/src/classes.rs @@ -30,9 +30,9 @@ pub struct ClassesTrie { trie: crate::BonsaiTrie, } -///////////////////////////////////////////////////// +////////////////////////////////////////////////////////////// // ClassesTrie implementations -///////////////////////////////////////////////////// +////////////////////////////////////////////////////////////// impl ClassesTrie { const BONSAI_IDENTIFIER: &'static [u8] = b"classes"; diff --git a/crates/katana/trie/src/contracts.rs b/crates/katana/trie/src/contracts.rs index 7c6da60cc4..95e7394ee6 100644 --- a/crates/katana/trie/src/contracts.rs +++ b/crates/katana/trie/src/contracts.rs @@ -10,6 +10,10 @@ pub struct ContractsTrie { trie: crate::BonsaiTrie, } +////////////////////////////////////////////////////////////// +// ContractsTrie implementations +////////////////////////////////////////////////////////////// + impl ContractsTrie { /// NOTE: The identifier value is only relevant if the underlying [`BonsaiDatabase`] /// implementation is shared across other tries. @@ -25,7 +29,6 @@ impl ContractsTrie { pub fn multiproof(&mut self, addresses: Vec) -> MultiProof { let keys = addresses.into_iter().map(Felt::from).collect::>(); - dbg!(&keys); self.trie.multiproof(Self::BONSAI_IDENTIFIER, keys) } } diff --git a/crates/katana/trie/src/storages.rs b/crates/katana/trie/src/storages.rs index bb3c153cf1..c7f6bbe912 100644 --- a/crates/katana/trie/src/storages.rs +++ b/crates/katana/trie/src/storages.rs @@ -8,24 +8,22 @@ use crate::id::CommitId; #[derive(Debug)] pub struct StoragesTrie { + /// The contract address the storage trie belongs to. + address: ContractAddress, trie: crate::BonsaiTrie, } impl StoragesTrie { - pub fn new(db: DB) -> Self { - Self { trie: crate::BonsaiTrie::new(db) } + pub fn new(db: DB, address: ContractAddress) -> Self { + Self { address, trie: crate::BonsaiTrie::new(db) } } - pub fn root(&self, address: ContractAddress) -> Felt { - self.trie.root(&address.to_bytes_be()) + pub fn root(&self) -> Felt { + self.trie.root(&self.address.to_bytes_be()) } - pub fn multiproof( - &mut self, - address: ContractAddress, - storage_keys: Vec, - ) -> MultiProof { - self.trie.multiproof(&address.to_bytes_be(), storage_keys) + pub fn multiproof(&mut self, storage_keys: Vec) -> MultiProof { + self.trie.multiproof(&self.address.to_bytes_be(), storage_keys) } } @@ -33,13 +31,8 @@ impl StoragesTrie where DB: BonsaiDatabase + BonsaiPersistentDatabase, { - pub fn insert( - &mut self, - address: ContractAddress, - storage_key: StorageKey, - storage_value: StorageValue, - ) { - self.trie.insert(&address.to_bytes_be(), storage_key, storage_value) + pub fn insert(&mut self, storage_key: StorageKey, storage_value: StorageValue) { + self.trie.insert(&self.address.to_bytes_be(), storage_key, storage_value) } pub fn commit(&mut self, block: BlockNumber) {