mithril_aggregator/services/
prover.rs

1use async_trait::async_trait;
2use rayon::prelude::*;
3use slog::{Logger, debug, info};
4use std::sync::Arc;
5
6use mithril_common::{
7    StdResult,
8    crypto_helper::{MKMap, MKMapNode, MKTreeStorer},
9    entities::{BlockNumber, BlockRange},
10    logging::LoggerExtensions,
11    signable_builder::BlockRangeRootRetriever,
12};
13use mithril_resource_pool::ResourcePool;
14
15use crate::services::TransactionsRetriever;
16
17/// Prover service is the cryptographic engine in charge of producing cryptographic proofs for transactions and blocks
18// TODO: finalize the trait implementation
19#[cfg_attr(test, mockall::automock)]
20#[async_trait]
21pub trait ProverService: Sync + Send {
22    /// Compute the cache
23    async fn compute_cache(&self, up_to: BlockNumber) -> StdResult<()>;
24}
25
26/// Mithril prover
27pub struct MithrilProverService<S: MKTreeStorer> {
28    _transaction_retriever: Arc<dyn TransactionsRetriever>,
29    block_range_root_retriever: Arc<dyn BlockRangeRootRetriever<S>>,
30    mk_map_pool: ResourcePool<MKMap<BlockRange, MKMapNode<BlockRange, S>, S>>,
31    logger: Logger,
32}
33
34impl<S: MKTreeStorer> MithrilProverService<S> {
35    /// Create a new Mithril prover
36    pub fn new(
37        _transaction_retriever: Arc<dyn TransactionsRetriever>,
38        block_range_root_retriever: Arc<dyn BlockRangeRootRetriever<S>>,
39        mk_map_pool_size: usize,
40        logger: Logger,
41    ) -> Self {
42        Self {
43            _transaction_retriever,
44            block_range_root_retriever,
45            mk_map_pool: ResourcePool::new(mk_map_pool_size, vec![]),
46            logger: logger.new_with_component_name::<Self>(),
47        }
48    }
49}
50
51#[async_trait]
52impl<S: MKTreeStorer> ProverService for MithrilProverService<S> {
53    async fn compute_cache(&self, up_to: BlockNumber) -> StdResult<()> {
54        let pool_size = self.mk_map_pool.size();
55        info!(
56            self.logger, "Starts computing the Merkle map pool resource of size {pool_size}";
57            "up_to_block_number" => *up_to,
58        );
59        let mk_map_cache = self
60            .block_range_root_retriever
61            .compute_merkle_map_from_block_range_roots(up_to)
62            .await?;
63        let mk_maps_new = (1..=pool_size)
64            .into_par_iter()
65            .map(|i| {
66                debug!(
67                    self.logger,
68                    "Computing the Merkle map pool resource {i}/{pool_size}"
69                );
70                mk_map_cache.clone()
71            })
72            .collect::<Vec<MKMap<_, _, _>>>();
73        debug!(self.logger, "Draining the Merkle map pool");
74        let discriminant_new = self.mk_map_pool.discriminant()? + 1;
75        self.mk_map_pool.set_discriminant(discriminant_new)?;
76        self.mk_map_pool.clear();
77        debug!(
78            self.logger,
79            "Giving back new resources to the Merkle map pool"
80        );
81        mk_maps_new
82            .into_iter()
83            .map(|mk_map| self.mk_map_pool.give_back_resource(mk_map, discriminant_new))
84            .collect::<StdResult<Vec<_>>>()?;
85        info!(
86            self.logger,
87            "Completed computing the Merkle map pool resource of size {pool_size}"
88        );
89
90        Ok(())
91    }
92}