blockifier,apollo_batcher: produce OS initial reads behind os_input

Add an os_input feature on blockifier and, behind it, compute the pre-block OS read values (CachedState::get_os_initial_reads) in finalize_block, expose them on BlockExecutionSummary/BlockExecutionArtifacts, and carry them into the batcher's CentralObjects decision-reached response. Not yet consumed downstream.

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

Author: yoavGrs

.github/workflows/apollo_storage_os_input_ci.yml

@@ -66,6 +66,7 @@ jobs:
       - run: cargo test -p starknet_committer --features os_input
       - run: cargo test -p apollo_committer_types --features os_input
       - run: cargo test -p apollo_committer --features os_input
+      - run: cargo build -p blockifier --features os_input
       - run: cargo build -p apollo_batcher --features os_input
       - run: cargo test -p apollo_batcher --features os_input
       - run: cargo test -p apollo_consensus_orchestrator --features os_input,testing

crates/apollo_batcher/Cargo.toml

@@ -13,6 +13,7 @@ os_input = [
   "apollo_committer_types/os_input",
   "apollo_reverts/os_input",
   "apollo_storage/os_input",
+  "blockifier/os_input",
   "dep:starknet_committer",
 ]
 testing = []

crates/apollo_batcher/src/batcher.rs

@@ -1025,6 +1025,8 @@ impl Batcher {
                 parent_proposal_commitment,
                 #[cfg(feature = "os_input")]
                 accessed_keys,
+                #[cfg(feature = "os_input")]
+                initial_reads: block_execution_artifacts.initial_reads,
             },
         })
     }

crates/apollo_batcher/src/block_builder.rs

@@ -30,6 +30,8 @@ use blockifier::blockifier_versioned_constants::VersionedConstants;
 use blockifier::bouncer::{BouncerWeights, CasmHashComputationData};
 use blockifier::concurrency::worker_pool::WorkerPool;
 use blockifier::context::BlockContext;
+#[cfg(feature = "os_input")]
+use blockifier::state::cached_state::StateMaps;
 use blockifier::state::cached_state::{CachedState, CommitmentStateDiff};
 use blockifier::state::contract_class_manager::ContractClassManager;
 use blockifier::state::errors::StateError;
