From f7c6093a5da24ccd063d84ff0bafa0bad58bf62b Mon Sep 17 00:00:00 2001
From: Barak <barakb@starkware.co>
Date: Thu, 1 Feb 2024 20:26:01 +0200
Subject: [PATCH] refactor(execution, native_blockifier): make
 TransactionExecutor.finalize() work without Py objects

---
 crates/native_blockifier/src/py_block_executor.rs  |  9 ++++++++-
 .../native_blockifier/src/transaction_executor.rs  | 14 +++++++++-----
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/crates/native_blockifier/src/py_block_executor.rs b/crates/native_blockifier/src/py_block_executor.rs
index e0cff99135..82527332c8 100644
--- a/crates/native_blockifier/src/py_block_executor.rs
+++ b/crates/native_blockifier/src/py_block_executor.rs
@@ -118,7 +118,14 @@ impl PyBlockExecutor {
     /// visited PC values.
     pub fn finalize(&mut self, is_pending_block: bool) -> (PyStateDiff, Vec<(PyFelt, Vec<usize>)>) {
         log::debug!("Finalizing execution...");
-        let finalized_state = self.tx_executor().finalize(is_pending_block);
+        let (commitment_state_diff, visited_pcs) = self.tx_executor().finalize(is_pending_block);
+        let visited_pcs = visited_pcs
+            .iter()
+            .map(|(class_hash, class_visited_pcs_vec)| {
+                (PyFelt::from(*class_hash), class_visited_pcs_vec.clone())
+            })
+            .collect();
+        let finalized_state = (PyStateDiff::from(commitment_state_diff), visited_pcs);
         log::debug!("Finalized execution.");
 
         finalized_state
diff --git a/crates/native_blockifier/src/transaction_executor.rs b/crates/native_blockifier/src/transaction_executor.rs
index 8f6a098c82..06043a6a2c 100644
--- a/crates/native_blockifier/src/transaction_executor.rs
+++ b/crates/native_blockifier/src/transaction_executor.rs
@@ -8,7 +8,8 @@ use blockifier::execution::call_info::{CallInfo, MessageL1CostInfo};
 use blockifier::execution::entry_point::ExecutionResources;
 use blockifier::fee::actual_cost::ActualCost;
 use blockifier::state::cached_state::{
-    CachedState, GlobalContractCache, StagedTransactionalState, StorageEntry, TransactionalState,
+    CachedState, CommitmentStateDiff, GlobalContractCache, StagedTransactionalState, StorageEntry,
+    TransactionalState,
 };
 use blockifier::state::state_api::{State, StateReader};
 use blockifier::transaction::account_transaction::AccountTransaction;
@@ -23,7 +24,7 @@ use starknet_api::core::ClassHash;
 
 use crate::errors::{NativeBlockifierError, NativeBlockifierResult};
 use crate::py_block_executor::{into_block_context, PyGeneralConfig};
-use crate::py_state_diff::{PyBlockInfo, PyStateDiff};
+use crate::py_state_diff::PyBlockInfo;
 use crate::py_transaction_execution_info::PyBouncerInfo;
 use crate::py_utils::PyFelt;
 
@@ -176,7 +177,10 @@ impl<S: StateReader> TransactionExecutor<S> {
 
     /// Returns the state diff and a list of contract class hash with the corresponding list of
     /// visited PC values.
-    pub fn finalize(&mut self, is_pending_block: bool) -> (PyStateDiff, Vec<(PyFelt, Vec<usize>)>) {
+    pub fn finalize(
+        &mut self,
+        is_pending_block: bool,
+    ) -> (CommitmentStateDiff, Vec<(ClassHash, Vec<usize>)>) {
         // Do not cache classes that were declared during a pending block.
         // They will be redeclared, and should not be cached since the content of this block is
         // transient.
@@ -192,11 +196,11 @@ impl<S: StateReader> TransactionExecutor<S> {
             .map(|(class_hash, class_visited_pcs)| {
                 let mut class_visited_pcs_vec: Vec<_> = class_visited_pcs.iter().cloned().collect();
                 class_visited_pcs_vec.sort();
-                (PyFelt::from(*class_hash), class_visited_pcs_vec)
+                (*class_hash, class_visited_pcs_vec)
             })
             .collect();
 
-        (PyStateDiff::from(self.state.to_state_diff()), visited_pcs)
+        (self.state.to_state_diff(), visited_pcs)
     }
 
     // Block pre-processing; see `block::pre_process_block` documentation.