blockifier,apollo_batcher: trim OS initial reads to the accessed-key set

crates/apollo_batcher/src/batcher.rs

@@ -989,6 +989,15 @@ impl Batcher {
             &tx_execution_infos,
         )?;
 
+        // The OS only needs the read values for the keys it accesses; drop the extra reads (e.g.
+        // reverted-tx reads).
+        #[cfg(feature = "os_input")]
+        let initial_reads = {
+            let mut initial_reads = block_execution_artifacts.initial_reads;
+            initial_reads.trim_to_accessed_keys(&accessed_keys);
+            initial_reads
+        };
+
         self.write_commitment_results_and_add_new_task(
             height,
             state_diff.clone(), // TODO(Nimrod): Remove the clone here.
@@ -1026,7 +1035,7 @@ impl Batcher {
                 #[cfg(feature = "os_input")]
                 accessed_keys,
                 #[cfg(feature = "os_input")]
-                initial_reads: block_execution_artifacts.initial_reads,
+                initial_reads,
             },
         })
     }

crates/apollo_committer/src/communication.rs

@@ -15,6 +15,7 @@ pub type LocalCommitterServer =
     LocalComponentServer<ApolloCommitter, CommitterRequest, CommitterResponse>;
 pub type RemoteCommitterServer = RemoteComponentServer<CommitterRequest, CommitterResponse>;
 
+// `CommitterRequest` without variant `ReadPathsAndCommitBlock` for `os_input` feature.
 #[cfg(not(feature = "os_input"))]
 #[async_trait]
 impl<S: StorageConstructor, ForestDB: ForestStorageWithEmptyReadContext<Storage = S>>

crates/blockifier/src/state/cached_state.rs

@@ -9,6 +9,7 @@ use starknet_types_core::felt::Felt;
 
 use crate::context::TransactionContext;
 use crate::execution::contract_class::RunnableCompiledClass;
+use crate::state::accessed_keys::AccessedKeys;
 use crate::state::errors::StateError;
 use crate::state::state_api::{State, StateReader, StateResult, UpdatableState};
 use crate::transaction::objects::TransactionExecutionInfo;
@@ -399,6 +400,15 @@ impl StateMaps {
         self.declared_contracts.extend(&other.declared_contracts)
     }
 
+    /// Removes every entry whose key is not in `accessed_keys`.
+    pub fn trim_to_accessed_keys(&mut self, accessed_keys: &AccessedKeys) {
+        self.storage.retain(|key, _| accessed_keys.storage_keys.contains(key));
+        self.nonces.retain(|address, _| accessed_keys.accessed_contracts.contains(address));
+        self.class_hashes.retain(|address, _| accessed_keys.accessed_contracts.contains(address));
+        self.compiled_class_hashes
+            .retain(|class_hash, _| accessed_keys.accessed_class_hashes.contains(class_hash));
+    }
+
     /// Subtracts other's mappings from self.
     /// Assumes (and enforces) other's keys contains self's.
     pub fn diff(&self, other: &Self) -> Self {

crates/starknet_os_flow_tests/src/test_manager.rs

@@ -901,7 +901,7 @@ impl<S: FlowTestState> TestBuilder<S> {
                 &event_expectations_per_tx,
                 &execution_outputs,
             );
-            let initial_reads = final_state.get_os_initial_reads().unwrap();
+            let mut initial_reads = final_state.get_os_initial_reads().unwrap();
             let state_diff = final_state.to_state_diff().unwrap().state_maps;
             // Update the wrapped state.
             state = final_state.state;
@@ -930,6 +930,7 @@ impl<S: FlowTestState> TestBuilder<S> {
             );
             // Commit the state diff, building the OS-input commitment infos from the Patricia
             // witness paths gathered during the commit.
+            initial_reads.trim_to_accessed_keys(&accessed_keys);
             let committer_state_diff =
                 commitment_state_diff_to_committer_state_diff(commitment_state_diff);
             let (new_state_roots, commitment_infos) = commit_state_diff_with_witnesses(