starknet_os_runner: build os input

Author: AvivYossef-starkware

crates/starknet_os_runner/Cargo.toml

@@ -9,14 +9,14 @@ description = "Runs transactions through the Starknet OS and returns Cairo PIE a
 [dependencies]
 blockifier.workspace = true
 blockifier_reexecution.workspace = true
+cairo-lang-starknet-classes.workspace = true
 indexmap.workspace = true
 starknet-rust.workspace = true
 starknet-rust-core.workspace = true
 starknet-types-core.workspace = true
 starknet_api.workspace = true
 starknet_os.workspace = true
 starknet_patricia.workspace = true
-cairo-lang-starknet-classes.workspace = true
 thiserror.workspace = true
 tokio.workspace = true
 url.workspace = true

crates/starknet_os_runner/src/classes_provider.rs

@@ -1,4 +1,4 @@
-use std::collections::{BTreeMap, HashMap, HashSet};
+use std::collections::{BTreeMap, HashSet};
 
 use cairo_lang_starknet_classes::casm_contract_class::CasmContractClass;
 use starknet_api::core::{ClassHash, CompiledClassHash};
@@ -21,6 +21,6 @@ pub trait ClassesProvider {
     /// Fetches all classes required for the OS run based on the executed class hashes.
     fn get_classes(
         &self,
-        executed_class_hashes: HashSet<ClassHash>,
+        executed_class_hashes: &HashSet<ClassHash>,
     ) -> Result<ClassesInput, ClassesProviderError>;
 }

crates/starknet_os_runner/src/errors.rs

@@ -13,9 +13,18 @@ pub enum VirtualBlockExecutorError {
 
     #[error("Block state unavailable after execution")]
     StateUnavailable,
+}
 
