arcium-primitives 0.6.0

use std::fmt::Debug;

use crate::{correlated_randomness::CorrelatedBatch, utils::TryFuture};

/// Generates preprocessing elements for a given preprocessing type.
pub trait CorrelationGenerator<PB: CorrelatedBatch>: Send {
    /// The interface to the network used to send/receive messages.
    type Net: Send;
    /// The error type returned by the generator.
    type Error: Debug + Clone + Send;

    /// Runs the preprocessing generation protocol and produces a batch of preprocessing elements
    /// with compile-time known size.
    fn run(&mut self, net: &mut Self::Net) -> impl TryFuture<Ok = PB, Error = Self::Error>;

    /// Runs the preprocessing generation protocol until at least `n_elements` are produced, and
    /// returns a vector of preprocessing elements.
    fn run_for(
        &mut self,
        n_elements: usize,
        net: &mut Self::Net,
    ) -> impl TryFuture<Ok = Vec<PB::Item>, Error = Self::Error> {
        let n_batches = n_elements.div_ceil(PB::batch_size());
        let mut result = Vec::with_capacity(n_batches * PB::batch_size());
        async move {
            for _ in 0..n_batches {
                let batch = self.run(net).await?;
                result.extend(batch.into_iter());
            }
            Ok(result)
        }
    }
}