added multiproof endpoint to partial tries

Author: michalpalkowski

Cargo.lock

@@ -252,9 +252,8 @@ where
             }
 
             info!(genesis_hash = %local_hash, "Genesis has already been initialized");
-        } else {
-            // Initialize the dev genesis block
-
+        } else if !is_forking {
+            // Initialize the dev genesis block (only for non-forked instances)
             let block = chain_spec.block();
             let states = chain_spec.state_updates();
 
@@ -273,6 +272,7 @@ where
 
             info!(genesis_hash = %outcome.block_hash, "Genesis initialized");
         }
+        // For forked instances, genesis is not created (fork already has its state)
 
         Ok(())
     }

crates/core/src/backend/mod.rs

@@ -422,7 +422,9 @@ where
         let provider = backend.storage.provider();
 
         let latest_num = provider.latest_number()?;
+        eprintln!("🔍 Creating executor for next block, latest_num: {}", latest_num);
         let updated_state = provider.latest()?;
+        eprintln!("✅ Got updated_state for block {}", latest_num + 1);
 
         let mut block_env = provider.block_env_at(latest_num.into())?.unwrap();
         backend.update_block_env(&mut block_env);

crates/core/src/service/block_producer.rs

@@ -422,9 +422,13 @@ where
         let state = self.state(&block_id)?;
 
         // Check that contract exist by checking the class hash of the contract,
-        // unless its address 0x1 which is special system contract and does not
-        // have a class. See https://docs.starknet.io/architecture-and-concepts/network-architecture/starknet-state/#address_0x1.
+        // unless its address 0x1 or 0x2 which are special system contracts and does not
+        // have a class.
+        // See:
+        //  https://docs.starknet.io/learn/protocol/state#address-0x1.
+        //  https://docs.starknet.io/learn/protocol/data-availability#v0-13-4
         if contract_address.0 != Felt::ONE
