diff --git a/crates/blockifier/src/concurrency/versioned_state.rs b/crates/blockifier/src/concurrency/versioned_state.rs index 9f1929a293..88e99eddc4 100644 --- a/crates/blockifier/src/concurrency/versioned_state.rs +++ b/crates/blockifier/src/concurrency/versioned_state.rs @@ -75,9 +75,9 @@ impl VersionedState { T: StateReader, { let writes = self.get_writes_up_to_index(from_index); - parent_state.update_cache(&writes); - parent_state.update_contract_class_cache( + parent_state.update_cache( + &writes, self.compiled_contract_classes.get_writes_up_to_index(from_index), ); } diff --git a/crates/blockifier/src/state/cached_state.rs b/crates/blockifier/src/state/cached_state.rs index 09cc649126..b96c689096 100644 --- a/crates/blockifier/src/state/cached_state.rs +++ b/crates/blockifier/src/state/cached_state.rs @@ -57,15 +57,17 @@ impl CachedState { Ok(self.to_state_diff()?.into()) } - pub fn update_cache(&mut self, write_updates: &StateMaps) { - let mut cache = self.cache.borrow_mut(); - cache.writes.extend(write_updates); - } - - pub fn update_contract_class_cache( + pub fn update_cache( &mut self, + write_updates: &StateMaps, local_contract_cache_updates: ContractClassMapping, ) { + // Check consistency between declared_contracts and class_hash_to_class. + for (&key, &value) in &write_updates.declared_contracts { + assert_eq!(value, local_contract_cache_updates.contains_key(&key)); + } + let mut cache = self.cache.borrow_mut(); + cache.writes.extend(write_updates); self.class_hash_to_class.get_mut().extend(local_contract_cache_updates); } @@ -110,9 +112,8 @@ impl UpdatableState for CachedState { class_hash_to_class: &ContractClassMapping, visited_pcs: &HashMap>, ) { - self.update_cache(writes); // TODO(OriF,15/5/24): Reconsider the clone. - self.update_contract_class_cache(class_hash_to_class.clone()); + self.update_cache(writes, class_hash_to_class.clone()); self.update_visited_pcs_cache(visited_pcs); } }