From a5bafd2e27f77dbe641ba01f72123528c44a33c4 Mon Sep 17 00:00:00 2001 From: Ayelet Zilber <138376632+ayeletstarkware@users.noreply.github.com> Date: Tue, 24 Dec 2024 18:02:26 +0200 Subject: [PATCH] refactor(starknet_batcher): delete ProposalOutput struct, use BlockExecutionArtifacts instead (#2917) --- crates/starknet_batcher/src/batcher.rs | 28 ++++++------ crates/starknet_batcher/src/batcher_test.rs | 14 +++--- crates/starknet_batcher/src/block_builder.rs | 40 +++++++++++++++- crates/starknet_batcher/src/utils.rs | 48 +------------------- 4 files changed, 62 insertions(+), 68 deletions(-) diff --git a/crates/starknet_batcher/src/batcher.rs b/crates/starknet_batcher/src/batcher.rs index 0aabda1a2a..c21ea9a6fb 100644 --- a/crates/starknet_batcher/src/batcher.rs +++ b/crates/starknet_batcher/src/batcher.rs @@ -43,6 +43,7 @@ use crate::block_builder::{ BlockBuilderFactory, BlockBuilderFactoryTrait, BlockBuilderTrait, + BlockExecutionArtifacts, BlockMetadata, }; use crate::config::BatcherConfig; @@ -51,7 +52,6 @@ use crate::utils::{ deadline_as_instant, proposal_status_from, verify_block_input, - ProposalOutput, ProposalResult, ProposalTask, }; @@ -81,7 +81,7 @@ pub struct Batcher { active_proposal_task: Option, // Holds all the proposals that completed execution in the current height. - executed_proposals: Arc>>>, + executed_proposals: Arc>>>, // The propose blocks transaction streams, used to stream out the proposal transactions. // Each stream is kept until all the transactions are streamed out, or a new height is started. @@ -408,13 +408,17 @@ impl Batcher { let proposal_id = input.proposal_id; let proposal_result = self.executed_proposals.lock().await.remove(&proposal_id); - let ProposalOutput { state_diff, nonces: address_to_nonce, tx_hashes, .. } = - proposal_result - .ok_or(BatcherError::ExecutedProposalNotFound { proposal_id })? - .map_err(|_| BatcherError::InternalError)?; - - self.commit_proposal_and_block(height, state_diff.clone(), address_to_nonce, tx_hashes) - .await?; + let block_execution_artifacts = proposal_result + .ok_or(BatcherError::ExecutedProposalNotFound { proposal_id })? + .map_err(|_| BatcherError::InternalError)?; + let state_diff = block_execution_artifacts.state_diff(); + self.commit_proposal_and_block( + height, + state_diff.clone(), + block_execution_artifacts.address_to_nonce(), + block_execution_artifacts.tx_hashes(), + ) + .await?; Ok(DecisionReachedResponse { state_diff }) } @@ -485,8 +489,7 @@ impl Batcher { let join_handle = tokio::spawn( async move { - let result = - block_builder.build_block().await.map(ProposalOutput::from).map_err(Arc::new); + let result = block_builder.build_block().await.map_err(Arc::new); // The proposal is done, clear the active proposal. // Keep the proposal result only if it is the same as the active proposal. @@ -512,9 +515,8 @@ impl Batcher { ) -> Option> { let guard = self.executed_proposals.lock().await; let proposal_result = guard.get(&proposal_id); - match proposal_result { - Some(Ok(output)) => Some(Ok(output.commitment)), + Some(Ok(artifacts)) => Some(Ok(artifacts.commitment())), Some(Err(e)) => Some(Err(e.clone())), None => None, } diff --git a/crates/starknet_batcher/src/batcher_test.rs b/crates/starknet_batcher/src/batcher_test.rs index 704f31f3ca..9ca17ce25d 100644 --- a/crates/starknet_batcher/src/batcher_test.rs +++ b/crates/starknet_batcher/src/batcher_test.rs @@ -45,7 +45,6 @@ use crate::block_builder::{ }; use crate::config::BatcherConfig; use crate::test_utils::{test_txs, FakeProposeBlockBuilder, FakeValidateBlockBuilder}; -use crate::utils::ProposalOutput; const INITIAL_HEIGHT: BlockNumber = BlockNumber(3); const STREAMING_CHUNK_SIZE: usize = 3; @@ -55,7 +54,7 @@ const BUILD_BLOCK_FAIL_ON_ERROR: BlockBuilderError = BlockBuilderError::FailOnError(FailOnErrorCause::BlockFull); fn proposal_commitment() -> ProposalCommitment { - ProposalOutput::from(BlockExecutionArtifacts::create_for_testing()).commitment + BlockExecutionArtifacts::create_for_testing().commitment() } fn propose_block_input(proposal_id: ProposalId) -> ProposeBlockInput { @@ -534,16 +533,15 @@ async fn add_sync_block_mismatch_block_number() { #[tokio::test] async fn decision_reached() { let mut mock_dependencies = MockDependencies::default(); - let expected_proposal_output = - ProposalOutput::from(BlockExecutionArtifacts::create_for_testing()); + let expected_artifacts = BlockExecutionArtifacts::create_for_testing(); mock_dependencies .mempool_client .expect_commit_block() .times(1) .with(eq(CommitBlockArgs { - address_to_nonce: expected_proposal_output.nonces, - tx_hashes: expected_proposal_output.tx_hashes, + address_to_nonce: expected_artifacts.address_to_nonce(), + tx_hashes: expected_artifacts.tx_hashes(), })) .returning(|_| Ok(())); @@ -551,7 +549,7 @@ async fn decision_reached() { .storage_writer .expect_commit_proposal() .times(1) - .with(eq(INITIAL_HEIGHT), eq(expected_proposal_output.state_diff.clone())) + .with(eq(INITIAL_HEIGHT), eq(expected_artifacts.state_diff())) .returning(|_, _| Ok(())); mock_create_builder_for_propose_block( @@ -567,7 +565,7 @@ async fn decision_reached() { let response = batcher.decision_reached(DecisionReachedInput { proposal_id: PROPOSAL_ID }).await.unwrap(); - assert_eq!(response.state_diff, expected_proposal_output.state_diff.clone()); + assert_eq!(response.state_diff, expected_artifacts.state_diff()); } #[rstest] diff --git a/crates/starknet_batcher/src/block_builder.rs b/crates/starknet_batcher/src/block_builder.rs index 7fc3595c08..cf3e9ec482 100644 --- a/crates/starknet_batcher/src/block_builder.rs +++ b/crates/starknet_batcher/src/block_builder.rs @@ -1,4 +1,4 @@ -use std::collections::BTreeMap; +use std::collections::{BTreeMap, HashMap, HashSet}; use async_trait::async_trait; use blockifier::blockifier::config::TransactionExecutorConfig; @@ -25,9 +25,13 @@ use papyrus_state_reader::papyrus_state::PapyrusReader; use papyrus_storage::StorageReader; use serde::{Deserialize, Serialize}; use starknet_api::block::{BlockHashAndNumber, BlockInfo}; +use starknet_api::block_hash::state_diff_hash::calculate_state_diff_hash; +use starknet_api::core::{ContractAddress, Nonce}; use starknet_api::executable_transaction::Transaction; use starknet_api::execution_resources::GasAmount; +use starknet_api::state::ThinStateDiff; use starknet_api::transaction::TransactionHash; +use starknet_batcher_types::batcher_types::ProposalCommitment; use thiserror::Error; use tracing::{debug, error, info, trace}; @@ -72,6 +76,40 @@ pub struct BlockExecutionArtifacts { pub l2_gas_used: GasAmount, } +impl BlockExecutionArtifacts { + pub fn address_to_nonce(&self) -> HashMap { + HashMap::from_iter( + self.commitment_state_diff + .address_to_nonce + .iter() + .map(|(address, nonce)| (*address, *nonce)), + ) + } + + pub fn tx_hashes(&self) -> HashSet { + HashSet::from_iter(self.execution_infos.keys().copied()) + } + + pub fn state_diff(&self) -> ThinStateDiff { + // TODO(Ayelet): Remove the clones. + let storage_diffs = self.commitment_state_diff.storage_updates.clone(); + let nonces = self.commitment_state_diff.address_to_nonce.clone(); + ThinStateDiff { + deployed_contracts: IndexMap::new(), + storage_diffs, + declared_classes: IndexMap::new(), + nonces, + // TODO: Remove this when the structure of storage diffs changes. + deprecated_declared_classes: Vec::new(), + replaced_classes: IndexMap::new(), + } + } + + pub fn commitment(&self) -> ProposalCommitment { + ProposalCommitment { state_diff_commitment: calculate_state_diff_hash(&self.state_diff()) } + } +} + /// The BlockBuilderTrait is responsible for building a new block from transactions provided by the /// tx_provider. The block building will stop at time deadline. /// The transactions that were added to the block will be streamed to the output_content_sender. diff --git a/crates/starknet_batcher/src/utils.rs b/crates/starknet_batcher/src/utils.rs index 4ea923a4db..275c13d726 100644 --- a/crates/starknet_batcher/src/utils.rs +++ b/crates/starknet_batcher/src/utils.rs @@ -1,18 +1,12 @@ -use std::collections::{HashMap, HashSet}; use std::sync::Arc; use blockifier::abi::constants; use chrono::Utc; -use indexmap::IndexMap; use starknet_api::block::{BlockHashAndNumber, BlockNumber}; -use starknet_api::block_hash::state_diff_hash::calculate_state_diff_hash; -use starknet_api::core::{ContractAddress, Nonce}; -use starknet_api::state::ThinStateDiff; -use starknet_api::transaction::TransactionHash; -use starknet_batcher_types::batcher_types::{BatcherResult, ProposalCommitment, ProposalStatus}; +use starknet_batcher_types::batcher_types::{BatcherResult, ProposalStatus}; use starknet_batcher_types::errors::BatcherError; -use crate::block_builder::{BlockBuilderError, BlockExecutionArtifacts}; +use crate::block_builder::BlockBuilderError; // BlockBuilderError is wrapped in an Arc since it doesn't implement Clone. pub(crate) type ProposalResult = Result>; @@ -23,44 +17,6 @@ pub(crate) struct ProposalTask { pub join_handle: tokio::task::JoinHandle<()>, } -#[derive(Debug, Default, PartialEq)] -pub(crate) struct ProposalOutput { - pub state_diff: ThinStateDiff, - pub commitment: ProposalCommitment, - pub tx_hashes: HashSet, - pub nonces: HashMap, -} - -impl From for ProposalOutput { - fn from(artifacts: BlockExecutionArtifacts) -> Self { - let commitment_state_diff = artifacts.commitment_state_diff; - let nonces = HashMap::from_iter( - commitment_state_diff - .address_to_nonce - .iter() - .map(|(address, nonce)| (*address, *nonce)), - ); - - // TODO: Get these from the transactions. - let deployed_contracts = IndexMap::new(); - let declared_classes = IndexMap::new(); - let state_diff = ThinStateDiff { - deployed_contracts, - storage_diffs: commitment_state_diff.storage_updates, - declared_classes, - nonces: commitment_state_diff.address_to_nonce, - // TODO: Remove this when the structure of storage diffs changes. - deprecated_declared_classes: Vec::new(), - replaced_classes: IndexMap::new(), - }; - let commitment = - ProposalCommitment { state_diff_commitment: calculate_state_diff_hash(&state_diff) }; - let tx_hashes = HashSet::from_iter(artifacts.execution_infos.keys().copied()); - - Self { state_diff, commitment, tx_hashes, nonces } - } -} - pub(crate) fn deadline_as_instant( deadline: chrono::DateTime, ) -> BatcherResult {