arcium-primitives 0.8.5

use std::{
    fmt::Debug,
    ops::{Add, AddAssign, Sub, SubAssign},
};

use crate::correlated_randomness::{
    bundler::errors::BundlerError,
    stream::{CorrelatedStreamError, PrefetchHandle},
};

pub mod errors;

/// The per-type correlation counts a [`BundleIterator`] holds (or that a fetch/prefetch requests).
pub type SizeOf<I> = <I as BundleIterator>::Size;

/// The error type of a [`BundleIterator`]'s correlation futures.
pub type ErrorOf<I> = <I as BundleIterator>::Error;

/// Bundles together multiple correlation generators/streams to fetch all correlations
///  required by a given consumer, e.g., a circuit, a protocol, a layer...
pub trait Bundler {
    /// A type to hold the futures to all requested correlations.
    type Iterator: BundleIterator;

    /// Fetches the requested correlations and returns a ready-to-use iterator.
    fn fetch(&mut self, size: &SizeOf<Self::Iterator>) -> Result<Self::Iterator, BundlerError>;

    /// Fetches the requested correlations and returns a ready-to-use iterator.
    fn fetch_for<Consumer: BundleConsumer<Iterator = Self::Iterator>>(
        &mut self,
        consumer: &Consumer,
    ) -> Result<Self::Iterator, BundlerError> {
        let size = consumer.required_preprocessing();
        self.fetch(&size)
    }

    /// Proactively generates the correlations for `size` so a later [`fetch`](Self::fetch) need not
    /// wait on generation. Fire-and-forget, and delivers nothing, so positions do not advance.
    /// `Ok` means the generation happened or was unnecessary.
    ///
    /// All-or-nothing per stream: a `size` beyond a stream's buffer resolves with
    /// [`RequestTooLarge`](CorrelatedStreamError::RequestTooLarge) or
    /// [`RateLimitExceeded`](CorrelatedStreamError::RateLimitExceeded).
    fn prefetch(&self, size: &SizeOf<Self::Iterator>) -> PrefetchHandle<ErrorOf<Self::Iterator>>;

    /// Prefetches everything `consumer` will require. See [`prefetch`](Self::prefetch).
    fn prefetch_for<Consumer: BundleConsumer<Iterator = Self::Iterator>>(
        &self,
        consumer: &Consumer,
    ) -> PrefetchHandle<ErrorOf<Self::Iterator>> {
        let size = consumer.required_preprocessing();
        self.prefetch(&size)
    }
}

impl<B: Bundler> Bundler for &mut B {
    type Iterator = B::Iterator;

    fn fetch(&mut self, size: &SizeOf<Self::Iterator>) -> Result<Self::Iterator, BundlerError> {
        (**self).fetch(size)
    }

    fn prefetch(&self, size: &SizeOf<Self::Iterator>) -> PrefetchHandle<ErrorOf<Self::Iterator>> {
        (**self).prefetch(size)
    }
}

/// A consumer of a bundle of correlations, e.g., a circuit, a protocol, a layer...
pub trait BundleConsumer {
    /// An iterator holding futures to all correlations required by the consumer
    ///  (e.g., F_q triples, binary singlets, ...)
    type Iterator: BundleIterator;

    /// Returns the size of the preprocessing required by the consumer.
    fn required_preprocessing(&self) -> SizeOf<Self::Iterator>;

    /// Fetches all preprocessing required by `consumer` and returns a ready-to-use iterator.
    fn fetch_preprocessing_from<PB: Bundler<Iterator = Self::Iterator>>(
        &self,
        bundler: &mut PB,
    ) -> Result<Self::Iterator, BundlerError> {
        let size = self.required_preprocessing();
        bundler.fetch(&size)
    }

    /// Prefetches all preprocessing this consumer will require. See [`Bundler::prefetch`].
    fn prefetch_preprocessing_from<PB: Bundler<Iterator = Self::Iterator>>(
        &self,
        bundler: &PB,
    ) -> PrefetchHandle<ErrorOf<Self::Iterator>> {
        let size = self.required_preprocessing();
        bundler.prefetch(&size)
    }
}

/// An iterator that holds futures to all correlations required by a consumer.
pub trait BundleIterator {
    /// The size of the iterator, i.e., the number of correlations per type that it holds.
    type Size: Debug + Clone + Default + Eq + Add + Sub + AddAssign + SubAssign;

    /// The error type of the correlation futures held by the iterator.
    type Error: Debug + Clone + From<CorrelatedStreamError> + Send + 'static;

    /// Returns the size of the iterator, i.e., the number of correlations per type that it holds.
    fn len(&self) -> Self::Size;

    /// Returns `true` if the iterator is empty, i.e., it holds no correlations.
    fn is_empty(&self) -> bool {
        self.len() == Self::Size::default()
    }
}

#[cfg(test)]
mod tests {
    use std::cell::Cell;

    use super::*;

    /// Minimal bundle whose `Size` is a plain count.
    struct Bundle(usize);

    impl BundleIterator for Bundle {
        type Size = usize;
        type Error = CorrelatedStreamError;

        fn len(&self) -> usize {
            self.0
        }
    }

    /// Records what it was asked to prefetch.
    #[derive(Default)]
    struct Recorder(Cell<Option<usize>>);

    impl Bundler for Recorder {
        type Iterator = Bundle;

        fn fetch(&mut self, size: &usize) -> Result<Bundle, BundlerError> {
            Ok(Bundle(*size))
        }

        fn prefetch(&self, size: &usize) -> PrefetchHandle<CorrelatedStreamError> {
            self.0.set(Some(*size));
            PrefetchHandle::ready(Ok(()))
        }
    }

    struct Consumer(usize);

    impl BundleConsumer for Consumer {
        type Iterator = Bundle;

        fn required_preprocessing(&self) -> usize {
            self.0
        }
    }

    /// Both consumer-driven entry points forward the consumer's required size to the bundler.
    #[tokio::test]
    async fn consumer_driven_prefetch_forwards_the_required_size() {
        let bundler = Recorder::default();
        bundler.prefetch_for(&Consumer(11)).await.unwrap();
        assert_eq!(bundler.0.get(), Some(11));

        let bundler = Recorder::default();
        Consumer(23)
            .prefetch_preprocessing_from(&bundler)
            .await
            .unwrap();
        assert_eq!(bundler.0.get(), Some(23));
    }

    /// `&mut B` forwards `prefetch` to `B`.
    #[tokio::test]
    async fn mut_ref_forwards_prefetch() {
        let mut bundler = Recorder::default();
        let forwarded = &mut bundler;
        forwarded.prefetch(&5).await.unwrap();
        assert_eq!(bundler.0.get(), Some(5));
    }
}