feat(katana): return contract storage root in RPC proof response (#2846)

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<Option<Felt>> {
+        self.0.storage_root(contract)
+    }
+
     fn state_root(&self) -> ProviderResult<Felt> {
         self.0.state_root()
     }

crates/katana/executor/src/implementation/blockifier/state.rs

@@ -274,6 +274,10 @@ impl<S: StateDb> StateRootProvider for CachedState<S> {
     fn contracts_root(&self) -> ProviderResult<Felt> {
         unimplemented!("not supported in executor's state")
     }
+
+    fn storage_root(&self, _: katana_primitives::ContractAddress) -> ProviderResult<Option<Felt>> {
+        unimplemented!("not supported in executor's state")
+    }
 }
 
 #[cfg(test)]

crates/katana/executor/src/implementation/noop.rs

@@ -204,4 +204,8 @@ impl StateRootProvider for NoopStateProvider {
     fn contracts_root(&self) -> ProviderResult<Felt> {
         Ok(Felt::ZERO)
     }
+
+    fn storage_root(&self, _: ContractAddress) -> ProviderResult<Option<Felt>> {
+        Ok(Some(Felt::ZERO))
+    }
 }

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,
 }

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 }

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<TrieDb<'a, tables::StoragesTrie, Tx>> {
-        katana_trie::StoragesTrie::new(TrieDb::new(self.tx.clone()))
+    pub fn storages_trie(
+        &self,
+        address: ContractAddress,
+    ) -> katana_trie::StoragesTrie<TrieDb<'a, tables::StoragesTrie, Tx>> {
+        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<SnapshotTrieDb<'a, tables::StoragesTrie, Tx>> {
         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)
     }
 }
 

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<StorageKey>,
     ) -> ProviderResult<katana_trie::MultiProof> {
-        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<Tx> StateRootProvider for LatestStateProvider<Tx>
 where
     Tx: DbTx + fmt::Debug + Send + Sync,
 {
-    fn classes_root(&self) -> ProviderResult<katana_primitives::Felt> {
+    fn classes_root(&self) -> ProviderResult<Felt> {
         let trie = TrieDbFactory::new(&self.0).latest().classes_trie();
         Ok(trie.root())
     }
 
-    fn contracts_root(&self) -> ProviderResult<katana_primitives::Felt> {
+    fn contracts_root(&self) -> ProviderResult<Felt> {
         let trie = TrieDbFactory::new(&self.0).latest().contracts_trie();
         Ok(trie.root())
     }
+
+    fn storage_root(&self, contract: ContractAddress) -> ProviderResult<Option<Felt>> {
+        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<Option<Felt>> {
+        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<katana_primitives::Felt> {
         let header = self.tx.get::<tables::Headers>(self.block_number)?.expect("should exist");
         Ok(header.state_root)

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<Db: Database> TrieWriter for DbProvider<Db> {
         self.0.update(|tx| {
             let mut contract_trie_db =
                 ContractsTrie::new(TrieDbMut::<tables::ContractsTrie, _>::new(tx));
-            let mut storage_trie_db =
-                StoragesTrie::new(TrieDbMut::<tables::StoragesTrie, _>::new(tx));
 
             let mut contract_leafs: HashMap<ContractAddress, ContractLeaf> = 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::<tables::StoragesTrie, _>::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<Db: Database> TrieWriter for DbProvider<Db> {
                 contract_leafs
                     .into_iter()
                     .map(|(address, mut leaf)| {
-                        let storage_root = storage_trie_db.root(address);
+                        let storage_trie = StoragesTrie::new(
+                            TrieDbMut::<tables::StoragesTrie, _>::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::<Vec<_>>()
+                    .collect::<Result<Vec<_>, ProviderError>>()?
             };
 
             for (k, v) in leaf_hashes {

crates/katana/storage/provider/src/providers/fork/backend.rs

@@ -695,6 +695,10 @@ impl StateRootProvider for SharedStateProvider {
     fn contracts_root(&self) -> ProviderResult<Felt> {
         unimplemented!("not supported in forked mode")
     }
+
+    fn storage_root(&self, _: ContractAddress) -> ProviderResult<Option<Felt>> {
+        unimplemented!("not supported in forked mode")
+    }
 }
 
 /// A helper function to convert a contract/class not found error returned by the RPC provider into

crates/katana/storage/provider/src/providers/fork/state.rs

@@ -131,6 +131,10 @@ impl StateRootProvider for ForkedStateDb {
     fn contracts_root(&self) -> ProviderResult<katana_primitives::Felt> {
         Ok(katana_primitives::Felt::ZERO)
     }
+
+    fn storage_root(&self, _: ContractAddress) -> ProviderResult<Option<katana_primitives::Felt>> {
+        Ok(Some(katana_primitives::Felt::ZERO))
+    }
 }
 
 #[derive(Debug)]
@@ -203,6 +207,10 @@ impl StateRootProvider for LatestStateProvider {
     fn contracts_root(&self) -> ProviderResult<katana_primitives::Felt> {
         Ok(katana_primitives::Felt::ZERO)
     }
+
+    fn storage_root(&self, _: ContractAddress) -> ProviderResult<Option<katana_primitives::Felt>> {
+        Ok(Some(katana_primitives::Felt::ZERO))
+    }
 }
 
 impl StateProvider for ForkedSnapshot {
@@ -306,6 +314,10 @@ impl StateRootProvider for ForkedSnapshot {
     fn contracts_root(&self) -> ProviderResult<katana_primitives::Felt> {
         Ok(katana_primitives::Felt::ZERO)
     }
+
+    fn storage_root(&self, _: ContractAddress) -> ProviderResult<Option<katana_primitives::Felt>> {
+        Ok(Some(katana_primitives::Felt::ZERO))
+    }
 }
 
 #[cfg(test)]

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<Felt>;
+
+    /// Retrieves the root of a contract's storage trie.
+    fn storage_root(&self, contract: ContractAddress) -> ProviderResult<Option<Felt>>;
 }
 
 #[auto_impl::auto_impl(&, Box, Arc)]

crates/katana/trie/src/classes.rs

@@ -30,9 +30,9 @@ pub struct ClassesTrie<DB: BonsaiDatabase> {
     trie: crate::BonsaiTrie<DB, Pedersen>,
 }
 
-/////////////////////////////////////////////////////
+//////////////////////////////////////////////////////////////
 // 	ClassesTrie implementations
-/////////////////////////////////////////////////////
+//////////////////////////////////////////////////////////////
 
 impl<DB: BonsaiDatabase> ClassesTrie<DB> {
     const BONSAI_IDENTIFIER: &'static [u8] = b"classes";

crates/katana/trie/src/contracts.rs

@@ -10,6 +10,10 @@ pub struct ContractsTrie<DB: BonsaiDatabase> {
     trie: crate::BonsaiTrie<DB, Pedersen>,
 }
 
+//////////////////////////////////////////////////////////////
+// 	ContractsTrie implementations
+//////////////////////////////////////////////////////////////
+
 impl<DB: BonsaiDatabase> ContractsTrie<DB> {
     /// NOTE: The identifier value is only relevant if the underlying [`BonsaiDatabase`]
     /// implementation is shared across other tries.
@@ -25,7 +29,6 @@ impl<DB: BonsaiDatabase> ContractsTrie<DB> {
 
     pub fn multiproof(&mut self, addresses: Vec<ContractAddress>) -> MultiProof {
         let keys = addresses.into_iter().map(Felt::from).collect::<Vec<Felt>>();
-        dbg!(&keys);
         self.trie.multiproof(Self::BONSAI_IDENTIFIER, keys)
     }
 }

crates/katana/trie/src/storages.rs

@@ -8,38 +8,31 @@ use crate::id::CommitId;
 
 #[derive(Debug)]
 pub struct StoragesTrie<DB: BonsaiDatabase> {
+    /// The contract address the storage trie belongs to.
+    address: ContractAddress,
     trie: crate::BonsaiTrie<DB, Pedersen>,
 }
 
 impl<DB: BonsaiDatabase> StoragesTrie<DB> {
-    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<StorageKey>,
-    ) -> MultiProof {
-        self.trie.multiproof(&address.to_bytes_be(), storage_keys)
+    pub fn multiproof(&mut self, storage_keys: Vec<StorageKey>) -> MultiProof {
+        self.trie.multiproof(&self.address.to_bytes_be(), storage_keys)
     }
 }
 
 impl<DB> StoragesTrie<DB>
 where
     DB: BonsaiDatabase + BonsaiPersistentDatabase<CommitId>,
 {
-    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) {