starknet_os_runner: virtual block executor

AvivYossef-starkware · AvivYossef-starkware · commit 1b993b5df146 · 2025-12-21T13:02:40.000+02:00
diff --git a/crates/starknet_os_runner/src/virtual_block_executor.rs b/crates/starknet_os_runner/src/virtual_block_executor.rs
@@ -6,12 +6,15 @@ use blockifier::blockifier::transaction_executor::{
 use blockifier::context::BlockContext;
 use blockifier::state::cached_state::{CachedState, StateMaps};
 use blockifier::state::contract_class_manager::ContractClassManager;
-use blockifier::state::state_reader_and_contract_manager::StateReaderAndContractManager;
+use blockifier::state::state_reader_and_contract_manager::{
+    FetchCompiledClasses,
+    StateReaderAndContractManager,
+};
 use blockifier::transaction::account_transaction::ExecutionFlags;
 use blockifier::transaction::transaction_execution::Transaction as BlockifierTransaction;
 use blockifier_reexecution::state_reader::rpc_state_reader::RpcStateReader;
 use starknet_api::block::BlockNumber;
-use starknet_api::core::ChainId;
+use starknet_api::transaction::fields::Fee;
 use starknet_api::transaction::{Transaction, TransactionHash};
 
 use crate::errors::VirtualBlockExecutorError;
@@ -32,8 +35,9 @@ pub struct VirtualBlockExecutionData {
 
 /// Executes a virtual block of transactions.
 ///
-/// A virtual block executor runs transactions without block preprocessing
-/// (`pre_process_block`), which is useful for simulating execution or generating
+/// A virtual block executor runs transactions on top of a given, finalized block.
+///  This means that some parts, like block preprocessing
+/// (`pre_process_block`), are skipped. Useful for simulating execution or generating
 /// OS input for proving.
 ///
 /// Implementations can fetch state from different sources (RPC nodes, local state,
@@ -47,13 +51,13 @@ pub struct VirtualBlockExecutionData {
 /// # Examples
 ///
 /// ```text
-/// let executor = RpcVirtualBlockExecutor::new(
+/// let executor = RpcStateReader::new_with_config_from_url(
 ///     "http://localhost:9545".to_string(),
-///     ChainId::Mainnet,
-///     contract_class_manager,
+///      ChainId::Mainnet,
+///      BlockNumber(1000),
 /// );
 ///
-/// let execution_data = executor.execute(block_number, transactions)?;
+/// let execution_data = executor.execute(block_number, contract_class_manager, transactions)?;
 /// // Use execution_data to build OS input for proving...
 /// ```
 pub trait VirtualBlockExecutor {
@@ -62,6 +66,7 @@ pub trait VirtualBlockExecutor {
     /// # Arguments
     ///
     /// * `block_number` - The block number to use for state and context
+    /// * `contract_class_manager` - Manager for compiled contract classes
     /// * `txs` - Invoke transactions to execute (with their hashes)
     ///
     /// # Returns
@@ -71,105 +76,16 @@ pub trait VirtualBlockExecutor {
     fn execute(
         &self,
         block_number: BlockNumber,
+        contract_class_manager: ContractClassManager,
         txs: Vec<(Transaction, TransactionHash)>,
     ) -> Result<VirtualBlockExecutionData, VirtualBlockExecutorError> {
         let blockifier_txs = Self::convert_invoke_txs(txs)?;
-        self.execute_inner(block_number, blockifier_txs)
-    }
-
-    fn execute_inner(
-        &self,
-        block_number: BlockNumber,
-        txs: Vec<BlockifierTransaction>,
-    ) -> Result<VirtualBlockExecutionData, VirtualBlockExecutorError>;
-
-    /// 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.
-    fn convert_invoke_txs(
-        txs: Vec<(Transaction, TransactionHash)>,
-    ) -> Result<Vec<BlockifierTransaction>, VirtualBlockExecutorError> {
-        // Skip validation, fee charging, and nonce check for virtual block execution.
-        let execution_flags = ExecutionFlags {
-            validate: true,
-            charge_fee: false,
-            strict_nonce_check: false,
-            only_query: false,
-        };
-
-        txs.into_iter()
-            .map(|(tx, tx_hash)| {
-                if !matches!(tx, Transaction::Invoke(_)) {
-                    return Err(VirtualBlockExecutorError::UnsupportedTransactionType);
-                }
-
-                BlockifierTransaction::from_api(
-                    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.clone(),
-                )
-                .map_err(|e| VirtualBlockExecutorError::TransactionExecutionError(e.to_string()))
-            })
-            .collect()
-    }
-}
-
-/// RPC-based virtual block executor.
-///
-/// This executor fetches historical state from an RPC node and executes transactions
-/// without block preprocessing. Validation and fee charging are always skipped,
-/// making it suitable for simulation and OS input generation.
-pub struct RpcVirtualBlockExecutor {
-    node_url: String,
-    chain_id: ChainId,
-    contract_class_manager: ContractClassManager,
-}
-
-impl RpcVirtualBlockExecutor {
-    /// Creates a new RPC-based virtual block executor.
-    ///
-    /// # Arguments
-    ///
-    /// * `node_url` - URL of the RPC node to fetch state from
-    /// * `chain_id` - The chain ID for transaction hash computation
-    /// * `contract_class_manager` - Manager for compiled contract classes
-    pub fn new(
-        node_url: String,
-        chain_id: ChainId,
-        contract_class_manager: ContractClassManager,
-    ) -> Self {
-        Self { node_url, chain_id, contract_class_manager }
-    }
-}
-
-impl VirtualBlockExecutor for RpcVirtualBlockExecutor {
-    fn execute_inner(
-        &self,
-        block_number: BlockNumber,
-        txs: Vec<BlockifierTransaction>,
-    ) -> Result<VirtualBlockExecutionData, VirtualBlockExecutorError> {
-        // Create RPC state reader for the given block.
-        let rpc_state_reader = RpcStateReader::new_with_config_from_url(
-            self.node_url.clone(),
-            self.chain_id.clone(),
-            block_number,
-        );
-
-        // Get block context from RPC.
-        let block_context = rpc_state_reader
-            .get_block_context()
-            .map_err(|e| VirtualBlockExecutorError::ReexecutionError(Box::new(e)))?;
+        let block_context = self.block_context(block_number)?;
+        let state_reader = self.state_reader(block_number)?;
 
         // Create state reader with contract manager.
-        let state_reader_and_contract_manager = StateReaderAndContractManager::new(
-            rpc_state_reader,
-            self.contract_class_manager.clone(),
-            None,
-        );
+        let state_reader_and_contract_manager =
+            StateReaderAndContractManager::new(state_reader, contract_class_manager, None);
 
         let block_state = CachedState::new(state_reader_and_contract_manager);
 
@@ -181,7 +97,7 @@ impl VirtualBlockExecutor for RpcVirtualBlockExecutor {
         );
 
         // Execute all transactions.
-        let execution_results = transaction_executor.execute_txs(&txs, None);
+        let execution_results = transaction_executor.execute_txs(&blockifier_txs, None);
 
         // Collect results, returning error if any transaction fails.
         let execution_outputs: Vec<TransactionExecutionOutput> = execution_results
@@ -203,4 +119,77 @@ impl VirtualBlockExecutor for RpcVirtualBlockExecutor {
 
         Ok(VirtualBlockExecutionData { execution_outputs, block_context, initial_reads })
     }
+
+    /// 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.
+    fn convert_invoke_txs(
+        txs: Vec<(Transaction, 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)
+                }
+            })
+            .collect()
+    }
+    /// Returns the block context for the given block number.
+    fn block_context(
+        &self,
+        block_number: BlockNumber,
+    ) -> Result<BlockContext, VirtualBlockExecutorError>;
+
+    /// Returns a state reader that implements `FetchCompiledClasses` for the given block number.
+    /// Must be `Send + Sync + 'static` to be used in the transaction executor.
+    fn state_reader(
+        &self,
+        block_number: BlockNumber,
+    ) -> Result<impl FetchCompiledClasses + Send + Sync + 'static, VirtualBlockExecutorError>;
+}
+
+/// RPC-based virtual block executor.
+///
+/// This executor fetches historical state from an RPC node and executes transactions
+/// without block preprocessing. Validation and fee charging are always skipped,
+/// making it suitable for simulation and OS input generation.
+impl VirtualBlockExecutor for RpcStateReader {
+    fn block_context(
+        &self,
+        _block_number: BlockNumber,
+    ) -> Result<BlockContext, VirtualBlockExecutorError> {
+        self.get_block_context()
+            .map_err(|e| VirtualBlockExecutorError::ReexecutionError(Box::new(e)))
+    }
+
+    fn state_reader(
+        &self,
+        _block_number: BlockNumber,
+    ) -> Result<impl FetchCompiledClasses + Send + Sync + 'static, VirtualBlockExecutorError> {
+        // RpcStateReader itself is the state reader
+        Ok(self.clone())
+    }
 }
diff --git a/crates/starknet_os_runner/src/virtual_block_executor_test.rs b/crates/starknet_os_runner/src/virtual_block_executor_test.rs
@@ -2,16 +2,17 @@ use std::env;
 
 use blockifier::blockifier::config::ContractClassManagerConfig;
 use blockifier::state::contract_class_manager::ContractClassManager;
-use blockifier::transaction::account_transaction::ExecutionFlags;
 use blockifier::transaction::transaction_execution::Transaction as BlockifierTransaction;
+use blockifier_reexecution::state_reader::rpc_state_reader::RpcStateReader;
 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;
+use starknet_api::transaction::{Transaction, TransactionHash};
 use starknet_api::{calldata, felt, invoke_tx_args};
 
-use crate::virtual_block_executor::{RpcVirtualBlockExecutor, VirtualBlockExecutor};
+use crate::errors::VirtualBlockExecutorError;
+use crate::virtual_block_executor::VirtualBlockExecutor;
 
 /// Block number to use for testing (mainnet block with known state).
 const TEST_BLOCK_NUMBER: u64 = 800000;
@@ -25,11 +26,53 @@ const STRK_TOKEN_ADDRESS: &str =
 /// [call_array_len, (to, selector, data_offset, data_len)..., calldata_len, calldata...]
 const SENDER_ADDRESS: &str = "0x01176a1bd84444c89232ec27754698e5d2e7e1a7f1539f12027f28b23ec9f3d8";
 
+/// Test wrapper for RpcStateReader that overrides execution flags to skip validation.
+struct TestRpcVirtualBlockExecutor(RpcStateReader);
+
+impl VirtualBlockExecutor for TestRpcVirtualBlockExecutor {
+    fn block_context(
+        &self,
+        block_number: BlockNumber,
+    ) -> Result<blockifier::context::BlockContext, VirtualBlockExecutorError> {
+        self.0.block_context(block_number)
+    }
+
+    fn state_reader(
+        &self,
+        block_number: BlockNumber,
+    ) -> Result<
+        impl blockifier::state::state_reader_and_contract_manager::FetchCompiledClasses
+        + Send
+        + Sync
+        + 'static,
+        VirtualBlockExecutorError,
+    > {
+        self.0.state_reader(block_number)
+    }
+
+    // Override the default implementation to skip validation.
+    fn convert_invoke_txs(
+        txs: Vec<(Transaction, TransactionHash)>,
+    ) -> Result<Vec<BlockifierTransaction>, VirtualBlockExecutorError> {
+        // Call the default trait implementation.
+        let mut blockifier_txs = RpcStateReader::convert_invoke_txs(txs)?;
+
+        // Modify validate flag to false for all transactions.
+        for tx in &mut blockifier_txs {
+            if let BlockifierTransaction::Account(account_tx) = tx {
+                account_tx.execution_flags.validate = false;
+            }
+        }
+
+        Ok(blockifier_txs)
+    }
+}
+
 /// Constructs an Invoke transaction that calls `balanceOf` on the STRK token contract.
 ///
 /// Since we skip validation and fee charging, we can use dummy values for signature,
 /// nonce, and resource bounds.
-fn construct_balance_of_invoke() -> BlockifierTransaction {
+fn construct_balance_of_invoke() -> (Transaction, TransactionHash) {
     let strk_token = ContractAddress::try_from(felt!(STRK_TOKEN_ADDRESS)).unwrap();
     let sender = ContractAddress::try_from(felt!(SENDER_ADDRESS)).unwrap();
 
@@ -56,24 +99,7 @@ fn construct_balance_of_invoke() -> BlockifierTransaction {
 
     let tx = Transaction::Invoke(invoke_tx);
     let tx_hash = tx.calculate_transaction_hash(&ChainId::Mainnet).unwrap();
-
-    // Skip fee charging, nonce check and validation.
-    let execution_flags = ExecutionFlags {
-        validate: false,
-        charge_fee: false,
-        strict_nonce_check: false,
-        only_query: false,
-    };
-
-    BlockifierTransaction::from_api(
-        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,
-    )
-    .unwrap()
+    (tx, tx_hash)
 }
 
 /// Integration test for RpcVirtualBlockExecutor with a constructed transaction.
@@ -93,22 +119,26 @@ fn construct_balance_of_invoke() -> BlockifierTransaction {
 /// NODE_URL=https://your-rpc-node cargo test -p starknet_os_runner -- --ignored
 /// ```
 #[test]
-#[ignore] // Requires RPC access - run with: cargo test -p starknet_os_runner -- --ignored
+#[ignore] // Requires RPC access 
 fn test_execute_constructed_balance_of_transaction() {
     let node_url =
         env::var("NODE_URL").expect("NODE_URL environment variable required for this test");
 
     // Construct a balanceOf transaction (with execution flags set).
-    let tx = construct_balance_of_invoke();
+    let (tx, tx_hash) = construct_balance_of_invoke();
 
     // Create the virtual block executor.
     let contract_class_manager = ContractClassManager::start(ContractClassManagerConfig::default());
-    let executor = RpcVirtualBlockExecutor::new(node_url, ChainId::Mainnet, contract_class_manager);
+    let executor = TestRpcVirtualBlockExecutor(RpcStateReader::new_with_config_from_url(
+        node_url,
+        ChainId::Mainnet,
+        BlockNumber(TEST_BLOCK_NUMBER),
+    ));
 
     // Execute the transaction.
     let result = executor
-        .execute_inner(BlockNumber(TEST_BLOCK_NUMBER), vec![tx])
-        .expect("Virtual block execution should succeed");
+        .execute(BlockNumber(TEST_BLOCK_NUMBER), contract_class_manager, vec![(tx, tx_hash)])
+        .unwrap();
 
     // Verify execution produced output.
     assert_eq!(result.execution_outputs.len(), 1, "Should have exactly one execution output");
