refactor(block-producer): remove Store traits (#571)

crates/block-producer/src/block_builder/mod.rs

@@ -12,11 +12,8 @@ use tokio::time::Duration;
 use tracing::{debug, info, instrument};
 
 use crate::{
-    batch_builder::batch::TransactionBatch,
-    errors::BuildBlockError,
-    mempool::SharedMempool,
-    store::{ApplyBlock, DefaultStore, Store},
-    COMPONENT, SERVER_BLOCK_FREQUENCY,
+    batch_builder::batch::TransactionBatch, errors::BuildBlockError, mempool::SharedMempool,
+    store::StoreClient, COMPONENT, SERVER_BLOCK_FREQUENCY,
 };
 
 pub(crate) mod prover;
@@ -36,12 +33,12 @@ pub struct BlockBuilder {
     /// Note: this _must_ be sign positive and less than 1.0.
     pub failure_rate: f32,
 
-    pub store: DefaultStore,
+    pub store: StoreClient,
     pub block_kernel: BlockProver,
 }
 
 impl BlockBuilder {
-    pub fn new(store: DefaultStore) -> Self {
+    pub fn new(store: StoreClient) -> Self {
         Self {
             block_interval: SERVER_BLOCK_FREQUENCY,
             // Note: The range cannot be empty.

crates/block-producer/src/block_builder/prover/tests.rs

@@ -23,7 +23,6 @@ use super::*;
 use crate::{
     batch_builder::TransactionBatch,
     block::{AccountWitness, BlockInputs},
-    store::Store,
     test_utils::{
         block::{build_actual_block_header, build_expected_block_header, MockBlockBuilder},
         MockProvenTxBuilder, MockStoreSuccessBuilder,

crates/block-producer/src/server/mod.rs

@@ -21,7 +21,7 @@ use crate::{
     domain::transaction::AuthenticatedTransaction,
     errors::{AddTransactionError, BlockProducerError, VerifyTxError},
     mempool::{BatchBudget, BlockBudget, BlockNumber, Mempool, SharedMempool},
-    store::{DefaultStore, Store},
+    store::StoreClient,
     COMPONENT, SERVER_MEMPOOL_STATE_RETENTION,
 };
 
@@ -38,7 +38,7 @@ pub struct BlockProducer {
     block_budget: BlockBudget,
     state_retention: usize,
     rpc_listener: TcpListener,
-    store: DefaultStore,
+    store: StoreClient,
     chain_tip: BlockNumber,
 }
 
@@ -49,7 +49,7 @@ impl BlockProducer {
     pub async fn init(config: BlockProducerConfig) -> Result<Self, ApiError> {
         info!(target: COMPONENT, %config, "Initializing server");
 
-        let store = DefaultStore::new(
+        let store = StoreClient::new(
             store_client::ApiClient::connect(config.store_url.to_string())
                 .await
                 .map_err(|err| ApiError::DatabaseConnectionFailed(err.to_string()))?,
@@ -171,7 +171,7 @@ struct BlockProducerRpcServer {
     /// the block-producers usage of the mempool.
     mempool: Mutex<SharedMempool>,
 
-    store: DefaultStore,
+    store: StoreClient,
 }
 
 #[tonic::async_trait]
@@ -189,7 +189,7 @@ impl api_server::Api for BlockProducerRpcServer {
 }
 
 impl BlockProducerRpcServer {
-    pub fn new(mempool: SharedMempool, store: DefaultStore) -> Self {
+    pub fn new(mempool: SharedMempool, store: StoreClient) -> Self {
         Self { mempool: Mutex::new(mempool), store }
     }
 

crates/block-producer/src/store/mod.rs

@@ -4,7 +4,6 @@ use std::{
     num::NonZeroU32,
 };
 
-use async_trait::async_trait;
 use itertools::Itertools;
 use miden_node_proto::{
     errors::{ConversionError, MissingFieldHelper},
@@ -32,31 +31,6 @@ use tracing::{debug, info, instrument};
 pub use crate::errors::{ApplyBlockError, BlockInputsError, TxInputsError};
 use crate::{block::BlockInputs, COMPONENT};
 
-// STORE TRAIT
-// ================================================================================================
-
-#[async_trait]
-pub trait Store: ApplyBlock {
-    /// Returns information needed from the store to verify a given proven transaction.
-    async fn get_tx_inputs(
-        &self,
-        proven_tx: &ProvenTransaction,
-    ) -> Result<TransactionInputs, TxInputsError>;
-
-    /// Returns information needed from the store to build a block.
-    async fn get_block_inputs(
-        &self,
-        updated_accounts: impl Iterator<Item = AccountId> + Send,
-        produced_nullifiers: impl Iterator<Item = &Nullifier> + Send,
-        notes: impl Iterator<Item = &NoteId> + Send,
-    ) -> Result<BlockInputs, BlockInputsError>;
-}
-
-#[async_trait]
-pub trait ApplyBlock: Send + Sync + 'static {
-    async fn apply_block(&self, block: &Block) -> Result<(), ApplyBlockError>;
-}
-
 // TRANSACTION INPUTS
 // ================================================================================================
 
@@ -139,15 +113,18 @@ impl TryFrom<GetTransactionInputsResponse> for TransactionInputs {
     }
 }
 
-// DEFAULT STORE IMPLEMENTATION
+// STORE CLIENT
 // ================================================================================================
 
+/// Interface to the store's gRPC API.
+///
+/// Essentially just a thin wrapper around the generated gRPC client which improves type safety.
 #[derive(Clone)]
-pub struct DefaultStore {
+pub struct StoreClient {
     store: store_client::ApiClient<Channel>,
 }
 
-impl DefaultStore {
+impl StoreClient {
     /// TODO: this should probably take store connection string and create a connection internally
     pub fn new(store: store_client::ApiClient<Channel>) -> Self {
         Self { store }
@@ -166,29 +143,9 @@ impl DefaultStore {
 
         BlockHeader::try_from(response.block_header.unwrap()).map_err(|err| err.to_string())
     }
-}
 
-#[async_trait]
-impl ApplyBlock for DefaultStore {
     #[instrument(target = "miden-block-producer", skip_all, err)]
-    async fn apply_block(&self, block: &Block) -> Result<(), ApplyBlockError> {
-        let request = tonic::Request::new(ApplyBlockRequest { block: block.to_bytes() });
-
-        let _ = self
-            .store
-            .clone()
-            .apply_block(request)
-            .await
-            .map_err(|status| ApplyBlockError::GrpcClientError(status.message().to_string()))?;
-
-        Ok(())
-    }
-}
-
-#[async_trait]
-impl Store for DefaultStore {
-    #[instrument(target = "miden-block-producer", skip_all, err)]
-    async fn get_tx_inputs(
+    pub async fn get_tx_inputs(
         &self,
         proven_tx: &ProvenTransaction,
     ) -> Result<TransactionInputs, TxInputsError> {
@@ -230,7 +187,7 @@ impl Store for DefaultStore {
         Ok(tx_inputs)
     }
 
-    async fn get_block_inputs(
+    pub async fn get_block_inputs(
         &self,
         updated_accounts: impl Iterator<Item = AccountId> + Send,
         produced_nullifiers: impl Iterator<Item = &Nullifier> + Send,
@@ -250,6 +207,20 @@ impl Store for DefaultStore {
             .map_err(|err| BlockInputsError::GrpcClientError(err.message().to_string()))?
             .into_inner();
 
-        Ok(store_response.try_into()?)
+        store_response.try_into()
+    }
+
+    #[instrument(target = "miden-block-producer", skip_all, err)]
+    pub async fn apply_block(&self, block: &Block) -> Result<(), ApplyBlockError> {
+        let request = tonic::Request::new(ApplyBlockRequest { block: block.to_bytes() });
+
+        let _ = self
+            .store
+            .clone()
+            .apply_block(request)
+            .await
+            .map_err(|status| ApplyBlockError::GrpcClientError(status.message().to_string()))?;
+
+        Ok(())
     }
 }

crates/block-producer/src/test_utils/block.rs

@@ -13,7 +13,6 @@ use crate::{
     batch_builder::TransactionBatch,
     block::BlockInputs,
     block_builder::prover::{block_witness::BlockWitness, BlockProver},
-    store::Store,
 };
 
 /// Constructs the block we expect to be built given the store state, and a set of transaction

crates/block-producer/src/test_utils/mod.rs

@@ -13,7 +13,7 @@ pub use proven_tx::{mock_proven_tx, MockProvenTxBuilder};
 
 mod store;
 
-pub use store::{MockStoreFailure, MockStoreSuccess, MockStoreSuccessBuilder};
+pub use store::{MockStoreSuccess, MockStoreSuccessBuilder};
 
 mod account;
 

crates/block-producer/src/test_utils/store.rs

@@ -4,7 +4,6 @@ use std::{
     ops::Not,
 };
 
-use async_trait::async_trait;
 use miden_node_proto::domain::{blocks::BlockInclusionProof, notes::NoteAuthenticationInfo};
 use miden_objects::{
     block::{Block, NoteBatch},
@@ -19,9 +18,7 @@ use super::*;
 use crate::{
     batch_builder::TransactionBatch,
     block::{AccountWitness, BlockInputs},
-    store::{
-        ApplyBlock, ApplyBlockError, BlockInputsError, Store, TransactionInputs, TxInputsError,
-    },
+    store::{ApplyBlockError, BlockInputsError, TransactionInputs, TxInputsError},
     test_utils::block::{
         block_output_notes, flatten_output_notes, note_created_smt_from_note_batches,
     },
@@ -185,11 +182,8 @@ impl MockStoreSuccess {
 
         locked_accounts.root()
     }
-}
 
-#[async_trait]
-impl ApplyBlock for MockStoreSuccess {
-    async fn apply_block(&self, block: &Block) -> Result<(), ApplyBlockError> {
+    pub async fn apply_block(&self, block: &Block) -> Result<(), ApplyBlockError> {
         // Intentionally, we take and hold both locks, to prevent calls to `get_tx_inputs()` from
         // going through while we're updating the store's data structure
         let mut locked_accounts = self.accounts.write().await;
@@ -240,11 +234,8 @@ impl ApplyBlock for MockStoreSuccess {
 
         Ok(())
     }
-}
 
-#[async_trait]
-impl Store for MockStoreSuccess {
-    async fn get_tx_inputs(
+    pub async fn get_tx_inputs(
         &self,
         proven_tx: &ProvenTransaction,
     ) -> Result<TransactionInputs, TxInputsError> {
@@ -290,7 +281,7 @@ impl Store for MockStoreSuccess {
         })
     }
 
-    async fn get_block_inputs(
+    pub async fn get_block_inputs(
         &self,
         updated_accounts: impl Iterator<Item = AccountId> + Send,
         produced_nullifiers: impl Iterator<Item = &Nullifier> + Send,
@@ -352,32 +343,3 @@ impl Store for MockStoreSuccess {
         })
     }
 }
-
-#[derive(Default)]
-pub struct MockStoreFailure;
-
-#[async_trait]
-impl ApplyBlock for MockStoreFailure {
-    async fn apply_block(&self, _block: &Block) -> Result<(), ApplyBlockError> {
-        Err(ApplyBlockError::GrpcClientError(String::new()))
-    }
-}
-
-#[async_trait]
-impl Store for MockStoreFailure {
-    async fn get_tx_inputs(
-        &self,
-        _proven_tx: &ProvenTransaction,
-    ) -> Result<TransactionInputs, TxInputsError> {
-        Err(TxInputsError::Dummy)
-    }
-
-    async fn get_block_inputs(
-        &self,
-        _updated_accounts: impl Iterator<Item = AccountId> + Send,
-        _produced_nullifiers: impl Iterator<Item = &Nullifier> + Send,
-        _notes: impl Iterator<Item = &NoteId> + Send,
-    ) -> Result<BlockInputs, BlockInputsError> {
-        Err(BlockInputsError::GrpcClientError(String::new()))
-    }
-}