diff --git a/crates/native_blockifier/src/py_transaction.rs b/crates/native_blockifier/src/py_transaction.rs index 5484c6c0eb..1e15e115a9 100644 --- a/crates/native_blockifier/src/py_transaction.rs +++ b/crates/native_blockifier/src/py_transaction.rs @@ -279,11 +279,68 @@ pub fn py_tx( } } +/// Wraps the transaction executor in an optional, to allow an explicit deallocation of it. +/// The explicit deallocation is needed since PyO3 can't track lifetimes within Python. #[pyclass] +pub struct PyTransactionExecutor { + pub executor: Option, +} + +#[pymethods] +impl PyTransactionExecutor { + #[new] + #[args(general_config, block_info, papyrus_storage)] + pub fn create( + general_config: &PyAny, + block_info: &PyAny, + papyrus_storage: &Storage, + ) -> NativeBlockifierResult { + log::debug!("Initializing Transaction Executor..."); + let executor = + PyTransactionExecutorInner::create(general_config, block_info, papyrus_storage)?; + log::debug!("Initialized Transaction Executor."); + + Ok(Self { executor: Some(executor) }) + } + + #[args(tx, raw_contract_class, enough_room_for_tx)] + pub fn execute( + &mut self, + tx: &PyAny, + raw_contract_class: Option<&str>, + // This is functools.partial(bouncer.add, tw_written=tx_written). + enough_room_for_tx: &PyAny, + ) -> NativeBlockifierResult<( + Py, + HashMap, + )> { + self.executor().execute(tx, raw_contract_class, enough_room_for_tx) + } + + pub fn finalize(&mut self) -> PyStateDiff { + log::debug!("Finalizing execution..."); + let state_diff = self.executor().finalize(); + self.close(); + log::debug!("Finalized execution."); + + state_diff + } + + pub fn close(&mut self) { + self.executor = None; + } +} + +impl PyTransactionExecutor { + fn executor(&mut self) -> &mut PyTransactionExecutorInner { + self.executor.as_mut().expect("Transaction executor should be initialized.") + } +} + // To access a field you must use `self.borrow_{field_name}()`. // Alternately, you can borrow the whole object using `self.with[_mut]()`. #[ouroboros::self_referencing] -pub struct PyTransactionExecutor { +pub struct PyTransactionExecutorInner { pub block_context: BlockContext, // State-related fields. @@ -297,31 +354,22 @@ pub struct PyTransactionExecutor { pub state: CachedState>, } -#[pymethods] -impl PyTransactionExecutor { - #[new] - #[args(general_config, block_info, papyrus_storage)] +impl PyTransactionExecutorInner { pub fn create( general_config: &PyAny, block_info: &PyAny, papyrus_storage: &Storage, ) -> NativeBlockifierResult { - log::debug!("Initializing Transaction Executor..."); - // Assumption: storage is aligned. let reader = papyrus_storage.reader().clone(); let block_context = py_block_context(general_config, block_info)?; - let build_result = build_tx_executor(block_context, reader); - log::debug!("Initialized Transaction Executor."); - - build_result + build_tx_executor(block_context, reader) } /// Executes the given transaction on the state maintained by the executor. /// Returns the execution trace, together with the compiled class hashes of executed classes /// (used for counting purposes). - #[args(tx, raw_contract_class, enough_room_for_tx)] pub fn execute( &mut self, tx: &PyAny, @@ -395,18 +443,14 @@ impl PyTransactionExecutor { /// Returns the state diff resulting in executing transactions. pub fn finalize(&mut self) -> PyStateDiff { - log::debug!("Finalizing execution..."); - let state_diff = PyStateDiff::from(self.borrow_state().to_state_diff()); - log::debug!("Finalized execution."); - - state_diff + PyStateDiff::from(self.borrow_state().to_state_diff()) } } pub fn build_tx_executor( block_context: BlockContext, storage_reader: papyrus_storage::StorageReader, -) -> NativeBlockifierResult { +) -> NativeBlockifierResult { // The following callbacks are required to capture the local lifetime parameter. fn storage_tx_builder( storage_reader: &papyrus_storage::StorageReader, @@ -425,7 +469,7 @@ pub fn build_tx_executor( let block_number = block_context.block_number; // The builder struct below is implicitly created by `ouroboros`. - let py_tx_executor_builder = PyTransactionExecutorTryBuilder { + let py_tx_executor_builder = PyTransactionExecutorInnerTryBuilder { block_context, storage_reader, storage_tx_builder,