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

Drop initial-reads entries whose keys are not in AccessedKeys (e.g. storage cells read only by a reverted transaction). The OS replays from the read values of the keys it accesses and needs nothing beyond that set, so the extra entries would only bloat the cende blob. Validated by the OS flow tests (113 passing), including all reverted-tx cases.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: yoavGrs

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/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;
@@ -400,6 +401,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

@@ -905,7 +905,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;
@@ -932,6 +932,7 @@ impl<S: FlowTestState> TestBuilder<S> {
                 &commitment_state_diff,
                 &accessed_keys_versioned_constants,
             );
+            initial_reads.trim_to_accessed_keys(&accessed_keys);
             // Commit the state diff.
             let committer_state_diff =
                 commitment_state_diff_to_committer_state_diff(commitment_state_diff);