@@ -129,6 +131,8 @@ pub struct BlockExecutionArtifacts {
     pub execution_data: BlockTransactionExecutionData,
     pub commitment_state_diff: CommitmentStateDiff,
     pub compressed_state_diff: Option<CommitmentStateDiff>,
+    #[cfg(feature = "os_input")]
+    pub initial_reads: StateMaps,
     pub bouncer_weights: BouncerWeights,
     pub l2_gas_used: GasAmount,
     pub casm_hash_computation_data_sierra_gas: CasmHashComputationData,
@@ -142,18 +146,23 @@ pub struct BlockExecutionArtifacts {
 
 impl BlockExecutionArtifacts {
     pub async fn new(
-        BlockExecutionSummary {
+        block_summary: BlockExecutionSummary,
+        execution_data: BlockTransactionExecutionData,
+        final_n_executed_txs: usize,
+    ) -> Self {
+        #[cfg(feature = "os_input")]
+        let initial_reads = block_summary.initial_reads;
+        let BlockExecutionSummary {
             state_diff: commitment_state_diff,
             compressed_state_diff,
             bouncer_weights,
             casm_hash_computation_data_sierra_gas,
             casm_hash_computation_data_proving_gas,
             compiled_class_hashes_for_migration,
             block_info,
-        }: BlockExecutionSummary,
-        execution_data: BlockTransactionExecutionData,
-        final_n_executed_txs: usize,
-    ) -> Self {
+            // TODO(Yoav): Remove the ".." when the os_input feature is removed.
+            ..
+        } = block_summary;
         let l1_da_mode = L1DataAvailabilityMode::from_use_kzg_da(block_info.use_kzg_da);
         let transactions_data =
             prepare_txs_hashing_data(&execution_data.execution_infos_and_signatures);
@@ -173,6 +182,8 @@ impl BlockExecutionArtifacts {
             execution_data,
             commitment_state_diff,
             compressed_state_diff,
+            #[cfg(feature = "os_input")]
+            initial_reads,
             bouncer_weights,
             l2_gas_used,
             casm_hash_computation_data_sierra_gas,

crates/apollo_batcher/src/block_builder_test.rs

@@ -45,6 +45,9 @@ use starknet_api::transaction::fields::{
     VIRTUAL_SNOS,
 };
 use starknet_api::transaction::TransactionHash;
+#[cfg(feature = "os_input")]
+use starknet_api::{contract_address, felt, nonce, proof_facts, storage_key, tx_hash};
+#[cfg(not(feature = "os_input"))]
 use starknet_api::{proof_facts, tx_hash};
 use starknet_types_core::felt::Felt;
 use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
@@ -100,6 +103,8 @@ async fn block_execution_artifacts(
     let block_summary = BlockExecutionSummary {
         state_diff: Default::default(),
         compressed_state_diff: Default::default(),
+        #[cfg(feature = "os_input")]
+        initial_reads: test_initial_reads(),
         bouncer_weights: BouncerWeights { l1_gas: 100, ..BouncerWeights::empty() },
         casm_hash_computation_data_sierra_gas: CasmHashComputationData::default(),
         casm_hash_computation_data_proving_gas: CasmHashComputationData::default(),
@@ -135,6 +140,14 @@ fn execution_info() -> TransactionExecutionInfo {
     }
 }
 
+#[cfg(feature = "os_input")]
+fn test_initial_reads() -> StateMaps {
+    let mut initial_reads = StateMaps::default();
+    initial_reads.nonces.insert(contract_address!("0x1"), nonce!(7_u64));
+    initial_reads.storage.insert((contract_address!("0x1"), storage_key!("0x2")), felt!(8_u8));
+    initial_reads
+}
+
 async fn one_chunk_test_expectations() -> TestExpectations {
     let input_txs = test_txs(0..3);
     let block_size = input_txs.len();
@@ -445,6 +458,8 @@ async fn transaction_failed_test_expectations() -> TestExpectations {
         Ok(BlockExecutionSummary {
             state_diff: expected_block_artifacts_copy.commitment_state_diff,
             compressed_state_diff: None,
+            #[cfg(feature = "os_input")]
+            initial_reads: test_initial_reads(),
             bouncer_weights: expected_block_artifacts_copy.bouncer_weights,
             casm_hash_computation_data_sierra_gas: expected_block_artifacts_copy
                 .casm_hash_computation_data_sierra_gas,
@@ -543,6 +558,8 @@ async fn set_close_block_expectations(
         Ok(BlockExecutionSummary {
             state_diff: output_block_artifacts.commitment_state_diff,
             compressed_state_diff: None,
+            #[cfg(feature = "os_input")]
+            initial_reads: test_initial_reads(),
             bouncer_weights: output_block_artifacts.bouncer_weights,
             casm_hash_computation_data_sierra_gas: output_block_artifacts
                 .casm_hash_computation_data_sierra_gas,
@@ -1080,6 +1097,8 @@ async fn failed_l1_handler_transaction_consumed() {
         Ok(BlockExecutionSummary {
             state_diff: Default::default(),
             compressed_state_diff: None,
+            #[cfg(feature = "os_input")]
+            initial_reads: test_initial_reads(),
             bouncer_weights: BouncerWeights::empty(),
             casm_hash_computation_data_sierra_gas: CasmHashComputationData::default(),
             casm_hash_computation_data_proving_gas: CasmHashComputationData::default(),
@@ -1141,6 +1160,8 @@ async fn partial_chunk_execution_proposer() {
         Ok(BlockExecutionSummary {
             state_diff: expected_block_artifacts.commitment_state_diff,
             compressed_state_diff: None,
+            #[cfg(feature = "os_input")]
+            initial_reads: test_initial_reads(),
             bouncer_weights: expected_block_artifacts.bouncer_weights,
             casm_hash_computation_data_sierra_gas: expected_block_artifacts
                 .casm_hash_computation_data_sierra_gas,

crates/apollo_batcher/src/test_utils.rs

@@ -23,6 +23,8 @@ use blockifier::blockifier::transaction_executor::BlockExecutionSummary;
 use blockifier::bouncer::{BouncerWeights, CasmHashComputationData};
 use blockifier::fee::receipt::TransactionReceipt;
 use blockifier::state::cached_state::CommitmentStateDiff;
+#[cfg(feature = "os_input")]
+use blockifier::state::cached_state::StateMaps;
 use blockifier::transaction::objects::TransactionExecutionInfo;
 use indexmap::{indexmap, IndexMap};
 use mockall::predicate::eq;
@@ -213,6 +215,8 @@ impl BlockExecutionArtifacts {
                 address_to_nonce: IndexMap::from_iter([(contract_address!("0x7"), nonce!(1_u64))]),
             },
             compressed_state_diff: Default::default(),
+            #[cfg(feature = "os_input")]
+            initial_reads: StateMaps::default(),
             bouncer_weights: BouncerWeights::empty(),
             casm_hash_computation_data_sierra_gas: CasmHashComputationData::empty(),
             casm_hash_computation_data_proving_gas: CasmHashComputationData::empty(),

crates/apollo_batcher_types/src/batcher_types.rs

@@ -6,6 +6,8 @@ use blockifier::bouncer::{BouncerWeights, CasmHashComputationData};
 #[cfg(feature = "os_input")]
 use blockifier::state::accessed_keys::AccessedKeys;
 use blockifier::state::cached_state::CommitmentStateDiff;
+#[cfg(feature = "os_input")]
+use blockifier::state::cached_state::StateMaps;
 use blockifier::transaction::objects::TransactionExecutionInfo;
 use chrono::prelude::*;
 use indexmap::IndexMap;
@@ -162,6 +164,9 @@ pub struct CentralObjects {
     pub parent_proposal_commitment: Option<ProposalCommitment>,
     #[cfg(feature = "os_input")]
     pub accessed_keys: AccessedKeys,
+    /// Pre-block read values the OS needs to replay the block.
+    #[cfg(feature = "os_input")]
+    pub initial_reads: StateMaps,
 }
 
 #[derive(Debug, Serialize, Deserialize, PartialEq)]

crates/blockifier/Cargo.toml

@@ -21,6 +21,7 @@ mocks = []
 native_blockifier = []
 node_api = []
 only-native = ["cairo_native"]
+os_input = []
 reexecution = ["transaction_serde"]
 testing = [
   "blockifier_test_utils",

crates/blockifier/src/blockifier/transaction_executor.rs

@@ -56,6 +56,8 @@ pub type CompiledClassHashesForMigration = Vec<CompiledClassHashV2ToV1>;
 pub struct BlockExecutionSummary {
     pub state_diff: CommitmentStateDiff,
     pub compressed_state_diff: Option<CommitmentStateDiff>,
+    #[cfg(feature = "os_input")]
+    pub initial_reads: StateMaps,
     pub bouncer_weights: BouncerWeights,
     pub casm_hash_computation_data_sierra_gas: CasmHashComputationData,
     pub casm_hash_computation_data_proving_gas: CasmHashComputationData,
@@ -272,6 +274,9 @@ pub(crate) fn finalize_block<S: StateReader>(
 
     let state_diff = block_state.to_state_diff()?.state_maps;
 
+    #[cfg(feature = "os_input")]
+    let initial_reads = block_state.get_os_initial_reads()?;
+
     let compressed_state_diff = if block_context.versioned_constants.enable_stateful_compression {
         Some(compress(&state_diff, block_state, alias_contract_address)?.into())
     } else {
@@ -299,6 +304,8 @@ pub(crate) fn finalize_block<S: StateReader>(
     Ok(BlockExecutionSummary {
         state_diff: state_diff.into(),
         compressed_state_diff,
+        #[cfg(feature = "os_input")]
+        initial_reads,
         bouncer_weights: *bouncer.get_bouncer_weights(),
         casm_hash_computation_data_sierra_gas,
         casm_hash_computation_data_proving_gas,

crates/blockifier/src/state/cached_state.rs

@@ -288,13 +288,39 @@ impl Default for CachedState<crate::test_utils::dict_state_reader::DictStateRead
     }
 }
 
-#[cfg(feature = "reexecution")]
+#[cfg(any(feature = "reexecution", feature = "os_input"))]
 impl<S: StateReader> CachedState<S> {
     pub fn get_initial_reads(&self) -> StateResult<StateMaps> {
         Ok(self.cache.borrow().initial_reads.clone())
     }
 }
 
+#[cfg(feature = "os_input")]
+impl<S: StateReader> CachedState<S> {
+    /// Returns the pre-block values the OS needs to replay the block: the values read during
+    /// execution, extended with the class hash and nonce of every accessed contract and the
+    /// compiled class hash of every accessed class (the OS reads these trie leaves even when
+    /// execution itself did not). `declared_contracts` is cleared, as the OS does not consume it.
+    pub fn get_os_initial_reads(&mut self) -> StateResult<StateMaps> {
+        // Back-fill write-only storage cells so their pre-block values are part of the initial
+        // reads.
+        self.update_initial_values_of_write_only_access()?;
+        let raw_initial_reads = self.get_initial_reads()?;
+        // Force-read the contract-trie and class-trie leaves so their pre-block values are cached
+        // in the initial reads.
+        for contract_address in raw_initial_reads.get_contract_addresses() {
+            self.get_class_hash_at(contract_address)?;
+            self.get_nonce_at(contract_address)?;
+        }
+        for class_hash in raw_initial_reads.declared_contracts.keys() {
+            self.get_compiled_class_hash(*class_hash)?;
+        }
+        let mut os_initial_reads = self.get_initial_reads()?;
+        os_initial_reads.declared_contracts.clear();
+        Ok(os_initial_reads)
+    }
+}
+
 pub type StorageEntry = (ContractAddress, StorageKey);
 
 #[derive(Debug, Default, derive_more::IntoIterator)]