+            && contract_address.0 != Felt::TWO
             && contract_address.0 != Felt::TWO
             && state.class_hash_of_contract(contract_address)?.is_none()
         {

crates/rpc/rpc-server/src/starknet/mod.rs

@@ -666,14 +666,17 @@ mod tests {
     use katana_core::service::block_producer::IntervalBlockProducer;
     use katana_db::Db;
     use katana_node::config::fork::ForkingConfig;
+    use katana_primitives::block::{
+        Block, BlockNumber, FinalityStatus, Header, SealedBlockWithStatus,
+    };
     use katana_primitives::chain::ChainId;
     use katana_primitives::class::ClassHash;
-    use katana_primitives::state::StateUpdates;
+    use katana_primitives::state::{StateUpdates, StateUpdatesWithClasses};
     use katana_primitives::ContractAddress;
     use katana_primitives::Felt;
-    use katana_provider::api::block::BlockHashProvider;
-    use katana_provider::api::block::BlockNumberProvider;
     use katana_provider::api::block::BlockWriter;
+    use katana_provider::api::block::{BlockHashProvider, BlockNumberProvider};
+    use katana_provider::api::env::BlockEnvProvider;
     use katana_provider::api::trie::TrieWriter;
     use katana_provider::MutableProvider;
     use katana_provider::{ForkProviderFactory, ProviderFactory};
@@ -801,6 +804,20 @@ mod tests {
         (address, class_hash, storage, nonce)
     }
 
+    fn create_test_block_with_state_updates(
+        block_number: BlockNumber,
+        _state_updates: StateUpdates,
+    ) -> SealedBlockWithStatus {
+        SealedBlockWithStatus {
+            block: Block {
+                header: Header { number: block_number, ..Default::default() },
+                body: Vec::new(),
+            }
+            .seal(),
+            status: FinalityStatus::AcceptedOnL2,
+        }
+    }
+
     /// To run this test you need to comment out global cache part in Node::build() "let global_class_cache = class_cache.build_global()?";
     #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
     async fn test_commit_new_state_root_two_katana_instances() {
@@ -811,11 +828,6 @@ mod tests {
 
         // initialize state and mine at least one block before starting the fork
         let mut producer = IntervalBlockProducer::new(backend.clone(), None);
-        let block_number = provider.latest_number().unwrap();
-        let state_updates = setup_mainnet_updates_randomized(5);
-        let provider_mut = backend.storage.provider_mut();
-        provider_mut.compute_state_root(block_number, &state_updates).unwrap();
-        provider_mut.commit().unwrap();
         producer.force_mine();
 
         let fork_from_block = provider.latest_number().unwrap();
@@ -842,34 +854,71 @@ mod tests {
         let fork_block_number = fork_provider.latest_number().unwrap();
 
         let fork_minimal_updates = setup_mainnet_updates_randomized(5);
+
+        // Insert block with state updates on fork
         let fork_provider_mut = fork_backend.storage.provider_mut();
-        let state_root =
-            fork_provider_mut.compute_state_root(fork_block_number, &fork_minimal_updates).unwrap();
+        let new_fork_block_number = fork_block_number + 1;
+        let fork_state_root = fork_provider_mut
+            .compute_state_root(new_fork_block_number, &fork_minimal_updates)
+            .unwrap();
+
+        // Create and insert block with the state updates
+        let fork_block = create_test_block_with_state_updates(
+            new_fork_block_number,
+            fork_minimal_updates.clone(),
+        );
+        fork_provider_mut
+            .insert_block_with_states_and_receipts(
+                fork_block,
+                StateUpdatesWithClasses {
+                    state_updates: fork_minimal_updates.clone(),
+                    ..Default::default()
+                },
+                vec![],
+                vec![],
+            )
+            .unwrap();
         fork_provider_mut.commit().unwrap();
+
+        // Insert block with same state updates on mainnet
         let provider_mut = backend.storage.provider_mut();
-        let mainnet_state_root_same_updates =
-            provider_mut.compute_state_root(block_number, &fork_minimal_updates).unwrap();
+        let new_block_number = block_number + 1;
+        let mainnet_state_root =
+            provider_mut.compute_state_root(new_block_number, &fork_minimal_updates).unwrap();
+
+        let mainnet_block =
+            create_test_block_with_state_updates(new_block_number, fork_minimal_updates.clone());
+        provider_mut
+            .insert_block_with_states_and_receipts(
+                mainnet_block,
+                StateUpdatesWithClasses {
+                    state_updates: fork_minimal_updates.clone(),
+                    ..Default::default()
+                },
+                vec![],
+                vec![],
+            )
+            .unwrap();
         provider_mut.commit().unwrap();
 
-        producer.force_mine();
-        fork_producer.force_mine();
-
         assert_eq!(
-            state_root, mainnet_state_root_same_updates,
+            fork_state_root, mainnet_state_root,
             "State roots do not match on first run: fork={:?}, mainnet={:?}",
-            state_root, mainnet_state_root_same_updates
+            fork_state_root, mainnet_state_root
         );
 
         let block_number = provider.latest_number().unwrap();
         let fork_block_number = fork_provider.latest_number().unwrap();
         let state_updates = setup_mainnet_updates_randomized(5);
         let fork_provider_mut = fork_backend.storage.provider_mut();
+        let new_fork_block_number = fork_block_number + 1;
         let fork_state_root =
-            fork_provider_mut.compute_state_root(fork_block_number, &state_updates).unwrap();
+            fork_provider_mut.compute_state_root(new_fork_block_number, &state_updates).unwrap();
         fork_provider_mut.commit().unwrap();
         let provider_mut = backend.storage.provider_mut();
+        let new_block_number = block_number + 1;
         let mainnet_state_root =
-            provider_mut.compute_state_root(block_number, &state_updates).unwrap();
+            provider_mut.compute_state_root(new_block_number, &state_updates).unwrap();
         provider_mut.commit().unwrap();
 
         producer.force_mine();
@@ -886,12 +935,14 @@ mod tests {
 
         let state_updates = setup_mainnet_updates_randomized(5);
         let fork_provider_mut = fork_backend.storage.provider_mut();
+        let new_fork_block_number = fork_block_number + 1;
         let fork_state_root =
-            fork_provider_mut.compute_state_root(fork_block_number, &state_updates).unwrap();
+            fork_provider_mut.compute_state_root(new_fork_block_number, &state_updates).unwrap();
         fork_provider_mut.commit().unwrap();
         let provider_mut = backend.storage.provider_mut();
+        let new_block_number = block_number + 1;
         let mainnet_state_root =
-            provider_mut.compute_state_root(block_number, &state_updates).unwrap();
+            provider_mut.compute_state_root(new_block_number, &state_updates).unwrap();
         provider_mut.commit().unwrap();
 
         assert_eq!(

crates/rpc/rpc-server/tests/forking.rs

@@ -64,11 +64,25 @@ where
         katana_trie::ContractsTrie::new(TrieDb::new(self.tx.clone()))
     }
 
+    /// Returns the partial contracts trie (for forked instances that use insert_with_proof).
+    pub fn partial_contracts_trie(
+        &self,
+    ) -> katana_trie::PartialContractsTrie<TrieDb<'a, tables::ContractsTrie, Tx>> {
+        katana_trie::PartialContractsTrie::new_partial(TrieDb::new(self.tx.clone()))
+    }
+
     /// Returns the classes trie.
     pub fn classes_trie(&self) -> katana_trie::ClassesTrie<TrieDb<'a, tables::ClassesTrie, Tx>> {
         katana_trie::ClassesTrie::new(TrieDb::new(self.tx.clone()))
     }
 
+    /// Returns the partial classes trie (for forked instances that use insert_with_proof).
+    pub fn partial_classes_trie(
+        &self,
+    ) -> katana_trie::PartialClassesTrie<TrieDb<'a, tables::ClassesTrie, Tx>> {
+        katana_trie::PartialClassesTrie::new_partial(TrieDb::new(self.tx.clone()))
+    }
+
     // TODO: makes this return an Option
     /// Returns the storages trie.
     pub fn storages_trie(
@@ -77,6 +91,14 @@ where
     ) -> katana_trie::StoragesTrie<TrieDb<'a, tables::StoragesTrie, Tx>> {
         katana_trie::StoragesTrie::new(TrieDb::new(self.tx.clone()), address)
     }
+
+    /// Returns the partial storages trie (for forked instances that use insert_with_proof).
+    pub fn partial_storages_trie(
+        &self,
+        address: ContractAddress,
+    ) -> katana_trie::PartialStoragesTrie<TrieDb<'a, tables::StoragesTrie, Tx>> {
+        katana_trie::PartialStoragesTrie::new_partial(TrieDb::new(self.tx.clone()), address)
+    }
 }
 
 /// Historical tries, allowing access to the state tries at each block.
