fix(vm): Update decommitted_code_hashes in prepare_to_decommit (m…

…atter-labs#2253)

## What ❔

Update `decommitted_code_hashes` in `prepare_to_decommit`

## Why ❔

Contract hashes that reached `prepare_to_decommit` should be returned by
`get_used_contracts`

## Checklist

<!-- Check your PR fulfills the following items. -->
<!-- For draft PRs check the boxes as you complete them. -->

- [ ] PR title corresponds to the body of PR (we generate changelog
entries from PRs).
- [ ] Tests for the changes have been added / updated.
- [ ] Documentation comments have been added / updated.
- [ ] Code has been formatted via `zk fmt` and `zk lint`.
- [ ] Spellcheck has been run via `zk spellcheck`.

core/lib/multivm/src/versions/vm_latest/old_vm/oracles/decommitter.rs

@@ -30,7 +30,7 @@ pub struct DecommitterOracle<const B: bool, S, H: HistoryMode> {
     /// Stores pages of memory where certain code hashes have already been decommitted.
     /// It is expected that they all are present in the DB.
     // `decommitted_code_hashes` history is necessary
-    pub decommitted_code_hashes: HistoryRecorder<HashMap<U256, u32>, HistoryEnabled>,
+    pub decommitted_code_hashes: HistoryRecorder<HashMap<U256, Option<u32>>, HistoryEnabled>,
     /// Stores history of decommitment requests.
     decommitment_requests: HistoryRecorder<Vec<()>, H>,
 }
@@ -89,7 +89,7 @@ impl<S: ReadStorage, const B: bool, H: HistoryMode> DecommitterOracle<B, S, H> {
 
     pub fn get_decommitted_code_hashes_with_history(
         &self,
-    ) -> &HistoryRecorder<HashMap<U256, u32>, HistoryEnabled> {
+    ) -> &HistoryRecorder<HashMap<U256, Option<u32>>, HistoryEnabled> {
         &self.decommitted_code_hashes
     }
 
@@ -108,7 +108,7 @@ impl<S: ReadStorage, const B: bool, H: HistoryMode> DecommitterOracle<B, S, H> {
             .map(|(_, value)| value.len() * std::mem::size_of::<U256>())
             .sum::<usize>();
         let decommitted_code_hashes_size =
-            self.decommitted_code_hashes.inner().len() * std::mem::size_of::<(U256, u32)>();
+            self.decommitted_code_hashes.inner().len() * std::mem::size_of::<(U256, Option<u32>)>();
 
         known_bytecodes_size + decommitted_code_hashes_size
     }
@@ -132,7 +132,7 @@ impl<S: ReadStorage, const B: bool, H: HistoryMode> DecommitterOracle<B, S, H> {
         );
         let decommitted_code_hashes_size =
             self.decommitted_code_hashes.borrow_history(|h| h.len(), 0)
-                * std::mem::size_of::<<HashMap<U256, u32> as WithHistory>::HistoryRecord>();
+                * std::mem::size_of::<<HashMap<U256, Option<u32>> as WithHistory>::HistoryRecord>();
 
         known_bytecodes_stack_size + known_bytecodes_heap_size + decommitted_code_hashes_size
     }
@@ -172,13 +172,16 @@ impl<S: ReadStorage + Debug, const B: bool, H: HistoryMode> DecommittmentProcess
             .inner()
             .get(&stored_hash)
             .copied()
+            .flatten()
         {
             partial_query.is_fresh = false;
             partial_query.memory_page = MemoryPage(memory_page);
 
             Ok(partial_query)
         } else {
             partial_query.is_fresh = true;
+            self.decommitted_code_hashes
+                .insert(stored_hash, None, partial_query.timestamp);
 
             Ok(partial_query)
         }
@@ -216,7 +219,7 @@ impl<S: ReadStorage + Debug, const B: bool, H: HistoryMode> DecommittmentProcess
             rw_flag: true,
         };
         self.decommitted_code_hashes
-            .insert(stored_hash, page_to_use.0, timestamp);
+            .insert(stored_hash, Some(page_to_use.0), timestamp);
 
         // Copy the bytecode (that is stored in 'values' Vec) into the memory page.
         if B {

core/lib/multivm/src/versions/vm_latest/tests/get_used_contracts.rs

@@ -1,6 +1,14 @@
-use std::collections::{HashMap, HashSet};
+use std::{
+    collections::{HashMap, HashSet},
+    str::FromStr,
+};
 
 use itertools::Itertools;
+use zk_evm_1_5_0::{
+    abstractions::DecommittmentProcessor,
+    aux_structures::{DecommittmentQuery, MemoryPage, Timestamp},
+    zkevm_opcode_defs::{VersionedHashHeader, VersionedHashNormalizedPreimage},
+};
 use zksync_state::WriteStorage;
 use zksync_system_constants::CONTRACT_DEPLOYER_ADDRESS;
 use zksync_test_account::Account;
@@ -91,6 +99,47 @@ fn test_get_used_contracts() {
     }
 }
 