-    #[error("Unsupported transaction type: only Invoke transactions are supported")]
-    UnsupportedTransactionType,
+#[derive(Debug, Error)]
+pub enum RunnerError {
+    #[error(transparent)]
+    ClassesProvider(#[from] ClassesProviderError),
+    #[error(transparent)]
+    ProofProvider(#[from] ProofProviderError),
+    #[error(transparent)]
+    VirtualBlockExecutor(#[from] VirtualBlockExecutorError),
+    #[error("OS Input generation failed: {0}")]
+    InputGenerationError(String),
 }
 
 #[derive(Debug, Error)]

crates/starknet_os_runner/src/lib.rs

@@ -1,5 +1,6 @@
 pub mod classes_provider;
 pub mod errors;
+pub mod runner;
 pub mod storage_proofs;
 pub mod virtual_block_executor;
 
@@ -9,3 +10,4 @@ mod storage_proofs_test;
 pub mod test_utils;
 #[cfg(test)]
 mod virtual_block_executor_test;
+

crates/starknet_os_runner/src/runner.rs

@@ -0,0 +1,108 @@
+use std::collections::HashMap;
+
+use blockifier::state::contract_class_manager::ContractClassManager;
+use starknet_api::block::{BlockHash, BlockNumber};
+use starknet_api::transaction::{InvokeTransaction, TransactionHash};
+use starknet_os::io::os_input::{OsBlockInput, StarknetOsInput};
+
+use crate::classes_provider::ClassesProvider;
+use crate::errors::RunnerError;
+use crate::storage_proofs::StorageProofProvider;
+use crate::virtual_block_executor::VirtualBlockExecutor;
+
+pub struct Runner<C, S, V>
+where
+    C: ClassesProvider,
+    S: StorageProofProvider,
+    V: VirtualBlockExecutor,
+{
+    pub classes_provider: C,
+    pub storage_proofs_provider: S,
+    pub virtual_block_executor: V,
+}
+
+impl<C, S, V> Runner<C, S, V>
+where
+    C: ClassesProvider,
+    S: StorageProofProvider,
+    V: VirtualBlockExecutor,
+{
+    pub fn new(classes_provider: C, storage_proofs_provider: S, virtual_block_executor: V) -> Self {
+        Self { classes_provider, storage_proofs_provider, virtual_block_executor }
+    }
+
+    /// Creates OS input from execution data and storage proofs.
+    pub fn create_os_input(
+        &self,
+        block_number: BlockNumber,
+        contract_class_manager: ContractClassManager,
+        txs: Vec<(InvokeTransaction, TransactionHash)>,
+    ) -> Result<StarknetOsInput, RunnerError> {
+        // 1. Execute virtual block and get execution data.
+        let mut execution_data = self.virtual_block_executor.execute(
+            block_number,
+            contract_class_manager,
+            txs.clone(),
+        )?;
+
+        // 2. Fetch classes (consuming executed_class_hashes).
+        let classes = self.classes_provider.get_classes(&execution_data.executed_class_hashes)?;
+
+        // 3. Fetch storage proofs (pass execution_data by reference to avoid moving yet).
+        let storage_proofs =
+            self.storage_proofs_provider.get_storage_proofs(block_number, &execution_data)?;
+
+        // 4. Convert execution outputs to CentralTransactionExecutionInfo (consuming).
+        let tx_execution_infos =
+            execution_data.execution_outputs.into_iter().map(|output| output.0.into()).collect();
+
+        // 5. Merge initial_reads with proof_state.
+        execution_data.initial_reads.extend(&storage_proofs.proof_state);
+
+        // 6. Assemble OsBlockInput.
+        let os_block_input = OsBlockInput {
+            contract_state_commitment_info: storage_proofs
+                .commitment_infos
+                .contracts_trie_commitment_info,
+            address_to_storage_commitment_info: storage_proofs
+                .commitment_infos
+                .storage_tries_commitment_infos,
+            contract_class_commitment_info: storage_proofs
+                .commitment_infos
+                .classes_trie_commitment_info,
+            transactions: txs
+                .into_iter()
+                .map(|(invoke_tx, tx_hash)| {
+                    starknet_api::executable_transaction::Transaction::Account(
+                        starknet_api::executable_transaction::AccountTransaction::Invoke(
+                            starknet_api::executable_transaction::InvokeTransaction {
+                                tx: invoke_tx,
+                                tx_hash,
+                            },
+                        ),
+                    )
+                })
+                .collect(),
+            tx_execution_infos,
+            // We assume that no classes are declared.
+            declared_class_hash_to_component_hashes: HashMap::new(),
+            // We assume that no classes are migrated.
+            class_hashes_to_migrate: HashMap::new(),
+            // Stored block hash buffer skipped for now.
+            old_block_number_and_hash: None,
+            // Block info and hashes from execution context.
+            block_info: execution_data.block_context.block_info().clone(),
+            // TODO(Aviv): Add prev and new block hashes.
+            prev_block_hash: BlockHash::default(),
+            new_block_hash: BlockHash::default(),
+            initial_reads: execution_data.initial_reads,
+        };
+
+        // 7. Final StarknetOsInput (consuming classes).
+        Ok(StarknetOsInput {
+            os_block_inputs: vec![os_block_input],
+            deprecated_compiled_classes: classes.deprecated_compiled_classes,
+            compiled_classes: classes.compiled_classes,
+        })
+    }
+}

crates/starknet_os_runner/src/virtual_block_executor.rs

@@ -18,7 +18,7 @@ use blockifier_reexecution::state_reader::rpc_state_reader::RpcStateReader;
 use starknet_api::block::BlockNumber;
 use starknet_api::core::ClassHash;
 use starknet_api::transaction::fields::Fee;
-use starknet_api::transaction::{Transaction, TransactionHash};
+use starknet_api::transaction::{InvokeTransaction, Transaction, TransactionHash};
 
 use crate::errors::VirtualBlockExecutorError;
 
@@ -77,12 +77,12 @@ pub trait VirtualBlockExecutor {
     /// # Returns
     ///
     /// Returns `VirtualBlockExecutionData` containing execution outputs for all
-    /// transactions, or an error if any transaction fails or is not an Invoke.
+    /// transactions, or an error if any transaction fails.
     fn execute(
         &self,
         block_number: BlockNumber,
         contract_class_manager: ContractClassManager,
-        txs: Vec<(Transaction, TransactionHash)>,
+        txs: Vec<(InvokeTransaction, TransactionHash)>,
     ) -> Result<VirtualBlockExecutionData, VirtualBlockExecutorError> {
         let blockifier_txs = Self::convert_invoke_txs(txs)?;
         let block_context = self.block_context(block_number)?;
@@ -138,38 +138,31 @@ pub trait VirtualBlockExecutor {
 
     /// Converts Invoke transactions to blockifier transactions.
     ///
-    /// Uses execution flags that skip fee charging and nonce check.
-    /// Returns an error if any transaction is not an Invoke.
+    /// Uses execution flags that skip strict nonce check for virtual block execution.
     fn convert_invoke_txs(
-        txs: Vec<(Transaction, TransactionHash)>,
+        txs: Vec<(InvokeTransaction, TransactionHash)>,
     ) -> Result<Vec<BlockifierTransaction>, VirtualBlockExecutorError> {
         txs.into_iter()
-            .map(|(tx, tx_hash)| {
-                if let Transaction::Invoke(invoke_tx) = tx {
-                    // Execute with validation, conditional fee charging based on resource bounds,
-                    // but skip strict nonce check for virtual block execution.
-                    let execution_flags = ExecutionFlags {
-                        only_query: false,
-                        charge_fee: invoke_tx.resource_bounds().max_possible_fee(invoke_tx.tip())
-                            > Fee(0),
-                        validate: true,
-                        strict_nonce_check: false,
-                    };
-
-                    BlockifierTransaction::from_api(
-                        Transaction::Invoke(invoke_tx),
-                        tx_hash,
-                        None, // class_info - not needed for Invoke.
-                        None, // paid_fee_on_l1 - not needed for Invoke.
-                        None, // deployed_contract_address - not needed for Invoke.
-                        execution_flags,
-                    )
-                    .map_err(|e| {
-                        VirtualBlockExecutorError::TransactionExecutionError(e.to_string())
-                    })
-                } else {
-                    Err(VirtualBlockExecutorError::UnsupportedTransactionType)
-                }
+            .map(|(invoke_tx, tx_hash)| {
+                // Execute with validation, conditional fee charging based on resource bounds,
+                // but skip strict nonce check for virtual block execution.
+                let execution_flags = ExecutionFlags {
+                    only_query: false,
+                    charge_fee: invoke_tx.resource_bounds().max_possible_fee(invoke_tx.tip())
+                        > Fee(0),
+                    validate: true,
+                    strict_nonce_check: false,
+                };
+
+                BlockifierTransaction::from_api(
+                    Transaction::Invoke(invoke_tx),
+                    tx_hash,
+                    None, // class_info - not needed for Invoke.
+                    None, // paid_fee_on_l1 - not needed for Invoke.
+                    None, // deployed_contract_address - not needed for Invoke.
+                    execution_flags,
+                )
+                .map_err(|e| VirtualBlockExecutorError::TransactionExecutionError(e.to_string()))
             })
             .collect()
     }

crates/starknet_os_runner/src/virtual_block_executor_test.rs

@@ -7,7 +7,7 @@ use starknet_api::abi::abi_utils::{get_storage_var_address, selector_from_name};
 use starknet_api::block::BlockNumber;
 use starknet_api::core::{ChainId, ContractAddress, Nonce};
 use starknet_api::test_utils::invoke::invoke_tx;
-use starknet_api::transaction::{Transaction, TransactionHash};
+use starknet_api::transaction::{InvokeTransaction, Transaction, TransactionHash};
 use starknet_api::{calldata, felt, invoke_tx_args};
 
 use crate::errors::VirtualBlockExecutorError;
@@ -40,7 +40,7 @@ impl VirtualBlockExecutor for TestRpcVirtualBlockExecutor {
 
     // Override the default implementation to skip validation.
     fn convert_invoke_txs(
-        txs: Vec<(Transaction, TransactionHash)>,
+        txs: Vec<(InvokeTransaction, TransactionHash)>,
     ) -> Result<Vec<BlockifierTransaction>, VirtualBlockExecutorError> {
         // Call the default trait implementation.
         let mut blockifier_txs = RpcStateReader::convert_invoke_txs(txs)?;
@@ -60,7 +60,7 @@ impl VirtualBlockExecutor for TestRpcVirtualBlockExecutor {
 ///
 /// Since we skip validation and fee charging, we can use dummy values for signature,
 /// nonce, and resource bounds.
-fn construct_balance_of_invoke() -> (Transaction, TransactionHash) {
+fn construct_balance_of_invoke() -> (InvokeTransaction, TransactionHash) {
     let strk_token = ContractAddress::try_from(STRK_TOKEN_ADDRESS).unwrap();
     let sender = ContractAddress::try_from(SENDER_ADDRESS).unwrap();
 
@@ -85,9 +85,10 @@ fn construct_balance_of_invoke() -> (Transaction, TransactionHash) {
         nonce: Nonce(felt!("0x1000000")),
     });
 
-    let tx = Transaction::Invoke(invoke_tx);
-    let tx_hash = tx.calculate_transaction_hash(&ChainId::Mainnet).unwrap();
-    (tx, tx_hash)
+    let tx_hash = Transaction::Invoke(invoke_tx.clone())
+        .calculate_transaction_hash(&ChainId::Mainnet)
+        .unwrap();
+    (invoke_tx, tx_hash)
 }
 
 /// Integration test for RpcVirtualBlockExecutor with a constructed transaction.