fix(blockifier): get alias from last state diff

starkware-libs · Dec 25, 2024 · 5937694 · 5937694
1 parent 357c73e
commit 5937694
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 10 deletions.
diff --git a/crates/blockifier/src/state/stateful_compression.rs b/crates/blockifier/src/state/stateful_compression.rs
@@ -144,7 +144,7 @@ pub fn compress<S: StateReader>(
     state: &S,
     alias_contract_address: ContractAddress,
 ) -> CompressionResult<StateMaps> {
-    let alias_compressor = AliasCompressor { state, alias_contract_address };
+    let alias_compressor = AliasCompressor::new(state, alias_contract_address, state_diff);
 
     let nonces = state_diff
         .nonces
@@ -181,9 +181,28 @@ pub fn compress<S: StateReader>(
 struct AliasCompressor<'a, S: StateReader> {
     state: &'a S,
     alias_contract_address: ContractAddress,
+    new_aliases: HashMap<AliasKey, Alias>,
 }
 
-impl<S: StateReader> AliasCompressor<'_, S> {
+impl<'a, S: StateReader> AliasCompressor<'a, S> {
+    fn new(state: &'a S, alias_contract_address: ContractAddress, state_diff: &StateMaps) -> Self {
+        Self {
+            state,
+            alias_contract_address,
+            new_aliases: state_diff
+                .storage
+                .iter()
+                .filter_map(|((contract_address, key), value)| {
+                    if *contract_address == alias_contract_address {
+                        Some((*key, *value))
+                    } else {
+                        None
+                    }
+                })
+                .collect(),
+        }
+    }
+
     fn compress_address(
         &self,
         contract_address: &ContractAddress,
@@ -210,7 +229,16 @@ impl<S: StateReader> AliasCompressor<'_, S> {
     }
 
     fn get_alias(&self, alias_key: AliasKey) -> CompressionResult<Alias> {
-        let alias = self.state.get_storage_at(self.alias_contract_address, alias_key)?;
-        if alias == Felt::ZERO { Err(CompressionError::MissedAlias(alias_key)) } else { Ok(alias) }
+        match self.new_aliases.get(&alias_key) {
+            Some(alias) => Ok(*alias),
+            None => {
+                let alias = self.state.get_storage_at(self.alias_contract_address, alias_key)?;
+                if alias == Felt::ZERO {
+                    Err(CompressionError::MissedAlias(alias_key))
+                } else {
+                    Ok(alias)
+                }
+            }
+        }
     }
 }
diff --git a/crates/blockifier/src/state/stateful_compression_test.rs b/crates/blockifier/src/state/stateful_compression_test.rs
@@ -277,8 +277,8 @@ fn test_read_only_state(#[values(0, 2)] n_existing_aliases: u8) {
 }
 
 /// Tests the range of alias keys that should be compressed.
-#[test]
-fn test_alias_compressor() {
+#[rstest]
+fn test_alias_compressor(#[values(true, false)] new_aliases_allocated: bool) {
     let alias = Felt::from(500_u16);
 
     let high_key = 200_u16;
@@ -291,10 +291,18 @@ fn test_alias_compressor() {
 
     let no_compression_contract_address = ContractAddress::from(10_u16);
 
-    let mut state_reader = DictStateReader::default();
-    insert_to_alias_contract(&mut state_reader.storage_view, high_storage_key, alias);
-    let alias_compressor =
-        AliasCompressor { state: &state_reader, alias_contract_address: *ALIAS_CONTRACT_ADDRESS };
+    let (state_reader, new_aliases) = if new_aliases_allocated {
+        let mut state_reader = DictStateReader::default();
+        insert_to_alias_contract(&mut state_reader.storage_view, high_storage_key, alias);
+        (state_reader, HashMap::new())
+    } else {
+        (DictStateReader::default(), vec![(high_storage_key, alias)].into_iter().collect())
+    };
+    let alias_compressor = AliasCompressor {
+        state: &state_reader,
+        alias_contract_address: *ALIAS_CONTRACT_ADDRESS,
+        new_aliases,
+    };
 
     assert_eq!(
         alias_compressor.compress_address(&high_contract_address).unwrap(),