test(katana): ensure contract non-existent is validated properly on `…

…StateProvider::nonce()` of forked provider (#1947)

# Description

<!--
A description of what this PR is solving.
-->

## Related issue

<!--
Please link related issues: Fixes #<issue_number>
More info: https://docs.github.com/en/free-pro-team@latest/github/managing-your-work-on-github/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword
-->

## Tests

<!--
Please refer to the CONTRIBUTING.md file to know more about the testing process. Ensure you've tested at least the package you're modifying if running all the tests consumes too much memory on your system.
-->

- [ ] Yes
- [ ] No, because they aren't needed
- [ ] No, because I need help

## Added to documentation?

<!--
If the changes are small, code comments are enough, otherwise, the documentation is needed. It
may be a README.md file added to your module/package, a DojoBook PR or both.
-->

- [ ] README.md
- [ ] [Dojo Book](https://github.com/dojoengine/book)
- [ ] No documentation needed

## Checklist

- [ ] I've formatted my code (`scripts/prettier.sh`, `scripts/rust_fmt.sh`, `scripts/cairo_fmt.sh`)
- [ ] I've linted my code (`scripts/clippy.sh`, `scripts/docs.sh`)
- [ ] I've commented my code
- [ ] I've requested a review after addressing the comments

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

@@ -372,7 +372,7 @@ impl BackendHandle {
 /// cache to avoid fetching it again. This is shared across multiple instances of
 /// [`ForkedStateDb`](super::state::ForkedStateDb).
 #[derive(Clone)]
-pub struct SharedStateProvider(Arc<CacheStateDb<BackendHandle>>);
+pub struct SharedStateProvider(pub(crate) Arc<CacheStateDb<BackendHandle>>);
 
 impl SharedStateProvider {
     pub(crate) fn new_with_backend(backend: BackendHandle) -> Self {
@@ -599,38 +599,26 @@ fn handle_not_found_err<T>(result: Result<T, BackendError>) -> Result<Option<T>,
 }
 
 #[cfg(test)]
-mod tests {
+pub(crate) mod test_utils {
 
     use std::sync::mpsc::sync_channel;
-    use std::time::Duration;
 
     use katana_primitives::block::BlockNumber;
-    use katana_primitives::contract::GenericContractInfo;
-    use starknet::macros::felt;
     use starknet::providers::jsonrpc::HttpTransport;
     use starknet::providers::JsonRpcClient;
     use tokio::net::TcpListener;
     use url::Url;
 
     use super::*;
 
-    const LOCAL_RPC_URL: &str = "http://localhost:5050";
-
-    const STORAGE_KEY: StorageKey = felt!("0x1");
-    const ADDR_1: ContractAddress = ContractAddress(felt!("0xADD1"));
-    const ADDR_1_NONCE: Nonce = felt!("0x1");
-    const ADDR_1_STORAGE_VALUE: StorageKey = felt!("0x8080");
-    const ADDR_1_CLASS_HASH: StorageKey = felt!("0x1");
-
-    fn create_forked_backend(rpc_url: String, block_num: BlockNumber) -> BackendHandle {
-        let url = Url::parse(&rpc_url).expect("valid url");
+    pub fn create_forked_backend(rpc_url: &str, block_num: BlockNumber) -> BackendHandle {
+        let url = Url::parse(rpc_url).expect("valid url");
         let provider = Arc::new(JsonRpcClient::new(HttpTransport::new(url)));
-        let block_id = BlockHashOrNumber::Num(block_num);
-        Backend::new(provider, block_id).unwrap()
+        Backend::new(provider, block_num.into()).unwrap()
     }
 
     // Starts a TCP server that never close the connection.
-    fn start_tcp_server() {
+    pub fn start_tcp_server() {
         use tokio::runtime::Builder;
 
         let (tx, rx) = sync_channel::<()>(1);
@@ -650,22 +638,36 @@ mod tests {
 
         rx.recv().unwrap();
     }
+}
+
+#[cfg(test)]
+mod tests {
+
+    use std::time::Duration;
+
+    use katana_primitives::contract::GenericContractInfo;
+    use starknet::macros::felt;
+
+    use super::test_utils::*;
+    use super::*;
+
+    const LOCAL_RPC_URL: &str = "http://localhost:5050";
+
+    const STORAGE_KEY: StorageKey = felt!("0x1");
+    const ADDR_1: ContractAddress = ContractAddress(felt!("0xADD1"));
+    const ADDR_1_NONCE: Nonce = felt!("0x1");
+    const ADDR_1_STORAGE_VALUE: StorageKey = felt!("0x8080");
+    const ADDR_1_CLASS_HASH: StorageKey = felt!("0x1");
 
-    const ERROR_INIT_BACKEND: &str = "Failed to create backend";
     const ERROR_SEND_REQUEST: &str = "Failed to send request to backend";
     const ERROR_STATS: &str = "Failed to get stats";
 
     #[test]
     fn handle_incoming_requests() {
-        let url = Url::try_from("http://127.0.0.1:8080").unwrap();
-        let provider = JsonRpcClient::new(HttpTransport::new(url));
-        let block_id = BlockHashOrNumber::Num(1);
-
         // start a mock remote network
         start_tcp_server();
 
-        // start backend
-        let handle = Backend::new(Arc::new(provider), block_id).expect(ERROR_INIT_BACKEND);
+        let handle = create_forked_backend("http://127.0.0.1:8080", 1);
 
         // check no pending requests
         let stats = handle.stats().expect(ERROR_STATS);
@@ -704,7 +706,7 @@ mod tests {
     #[test]
     fn get_from_cache_if_exist() {
         // setup
-        let backend = create_forked_backend(LOCAL_RPC_URL.into(), 1);
+        let backend = create_forked_backend(LOCAL_RPC_URL, 1);
         let state_db = CacheStateDb::new(backend);
 
         state_db
@@ -734,7 +736,7 @@ mod tests {
 
     #[test]
     fn fetch_from_fork_will_err_if_backend_thread_not_running() {
-        let backend = create_forked_backend(LOCAL_RPC_URL.into(), 1);
+        let backend = create_forked_backend(LOCAL_RPC_URL, 1);
         let provider = SharedStateProvider(Arc::new(CacheStateDb::new(backend)));
         assert!(StateProvider::nonce(&provider, ADDR_1).is_err())
     }
@@ -751,7 +753,7 @@ mod tests {
     #[test]
     #[ignore]
     fn fetch_from_fork_if_not_in_cache() {
-        let backend = create_forked_backend(FORKED_URL.into(), 908622);
+        let backend = create_forked_backend(FORKED_URL, 908622);
         let provider = SharedStateProvider(Arc::new(CacheStateDb::new(backend)));
 
         // fetch from remote

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

@@ -39,9 +39,23 @@ impl StateProvider for ForkedStateDb {
         StateProvider::class_hash_of_contract(&self.db, address)
     }
 
+    // When reading from local storage, we only consider entries that have non-zero nonce
+    // values OR non-zero class hashes.
+    //
+    // Nonce == 0 && ClassHash == 0
+    // - Contract does not exist locally (so try find from remote state)
+    // Nonce != 0 && ClassHash == 0
+    // - Contract exists and was deployed remotely but new nonce was set locally (so no need to read
+    //   from remote state anymore)
+    // Nonce == 0 && ClassHash != 0
+    // - Contract exists and was deployed locally (always read from local state)
     fn nonce(&self, address: ContractAddress) -> ProviderResult<Option<Nonce>> {
-        if let nonce @ Some(_) =
-            self.contract_state.read().get(&address).map(|i| i.nonce).filter(|n| n != &Nonce::ZERO)
+        if let nonce @ Some(_) = self
+            .contract_state
+            .read()
+            .get(&address)
+            .filter(|c| c.nonce != Nonce::default() || c.class_hash != ClassHash::default())
+            .map(|c| c.nonce)
         {
             return Ok(nonce);
         }
@@ -134,8 +148,8 @@ impl StateProvider for ForkedSnapshot {
             .inner
             .contract_state
             .get(&address)
-            .map(|info| info.nonce)
-            .filter(|n| n != &Nonce::ZERO)
+            .filter(|c| c.nonce != Nonce::default() || c.class_hash != ClassHash::default())
+            .map(|c| c.nonce)
         {
             return Ok(nonce);
         }
@@ -199,3 +213,135 @@ impl ContractClassProvider for ForkedSnapshot {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashMap;
+
+    use katana_primitives::state::{StateUpdates, StateUpdatesWithDeclaredClasses};
+    use starknet::macros::felt;
+
+    use super::*;
+    use crate::providers::fork::backend::test_utils::create_forked_backend;
+
+    #[test]
+    fn test_get_nonce() {
+        let backend = create_forked_backend("http://localhost:8080", 1);
+
+        let address: ContractAddress = felt!("1").into();
+        let class_hash = felt!("11");
+        let remote_nonce = felt!("111");
+        let local_nonce = felt!("1111");
+
+        // Case: contract doesn't exist at all
+        {
+            let remote = SharedStateProvider::new_with_backend(backend.clone());
+            let local = ForkedStateDb::new(remote.clone());
+
+            // asserts that its error for now
+            assert!(local.nonce(address).is_err());
+            assert!(remote.nonce(address).is_err());
+
+            // make sure the snapshot maintains the same behavior
+            let snapshot = local.create_snapshot();
+            assert!(snapshot.nonce(address).is_err());
+        }
+
+        // Case: contract exist remotely
+        {
+            let remote = SharedStateProvider::new_with_backend(backend.clone());
+            let local = ForkedStateDb::new(remote.clone());
+
+            let nonce_updates = HashMap::from([(address, remote_nonce)]);
+            let updates = StateUpdatesWithDeclaredClasses {
+                state_updates: StateUpdates { nonce_updates, ..Default::default() },
+                ..Default::default()
+            };
+            remote.0.insert_updates(updates);
+
+            assert_eq!(local.nonce(address).unwrap(), Some(remote_nonce));
+            assert_eq!(remote.nonce(address).unwrap(), Some(remote_nonce));
+
+            // make sure the snapshot maintains the same behavior
+            let snapshot = local.create_snapshot();
+            assert_eq!(snapshot.nonce(address).unwrap(), Some(remote_nonce));
+        }
+
+        // Case: contract exist remotely but nonce was updated locally
+        {
+            let remote = SharedStateProvider::new_with_backend(backend.clone());
+            let local = ForkedStateDb::new(remote.clone());
+
+            let nonce_updates = HashMap::from([(address, remote_nonce)]);
+            let contract_updates = HashMap::from([(address, class_hash)]);
+            let updates = StateUpdatesWithDeclaredClasses {
+                state_updates: StateUpdates {
+                    nonce_updates,
+                    contract_updates,
+                    ..Default::default()
+                },
+                ..Default::default()
+            };
+            remote.0.insert_updates(updates);
+
+            let nonce_updates = HashMap::from([(address, local_nonce)]);
+            let updates = StateUpdatesWithDeclaredClasses {
+                state_updates: StateUpdates { nonce_updates, ..Default::default() },
+                ..Default::default()
+            };
+            local.insert_updates(updates);
+
+            assert_eq!(local.nonce(address).unwrap(), Some(local_nonce));
+            assert_eq!(remote.nonce(address).unwrap(), Some(remote_nonce));
+
+            // make sure the snapshot maintains the same behavior
+            let snapshot = local.create_snapshot();
+            assert_eq!(snapshot.nonce(address).unwrap(), Some(local_nonce));
+        }
+
+        // Case: contract was deployed locally only and has non-zero nonce
+        {
+            let remote = SharedStateProvider::new_with_backend(backend.clone());
+            let local = ForkedStateDb::new(remote.clone());
+
+            let contract_updates = HashMap::from([(address, class_hash)]);
+            let nonce_updates = HashMap::from([(address, local_nonce)]);
+            let updates = StateUpdatesWithDeclaredClasses {
+                state_updates: StateUpdates {
+                    nonce_updates,
+                    contract_updates,
+                    ..Default::default()
+                },
+                ..Default::default()
+            };
+            local.insert_updates(updates);
+
+            assert_eq!(local.nonce(address).unwrap(), Some(local_nonce));
+            assert!(remote.nonce(address).is_err());
+
+            // make sure the snapshot maintains the same behavior
+            let snapshot = local.create_snapshot();
+            assert_eq!(snapshot.nonce(address).unwrap(), Some(local_nonce));
+        }
+
+        // Case: contract was deployed locally only and has zero nonce
+        {
+            let remote = SharedStateProvider::new_with_backend(backend.clone());
+            let local = ForkedStateDb::new(remote.clone());
+
+            let contract_updates = HashMap::from([(address, class_hash)]);
+            let updates = StateUpdatesWithDeclaredClasses {
+                state_updates: StateUpdates { contract_updates, ..Default::default() },
+                ..Default::default()
+            };
+            local.insert_updates(updates);
+
+            assert_eq!(local.nonce(address).unwrap(), Some(Default::default()));
+            assert!(remote.nonce(address).is_err());
+
+            // make sure the snapshot maintains the same behavior
+            let snapshot = local.create_snapshot();
+            assert_eq!(snapshot.nonce(address).unwrap(), Some(Default::default()));
+        }
+    }
+}