+#[test]
+fn test_contract_is_used_right_after_prepare_to_decommit() {
+    let mut vm = VmTesterBuilder::new(HistoryDisabled)
+        .with_empty_in_memory_storage()
+        .with_execution_mode(TxExecutionMode::VerifyExecute)
+        .build();
+
+    assert!(vm.vm.get_used_contracts().is_empty());
+
+    let bytecode_hash =
+        U256::from_str("0x100067ff3124f394104ab03481f7923f0bc4029a2aa9d41cc1d848c81257185")
+            .unwrap();
+    vm.vm
+        .state
+        .decommittment_processor
+        .populate(vec![(bytecode_hash, vec![])], Timestamp(0));
+
+    let header = hex::decode("0100067f").unwrap();
+    let normalized_preimage =
+        hex::decode("f3124f394104ab03481f7923f0bc4029a2aa9d41cc1d848c81257185").unwrap();
+    vm.vm
+        .state
+        .decommittment_processor
+        .prepare_to_decommit(
+            0,
+            DecommittmentQuery {
+                header: VersionedHashHeader(header.try_into().unwrap()),
+                normalized_preimage: VersionedHashNormalizedPreimage(
+                    normalized_preimage.try_into().unwrap(),
+                ),
+                timestamp: Timestamp(0),
+                memory_page: MemoryPage(0),
+                decommitted_length: 0,
+                is_fresh: false,
+            },
+        )
+        .unwrap();
+
+    assert_eq!(vm.vm.get_used_contracts(), vec![bytecode_hash]);
+}
+
 fn known_bytecodes_without_aa_code<S: WriteStorage, H: HistoryMode>(
     vm: &Vm<S, H>,
 ) -> HashMap<U256, Vec<U256>> {

core/lib/multivm/src/versions/vm_latest/tests/tester/inner_state.rs

@@ -43,7 +43,7 @@ pub(crate) struct DecommitterTestInnerState<H: HistoryMode> {
     /// so we just compare the modified keys. This is reasonable enough.
     pub(crate) modified_storage_keys: ModifiedKeysMap,
     pub(crate) known_bytecodes: HistoryRecorder<HashMap<U256, Vec<U256>>, H>,
-    pub(crate) decommitted_code_hashes: HistoryRecorder<HashMap<U256, u32>, HistoryEnabled>,
+    pub(crate) decommitted_code_hashes: HistoryRecorder<HashMap<U256, Option<u32>>, HistoryEnabled>,
 }
 
 #[derive(Clone, PartialEq, Debug)]

core/lib/multivm/src/versions/vm_latest/tracers/circuits_tracer.rs

@@ -177,20 +177,21 @@ impl<S: WriteStorage, H: HistoryMode> CircuitsTracer<S, H> {
             .decommitted_code_hashes
             .history();
         for (_, history_event) in &history[last_decommitment_history_entry_checked..] {
-            // We assume that only insertions may happen during a single VM inspection.
-            assert!(history_event.value.is_none());
-            let bytecode_len = state
-                .decommittment_processor
-                .known_bytecodes
-                .inner()
-                .get(&history_event.key)
-                .expect("Bytecode must be known at this point")
-                .len();
+            // We update cycles once per bytecode when it is actually decommitted.
+            if history_event.value.is_some() {
+                let bytecode_len = state
+                    .decommittment_processor
+                    .known_bytecodes
+                    .inner()
+                    .get(&history_event.key)
+                    .expect("Bytecode must be known at this point")
+                    .len();
 
-            // Each cycle of `CodeDecommitter` processes 2 words.
-            // If the number of words in bytecode is odd, then number of cycles must be rounded up.
-            let decommitter_cycles_used = (bytecode_len + 1) / 2;
-            self.statistics.code_decommitter_cycles += decommitter_cycles_used as u32;
+                // Each cycle of `CodeDecommitter` processes 2 words.
+                // If the number of words in bytecode is odd, then number of cycles must be rounded up.
+                let decommitter_cycles_used = (bytecode_len + 1) / 2;
+                self.statistics.code_decommitter_cycles += decommitter_cycles_used as u32;
+            }
         }
         self.last_decommitment_history_entry_checked = Some(history.len());
     }