BREAKING: StateReader::get_storage_at return zero by default (#1011)

* update InMemoryStateReader & CachedState

* Add comments

* Apply changes to State implementation too

* Fix behaviour

* Add test

Author: fmoletta

src/state/cached_state.rs

@@ -105,28 +105,12 @@ impl<T: StateReader> StateReader for CachedState<T> {
     }
 
     /// Returns storage data for a given storage entry.
+    /// Returns zero as default value if missing
     fn get_storage_at(&self, storage_entry: &StorageEntry) -> Result<Felt252, StateError> {
-        if self.cache.get_storage(storage_entry).is_none() {
-            match self.state_reader.get_storage_at(storage_entry) {
-                Ok(storage) => {
-                    return Ok(storage);
-                }
-                Err(
-                    StateError::EmptyKeyInStorage
-                    | StateError::NoneStoragLeaf(_)
-                    | StateError::NoneStorage(_)
-                    | StateError::NoneContractState(_),
-                ) => return Ok(Felt252::zero()),
-                Err(e) => {
-                    return Err(e);
-                }
-            }
-        }
-
         self.cache
             .get_storage(storage_entry)
-            .ok_or_else(|| StateError::NoneStorage(storage_entry.clone()))
-            .cloned()
+            .map(|v| Ok(v.clone()))
+            .unwrap_or_else(|| self.state_reader.get_storage_at(storage_entry))
     }
 
     // TODO: check if that the proper way to store it (converting hash to address)
@@ -353,27 +337,20 @@ impl<T: StateReader> State for CachedState<T> {
             .clone())
     }
 
+    /// Returns storage data for a given storage entry.
+    /// Returns zero as default value if missing
+    /// Adds the value to the cache's inital_values if not present
     fn get_storage_at(&mut self, storage_entry: &StorageEntry) -> Result<Felt252, StateError> {
-        if self.cache.get_storage(storage_entry).is_none() {
-            let value = match self.state_reader.get_storage_at(storage_entry) {
-                Ok(value) => value,
-                Err(
-                    StateError::EmptyKeyInStorage
-                    | StateError::NoneStoragLeaf(_)
-                    | StateError::NoneStorage(_)
-                    | StateError::NoneContractState(_),
-                ) => Felt252::zero(),
-                Err(e) => return Err(e),
-            };
-            self.cache
-                .storage_initial_values
-                .insert(storage_entry.clone(), value);
+        match self.cache.get_storage(storage_entry) {
+            Some(value) => Ok(value.clone()),
+            None => {
+                let value = self.state_reader.get_storage_at(storage_entry)?;
+                self.cache
+                    .storage_initial_values
+                    .insert(storage_entry.clone(), value.clone());
+                Ok(value)
+            }
         }
-
-        self.cache
-            .get_storage(storage_entry)
-            .ok_or_else(|| StateError::NoneStorage(storage_entry.clone()))
-            .cloned()
     }
 
     // TODO: check if that the proper way to store it (converting hash to address)

src/state/in_memory_state_reader.rs

@@ -97,11 +97,11 @@ impl StateReader for InMemoryStateReader {
     }
 
     fn get_storage_at(&self, storage_entry: &StorageEntry) -> Result<Felt252, StateError> {
-        let storage = self
+        Ok(self
             .address_to_storage
             .get(storage_entry)
-            .ok_or_else(|| StateError::NoneStorage(storage_entry.clone()));
-        storage.cloned()
+            .cloned()
+            .unwrap_or_default())
     }
 
     fn get_compiled_class_hash(
@@ -132,10 +132,21 @@ impl StateReader for InMemoryStateReader {
 
 #[cfg(test)]
 mod tests {
+    use num_traits::One;
+
     use super::*;
     use crate::services::api::contract_classes::deprecated_contract_class::ContractClass;
     use std::sync::Arc;
 
+    #[test]
+    fn get_storage_returns_zero_if_missing() {
+        let state_reader = InMemoryStateReader::default();
+        assert!(state_reader
+            .get_storage_at(&(Address(Felt252::one()), Felt252::one().to_be_bytes()))
+            .unwrap()
+            .is_zero())
+    }
+
     #[test]
     fn get_contract_state_test() {
         let mut state_reader = InMemoryStateReader::new(

src/state/state_api.rs

@@ -16,6 +16,7 @@ pub trait StateReader {
     /// Returns the nonce of the given contract instance.
     fn get_nonce_at(&self, contract_address: &Address) -> Result<Felt252, StateError>;
     /// Returns the storage value under the given key in the given contract instance.
+    /// Returns zero by default if the value is not present
     fn get_storage_at(&self, storage_entry: &StorageEntry) -> Result<Felt252, StateError>;
     /// Return the class hash of the given casm contract class
     fn get_compiled_class_hash(
@@ -65,6 +66,8 @@ pub trait State {
     /// Default: 0 for an uninitialized contract address.
     fn get_nonce_at(&mut self, contract_address: &Address) -> Result<Felt252, StateError>;
 
+    /// Returns storage data for a given storage entry.
+    /// Returns zero as default value if missing
     fn get_storage_at(&mut self, storage_entry: &StorageEntry) -> Result<Felt252, StateError>;
 
     fn get_compiled_class_hash(&mut self, class_hash: &ClassHash) -> Result<ClassHash, StateError>;