arcium-primitives 0.8.5

use std::fmt::Debug;

use futures::future::BoxFuture;

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;

    /// Whether one party may skip ahead independently, without all parties drawing identical
    /// amounts in lockstep: `true` for dealer-style generators (deterministic per index), `false`
    /// for interactive ones (where an asymmetric skip would desync the joint protocol).
    ///
    /// If set to `true` you **must** override [`skip`](Self::skip) — the default panics (it can't
    /// construct a `Self::Error`). When `false`, `skip` is never called.
    const SUPPORTS_UNILATERAL_SKIP: bool = false;

    /// 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>;

    /// Advances the logical position by `n_elements` without materializing the skipped elements,
    /// to burn correlated elements during resync. Only called when
    /// [`SUPPORTS_UNILATERAL_SKIP`](Self::SUPPORTS_UNILATERAL_SKIP) is `true`; the default panics.
    fn skip(
        &mut self,
        _n_elements: usize,
        _net: &mut Self::Net,
    ) -> impl TryFuture<Ok = (), Error = Self::Error> {
        async move {
            unreachable!(
                "skip() called on a generator without SUPPORTS_UNILATERAL_SKIP; the stream \
                 dispatcher must guard on the const before ever calling skip()"
            )
        }
    }

    /// 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)
        }
    }
}

/// A generator whose orders can be dispatched without awaiting their results: [`issue_for`]
/// synchronously puts the order in flight (e.g. writes the request to the wire) and returns a
/// detached future for its outcome, so several orders may be pending at once.
///
/// Implementations must resolve results in issue (FIFO) order and produce elements from a single
/// logically-ordered sequence — order `k+1`'s elements follow order `k`'s regardless of when the
/// futures resolve. Interactive generators (all parties drawing in lockstep) cannot uphold this
/// independently and should not implement this trait.
///
/// [`issue_for`]: Self::issue_for
pub trait PipelinedCorrelationGenerator<PB: CorrelatedBatch>: CorrelationGenerator<PB> {
    /// Dispatches an order for at least `n_elements` and returns a future resolving to them.
    ///
    /// The returned future is `'static`, so it cannot borrow `net`: implementors may only use it
    /// *during* the call (to write the order out) and must capture a cloneable/`Arc` handle if the
    /// pending part of the order needs the network too. Implementors whose ordering path needs no
    /// network at all (e.g. a client whose event loop owns the link) take `Net = ()`.
    fn issue_for(
        &mut self,
        n_elements: usize,
        net: &mut Self::Net,
    ) -> BoxFuture<'static, Result<Vec<PB::Item>, Self::Error>>;
}