refactor(starknet_batcher): delete ProposalOutput struct, use BlockEx…

…ecutionArtifacts instead (#2917)

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<ProposalTask>,
 
     // Holds all the proposals that completed execution in the current height.
-    executed_proposals: Arc<Mutex<HashMap<ProposalId, ProposalResult<ProposalOutput>>>>,
+    executed_proposals: Arc<Mutex<HashMap<ProposalId, ProposalResult<BlockExecutionArtifacts>>>>,
 
     // 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<ProposalResult<ProposalCommitment>> {
         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,
         }

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,24 +533,23 @@ 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(()));
 
     mock_dependencies
         .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]

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<ContractAddress, Nonce> {
+        HashMap::from_iter(
+            self.commitment_state_diff
+                .address_to_nonce
+                .iter()
+                .map(|(address, nonce)| (*address, *nonce)),
+        )
+    }
+
+    pub fn tx_hashes(&self) -> HashSet<TransactionHash> {
+        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.

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<T> = Result<T, Arc<BlockBuilderError>>;
@@ -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<TransactionHash>,
-    pub nonces: HashMap<ContractAddress, Nonce>,
-}
-
-impl From<BlockExecutionArtifacts> 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<Utc>,
 ) -> BatcherResult<tokio::time::Instant> {