@@ -101,6 +123,14 @@ where
         katana_trie::ContractsTrie::new(SnapshotTrieDb::new(self.tx.clone(), commit))
     }
 
+    /// Returns the partial historical contracts trie (for forked instances that use insert_with_proof).
+    pub fn partial_contracts_trie(
+        &self,
+    ) -> katana_trie::PartialContractsTrie<SnapshotTrieDb<'a, tables::ContractsTrie, Tx>> {
+        let commit = CommitId::new(self.block);
+        katana_trie::PartialContractsTrie::new_partial(SnapshotTrieDb::new(self.tx.clone(), commit))
+    }
+
     /// Returns the historical classes trie.
     pub fn classes_trie(
         &self,
@@ -109,6 +139,14 @@ where
         katana_trie::ClassesTrie::new(SnapshotTrieDb::new(self.tx.clone(), commit))
     }
 
+    /// Returns the partial historical classes trie (for forked instances that use insert_with_proof).
+    pub fn partial_classes_trie(
+        &self,
+    ) -> katana_trie::PartialClassesTrie<SnapshotTrieDb<'a, tables::ClassesTrie, Tx>> {
+        let commit = CommitId::new(self.block);
+        katana_trie::PartialClassesTrie::new_partial(SnapshotTrieDb::new(self.tx.clone(), commit))
+    }
+
     // TODO: makes this return an Option
     /// Returns the historical storages trie.
     pub fn storages_trie(
@@ -118,6 +156,18 @@ where
         let commit = CommitId::new(self.block);
         katana_trie::StoragesTrie::new(SnapshotTrieDb::new(self.tx.clone(), commit), address)
     }
+
+    /// Returns the partial historical storages trie (for forked instances that use insert_with_proof).
+    pub fn partial_storages_trie(
+        &self,
+        address: ContractAddress,
+    ) -> katana_trie::PartialStoragesTrie<SnapshotTrieDb<'a, tables::StoragesTrie, Tx>> {
+        let commit = CommitId::new(self.block);
+        katana_trie::PartialStoragesTrie::new_partial(
+            SnapshotTrieDb::new(self.tx.clone(), commit),
+            address,
+        )
+    }
 }
 
 // --- Trie's database implementations. These are implemented based on the Bonsai Trie

crates/storage/db/src/trie/mod.rs

@@ -21,6 +21,7 @@ use katana_provider_api::block::{
 };
 use katana_provider_api::env::BlockEnvProvider;
 use katana_provider_api::stage::StageCheckpointProvider;
