apollo_batcher: thread native classes whitelist from create block builder to manager

Author: ArniStarkware

crates/apollo_batcher/src/batcher.rs

@@ -354,6 +354,12 @@ impl Batcher {
                         .block_builder_config
                         .proposer_idle_detection_delay_millis,
                 },
+                self.config
+                    .static_config
+                    .contract_class_manager_config
+                    .cairo_native_run_config
+                    .native_classes_whitelist
+                    .clone(),
                 Box::new(tx_provider),
                 Some(output_tx_sender),
                 Some(candidate_tx_sender),
@@ -442,6 +448,12 @@ impl Batcher {
                         .block_builder_config
                         .proposer_idle_detection_delay_millis,
                 },
+                self.config
+                    .static_config
+                    .contract_class_manager_config
+                    .cairo_native_run_config
+                    .native_classes_whitelist
+                    .clone(),
                 Box::new(tx_provider),
                 None,
                 None,

crates/apollo_batcher/src/batcher_test.rs

@@ -292,7 +292,7 @@ fn mock_create_builder_for_validate_block(
     build_block_result: BlockBuilderResult<BlockExecutionArtifacts>,
 ) {
     block_builder_factory.expect_create_block_builder().times(1).return_once(
-        |_, _, tx_provider, _, _, _, _| {
+        |_, _, _, tx_provider, _, _, _, _| {
             let block_builder = FakeValidateBlockBuilder {
                 tx_provider,
                 build_block_result: Some(build_block_result),
@@ -317,7 +317,7 @@ fn mock_create_builder_for_propose_block(
     build_block_result: BlockBuilderResult<BlockExecutionArtifacts>,
 ) {
     block_builder_factory.expect_create_block_builder().times(1).return_once(
-        move |_, _, tx_provider, output_content_sender, _, _, _| {
+        move |_, _, _, tx_provider, output_content_sender, _, _, _| {
             let block_builder = FakeProposeBlockBuilder {
                 output_content_sender: output_content_sender.unwrap(),
                 output_txs,

crates/apollo_batcher/src/block_builder.rs

@@ -17,6 +17,7 @@ use apollo_state_reader::apollo_state::{ApolloReader, ClassReader};
 use apollo_storage::StorageReader;
 use async_trait::async_trait;
 use blockifier::blockifier::concurrent_transaction_executor::ConcurrentTransactionExecutor;
+use blockifier::blockifier::config::NativeClassesWhitelist;
 use blockifier::blockifier::transaction_executor::{
     BlockExecutionSummary,
     CompiledClassHashesForMigration,
@@ -759,6 +760,7 @@ pub trait BlockBuilderFactoryTrait: Send + Sync {
         &self,
         block_metadata: BlockMetadata,
         execution_params: BlockBuilderExecutionParams,
+        native_classes_whitelist: NativeClassesWhitelist,
         tx_provider: Box<dyn TransactionProvider>,
         output_content_sender: Option<
             tokio::sync::mpsc::UnboundedSender<InternalConsensusTransaction>,
@@ -782,6 +784,7 @@ impl BlockBuilderFactory {
     fn preprocess_and_create_transaction_executor(
         &self,
         block_metadata: BlockMetadata,
+        native_classes_whitelist: NativeClassesWhitelist,
         runtime: tokio::runtime::Handle,
     ) -> BlockBuilderResult<ConcurrentTransactionExecutor<ApolloStateReaderAndContractManager>>
     {
@@ -804,9 +807,10 @@ impl BlockBuilderFactory {
         let class_reader = Some(ClassReader { reader: self.class_manager_client.clone(), runtime });
         let apollo_reader =
             ApolloReader::new_with_class_reader(self.storage_reader.clone(), height, class_reader);
-        let state_reader = StateReaderAndContractManager::new(
+        let state_reader = StateReaderAndContractManager::new_with_native_classes_whitelist(
             apollo_reader,
             self.contract_class_manager.clone(),
+            native_classes_whitelist,
             Some(BATCHER_CLASS_CACHE_METRICS),
         );
 
@@ -827,6 +831,7 @@ impl BlockBuilderFactoryTrait for BlockBuilderFactory {
         &self,
         block_metadata: BlockMetadata,
         execution_params: BlockBuilderExecutionParams,
+        native_classes_whitelist: NativeClassesWhitelist,
         tx_provider: Box<dyn TransactionProvider>,
         output_content_sender: Option<
             tokio::sync::mpsc::UnboundedSender<InternalConsensusTransaction>,
@@ -835,7 +840,11 @@ impl BlockBuilderFactoryTrait for BlockBuilderFactory {
         pre_confirmed_tx_sender: Option<PreconfirmedTxSender>,
         runtime: tokio::runtime::Handle,
     ) -> BlockBuilderResult<(Box<dyn BlockBuilderTrait>, AbortSignalSender)> {
-        let executor = self.preprocess_and_create_transaction_executor(block_metadata, runtime)?;
+        let executor = self.preprocess_and_create_transaction_executor(
+            block_metadata,
+            native_classes_whitelist,
+            runtime,
+        )?;
         let (abort_signal_sender, abort_signal_receiver) = tokio::sync::oneshot::channel();
         let transaction_converter = TransactionConverter::new(
             self.class_manager_client.clone(),

crates/blockifier/src/state/contract_class_manager.rs

@@ -10,7 +10,7 @@ pub mod trivial_class_manager {
     use starknet_api::class_cache::GlobalContractCache;
     use starknet_api::core::{ClassHash, CompiledClassHash};
 
-    use crate::blockifier::config::ContractClassManagerConfig;
+    use crate::blockifier::config::{ContractClassManagerConfig, NativeClassesWhitelist};
     use crate::execution::contract_class::RunnableCompiledClass;
     use crate::state::global_cache::{CompiledClasses, RawClassCache};
 
@@ -33,7 +33,11 @@ pub mod trivial_class_manager {
             }
         }
 
-        pub fn get_runnable(&self, class_hash: &ClassHash) -> Option<RunnableCompiledClass> {
+        pub fn get_runnable(
+            &self,
+            class_hash: &ClassHash,
+            _native_classes_whitelist: &NativeClassesWhitelist,
+        ) -> Option<RunnableCompiledClass> {
             Some(self.class_cache.get(class_hash)?.to_runnable())
         }
 

crates/blockifier/src/state/native_class_manager.rs

@@ -13,7 +13,11 @@ use starknet_api::core::{ClassHash, CompiledClassHash};
 use starknet_api::state::SierraContractClass;
 use thiserror::Error;
 
-use crate::blockifier::config::{CairoNativeRunConfig, ContractClassManagerConfig};
+use crate::blockifier::config::{
+    CairoNativeRunConfig,
+    ContractClassManagerConfig,
+    NativeClassesWhitelist,
+};
 use crate::execution::contract_class::{CompiledClassV1, RunnableCompiledClass};
 use crate::execution::native::contract_class::NativeCompiledClassV1;
 use crate::metrics::NATIVE_COMPILATION_ERROR;
@@ -117,7 +121,11 @@ impl NativeClassManager {
     }
 
     /// Returns the runnable compiled class for the given class hash, if it exists in class_cache.
-    pub fn get_runnable(&self, class_hash: &ClassHash) -> Option<RunnableCompiledClass> {
+    pub fn get_runnable(
+        &self,
+        class_hash: &ClassHash,
+        native_classes_whitelist: &NativeClassesWhitelist,
+    ) -> Option<RunnableCompiledClass> {
         let cached_class = self.class_cache.get(class_hash)?;
 
         let cached_class = match cached_class {
@@ -131,7 +139,7 @@ impl NativeClassManager {
                 cached_class
             }
             CompiledClasses::V1Native(CachedCairoNative::Compiled(native))
-                if !self.run_class_with_cairo_native(class_hash) =>
+                if !native_classes_whitelist.contains(class_hash) =>
             {
                 CompiledClasses::V1(native.casm(), Arc::new(SierraContractClass::default()))
             }
@@ -219,11 +227,6 @@ impl NativeClassManager {
         self.cairo_native_run_config.wait_on_native_compilation
     }
 
-    /// Determines if a contract should run with cairo native based on the whitelist.
-    pub fn run_class_with_cairo_native(&self, class_hash: &ClassHash) -> bool {
-        self.cairo_native_run_config.native_classes_whitelist.contains(class_hash)
-    }
-
     /// Clears the contract class_cache.
     pub fn clear(&mut self) {
         self.class_cache.clear();

crates/blockifier/src/state/native_class_manager_test.rs

@@ -254,7 +254,7 @@ fn test_native_classes_whitelist(
         wait_on_native_compilation: true,
         panic_on_compilation_failure: true,
         channel_size: TEST_CHANNEL_SIZE,
-        native_classes_whitelist: whitelist,
+        native_classes_whitelist: whitelist.clone(),
     };
     let manager = NativeClassManager::create_for_testing(native_config);
 
@@ -264,12 +264,12 @@ fn test_native_classes_whitelist(
 
     match allow_run_native {
         true => assert_matches!(
-            manager.get_runnable(&class_hash),
+            manager.get_runnable(&class_hash, &whitelist),
             Some(RunnableCompiledClass::V1Native(_))
         ),
         false => {
             assert_matches!(
-                manager.get_runnable(&class_hash).unwrap(),
+                manager.get_runnable(&class_hash, &whitelist).unwrap(),
                 RunnableCompiledClass::V1(_)
             )
         }

crates/blockifier/src/state/state_reader_and_contract_manager.rs

@@ -1,6 +1,7 @@
 use starknet_api::core::{ClassHash, CompiledClassHash};
 use starknet_types_core::felt::Felt;
 
+use crate::blockifier::config::NativeClassesWhitelist;
 use crate::execution::contract_class::RunnableCompiledClass;
 use crate::metrics::CacheMetrics;
 use crate::state::contract_class_manager::ContractClassManager;
@@ -33,6 +34,7 @@ impl<T: FetchCompiledClasses + ?Sized> FetchCompiledClasses for Box<T> {
 pub struct StateReaderAndContractManager<S: FetchCompiledClasses> {
     pub state_reader: S,
     contract_class_manager: ContractClassManager,
+    native_classes_whitelist: NativeClassesWhitelist,
     class_cache_metrics: Option<CacheMetrics>,
 }
 
@@ -42,7 +44,21 @@ impl<S: FetchCompiledClasses> StateReaderAndContractManager<S> {
         contract_class_manager: ContractClassManager,
         class_cache_metrics: Option<CacheMetrics>,
     ) -> Self {
-        Self { state_reader, contract_class_manager, class_cache_metrics }
+        Self::new_with_native_classes_whitelist(
+            state_reader,
+            contract_class_manager,
+            NativeClassesWhitelist::All,
+            class_cache_metrics,
+        )
+    }
+
+    pub fn new_with_native_classes_whitelist(
+        state_reader: S,
+        contract_class_manager: ContractClassManager,
+        native_classes_whitelist: NativeClassesWhitelist,
+        class_cache_metrics: Option<CacheMetrics>,
+    ) -> Self {
+        Self { state_reader, contract_class_manager, native_classes_whitelist, class_cache_metrics }
     }
 }
 
@@ -51,7 +67,9 @@ impl<S: FetchCompiledClasses> StateReaderAndContractManager<S> {
         &self,
         class_hash: ClassHash,
     ) -> StateResult<RunnableCompiledClass> {
-        if let Some(runnable_class) = self.contract_class_manager.get_runnable(&class_hash) {
+        if let Some(runnable_class) =
+            self.contract_class_manager.get_runnable(&class_hash, &self.native_classes_whitelist)
+        {
             match &runnable_class {
                 RunnableCompiledClass::V0(_) => {}
                 _ => {
@@ -72,8 +90,10 @@ impl<S: FetchCompiledClasses> StateReaderAndContractManager<S> {
         let compiled_class = self.state_reader.get_compiled_classes(class_hash)?;
         self.contract_class_manager.set_and_compile(class_hash, compiled_class.clone());
         // Access the cache again in case the class was compiled.
-        let runnable_class =
-            self.contract_class_manager.get_runnable(&class_hash).unwrap_or_else(|| {
+        let runnable_class = self
+            .contract_class_manager
+            .get_runnable(&class_hash, &self.native_classes_whitelist)
+            .unwrap_or_else(|| {
                 // Edge case that should not be happen if the cache size is big enough.
                 // TODO(Yoni): consider having an atomic set-and-get.
                 log::error!("Class is missing immediately after being cached.");

crates/blockifier/src/state/state_reader_and_contract_manager_test.rs

@@ -13,7 +13,7 @@ use starknet_api::core::ClassHash;
 
 #[cfg(not(feature = "cairo_native"))]
 use crate::blockifier::config::CairoNativeRunConfig;
-use crate::blockifier::config::ContractClassManagerConfig;
+use crate::blockifier::config::{ContractClassManagerConfig, NativeClassesWhitelist};
 use crate::execution::contract_class::RunnableCompiledClass;
 use crate::state::contract_class_manager::ContractClassManager;
 #[cfg(not(feature = "cairo_native"))]
@@ -80,7 +80,12 @@ fn test_get_compiled_class_without_native_in_cache(
         build_reader_and_declare_contract(test_contract.into(), contract_manager_config);
 
     // Sanity check - the class manager's cache is empty.
-    assert!(state_reader.contract_class_manager.get_runnable(&test_class_hash).is_none());
+    assert!(
+        state_reader
+            .contract_class_manager
+            .get_runnable(&test_class_hash, &NativeClassesWhitelist::All)
+            .is_none()
+    );
 
     let compiled_class = state_reader.get_compiled_class(test_class_hash).unwrap();