+use katana_provider_api::state::StateFactoryProvider;
 use katana_provider_api::state_update::StateUpdateProvider;
 use katana_provider_api::transaction::{
     ReceiptProvider, TransactionProvider, TransactionStatusProvider, TransactionTraceProvider,
@@ -207,22 +208,53 @@ impl<Tx1: DbTx> BlockNumberProvider for ForkedProvider<Tx1> {
     }
 
     fn latest_number(&self) -> ProviderResult<BlockNumber> {
-        match self.local_db.latest_number() {
-            Ok(num) => Ok(num),
-            // return the fork block number if local db return this error. this can only happen whne
-            // the ForkedProvider is constructed without inserting any locally produced
-            // blocks.
-            Err(ProviderError::MissingLatestBlockNumber) => Ok(self.block_id()),
-            Err(err) => Err(err),
-        }
+        let fork_point = self.block_id();
+        let local_latest = match self.local_db.latest_number() {
+            Ok(num) => num,
+            Err(ProviderError::MissingLatestBlockNumber) => fork_point,
+            Err(err) => return Err(err),
+        };
+        Ok(local_latest.max(fork_point))
     }
 }
 
 impl<Tx1: DbTx> BlockIdReader for ForkedProvider<Tx1> {}
 
 impl<Tx1: DbTx> BlockHashProvider for ForkedProvider<Tx1> {
     fn latest_hash(&self) -> ProviderResult<BlockHash> {
-        self.local_db.latest_hash()
+        // Use the same logic as latest_number() - if local_db has blocks, use local hash
+        let fork_point = self.block_id();
+        let local_latest = match self.local_db.latest_number() {
+            Ok(num) => num,
+            Err(ProviderError::MissingLatestBlockNumber) => fork_point,
+            Err(err) => return Err(err),
+        };
+
+        // If we have local blocks after fork point, use local hash
+        if local_latest > fork_point {
+            return self.local_db.latest_hash();
+        }
+
+        // Otherwise, use fork point hash (either local_latest == fork_point or local_db is empty)
+        if let Ok(hash) = self.local_db.latest_hash() {
+            Ok(hash)
+        } else {
+            // If local_db is empty, return the hash of the fork point block
+            if let Some(hash) = self.fork_db.db.provider().block_hash_by_num(fork_point)? {
+                Ok(hash)
+            } else {
+                // Fetch the fork point block if not cached
+                if self.fork_db.fetch_historical_blocks(fork_point.into())? {
+                    self.fork_db
+                        .db
+                        .provider()
+                        .block_hash_by_num(fork_point)?
+                        .ok_or(ProviderError::MissingLatestBlockHash)
+                } else {
+                    Err(ProviderError::MissingLatestBlockHash)
+                }
+            }
+        }
     }
 
     fn block_hash_by_num(&self, num: BlockNumber) -> ProviderResult<Option<BlockHash>> {
@@ -668,6 +700,28 @@ impl<Tx1: DbTxMut> BlockWriter for ForkedProvider<Tx1> {
         receipts: Vec<Receipt>,
         executions: Vec<TypedTransactionExecutionInfo>,
     ) -> ProviderResult<()> {
+        // BUGFIX: Before inserting state updates, ensure all contracts referenced in nonce_updates
+        // have their ContractInfo in local_db. For forked contracts, the class_hash may only exist
+        // in fork_cache (in-memory) or on the remote fork. We need to copy it to local_db first.
+        use katana_db::tables;
+        use katana_provider_api::state::StateProvider;
+
+        for addr in states.state_updates.nonce_updates.keys() {
+            // Check if ContractInfo exists in local_db
+            if self.local_db.tx().get::<tables::ContractInfo>(*addr)?.is_none() {
+                // Contract info not in local_db, check if it exists in fork cache or remote
+                // Create a state provider to search in fork
+                let state = self.latest()?;
+                if let Some(class_hash) = state.class_hash_of_contract(*addr)? {
+                    // Found in fork - copy to local_db before processing state updates
+                    let nonce = state.nonce(*addr)?.unwrap_or_default();
+                    let contract_info =
+                        katana_primitives::contract::GenericContractInfo { class_hash, nonce };
+                    self.local_db.tx().put::<tables::ContractInfo>(*addr, contract_info)?;
+                }
+            }
+        }
+
         self.local_db.insert_block_with_states_and_receipts(block, states, receipts, executions)
     }
 }