arcium-core-utils 0.8.5

use primitives::{
    algebra::field::binary::Gf2_128,
    correlated_randomness::{
        bundler::{errors::BundlerError, Bundler},
        dabits::DaBit,
        singlets::Singlet,
        stream::{
            Buffer,
            CorrelatedStream,
            CorrelatedStreamError,
            PrefetchHandle,
            ResyncHandle,
            SharedBufferConfig,
        },
        triples::Triple,
        Correlation,
    },
};

use crate::{
    circuit::preprocessing::CircuitPreprocessing,
    config::{BaseFieldOf, MpcConfig, MpcFieldOf, ScalarFieldOf},
    errors::AbortError,
    preprocessing::{iterator::PreprocessingIterator, PerStream, PreprocessingKind},
};

// ── Fan-out over the bundler's 11 streams ───────────────────────────────────────────────────── //

/// Fetches all 11 streams into a `PreprocessingIterator`.
macro_rules! build_iterator {
    ($self:ident, $req:ident, $(($field:ident, $stream:ident, $kind:ident)),+ $(,)?) => {
        PreprocessingIterator {
            $( $field: $self.$stream.next_n($req[PreprocessingKind::$kind])?.into_iter() ),+
        }
    };
}

/// Applies a method call to the stream serving `kind`, e.g. `with_stream!(self, kind, position())`.
/// The call's result type must not vary with the stream.
macro_rules! with_stream {
    ($self:ident, $kind:expr, $($call:tt)+) => {
        match $kind {
            PreprocessingKind::BitSinglets       => $self.binary_singlet_stream.$($call)+,
            PreprocessingKind::BitTriples        => $self.binary_triple_stream.$($call)+,
            PreprocessingKind::BaseFieldSinglets => $self.basefield_singlet_stream.$($call)+,
            PreprocessingKind::BaseFieldTriples  => $self.basefield_triple_stream.$($call)+,
            PreprocessingKind::BaseFieldDaBits   => $self.basefield_dabit_stream.$($call)+,
            PreprocessingKind::ScalarSinglets    => $self.scalar_singlet_stream.$($call)+,
            PreprocessingKind::ScalarTriples     => $self.scalar_triple_stream.$($call)+,
            PreprocessingKind::ScalarDaBits      => $self.scalar_dabit_stream.$($call)+,
            PreprocessingKind::MpcFieldSinglets  => $self.mpc_field_singlet_stream.$($call)+,
            PreprocessingKind::MpcFieldTriples   => $self.mpc_field_triple_stream.$($call)+,
            PreprocessingKind::MpcFieldDaBits    => $self.mpc_field_dabit_stream.$($call)+,
        }
    };
}

/// A stream that can fill a [`StreamBundler`] slot.
pub trait BundledStream<P: Correlation>: CorrelatedStream<P, Error = AbortError> + Buffer {}

impl<P: Correlation, T: CorrelatedStream<P, Error = AbortError> + Buffer> BundledStream<P> for T {}

/// One type-erased bundler slot.
type Slot<P> = Box<dyn BundledStream<P>>;

/// Stream bundler, holding one stream per preprocessing type, to provide all preprocessing for a
/// circuit via its streams.
pub struct StreamBundler<C: MpcConfig> {
    // Base field
    pub basefield_dabit_stream: Slot<DaBit<BaseFieldOf<C>>>,
    pub basefield_singlet_stream: Slot<Singlet<BaseFieldOf<C>>>,
    pub basefield_triple_stream: Slot<Triple<BaseFieldOf<C>>>,
    // Binary (Gf2_128)
    pub binary_singlet_stream: Slot<Singlet<Gf2_128>>,
    pub binary_triple_stream: Slot<Triple<Gf2_128>>,
    // MPC field
    pub mpc_field_dabit_stream: Slot<DaBit<MpcFieldOf<C>>>,
    pub mpc_field_singlet_stream: Slot<Singlet<MpcFieldOf<C>>>,
    pub mpc_field_triple_stream: Slot<Triple<MpcFieldOf<C>>>,
    // Scalar
    pub scalar_dabit_stream: Slot<DaBit<ScalarFieldOf<C>>>,
    pub scalar_singlet_stream: Slot<Singlet<ScalarFieldOf<C>>>,
    pub scalar_triple_stream: Slot<Triple<ScalarFieldOf<C>>>,
}

impl<C: MpcConfig> StreamBundler<C> {
    /// Creates a new bundler from the given streams.
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        basefield_dabit_stream: impl BundledStream<DaBit<BaseFieldOf<C>>> + 'static,
        basefield_singlet_stream: impl BundledStream<Singlet<BaseFieldOf<C>>> + 'static,
        basefield_triple_stream: impl BundledStream<Triple<BaseFieldOf<C>>> + 'static,
        binary_singlet_stream: impl BundledStream<Singlet<Gf2_128>> + 'static,
        binary_triple_stream: impl BundledStream<Triple<Gf2_128>> + 'static,
        mpc_field_dabit_stream: impl BundledStream<DaBit<MpcFieldOf<C>>> + 'static,
        mpc_field_singlet_stream: impl BundledStream<Singlet<MpcFieldOf<C>>> + 'static,
        mpc_field_triple_stream: impl BundledStream<Triple<MpcFieldOf<C>>> + 'static,
        scalar_dabit_stream: impl BundledStream<DaBit<ScalarFieldOf<C>>> + 'static,
        scalar_singlet_stream: impl BundledStream<Singlet<ScalarFieldOf<C>>> + 'static,
        scalar_triple_stream: impl BundledStream<Triple<ScalarFieldOf<C>>> + 'static,
    ) -> Self {
        Self {
            basefield_dabit_stream: Box::new(basefield_dabit_stream),
            basefield_singlet_stream: Box::new(basefield_singlet_stream),
            basefield_triple_stream: Box::new(basefield_triple_stream),
            binary_singlet_stream: Box::new(binary_singlet_stream),
            binary_triple_stream: Box::new(binary_triple_stream),
            mpc_field_dabit_stream: Box::new(mpc_field_dabit_stream),
            mpc_field_singlet_stream: Box::new(mpc_field_singlet_stream),
            mpc_field_triple_stream: Box::new(mpc_field_triple_stream),
            scalar_dabit_stream: Box::new(scalar_dabit_stream),
            scalar_singlet_stream: Box::new(scalar_singlet_stream),
            scalar_triple_stream: Box::new(scalar_triple_stream),
        }
    }

    /// Elements delivered by every stream, per type.
    pub fn positions(&self) -> CircuitPreprocessing {
        let mut pos = CircuitPreprocessing::default();
        for kind in PreprocessingKind::ALL {
            pos[kind] = with_stream!(self, kind, position()) as usize;
        }
        pos
    }

    /// Elements currently buffered in every stream, per type.
    pub fn buffered(&self) -> CircuitPreprocessing {
        let mut buf = CircuitPreprocessing::default();
        for kind in PreprocessingKind::ALL {
            buf[kind] = with_stream!(self, kind, buffered()) as usize;
        }
        buf
    }

    /// Advances every stream to its per-type `target`. Rejects the whole call if any target is
    /// behind its stream's position; otherwise dispatches all resyncs and joins them.
    pub async fn resync(&self, targets: &CircuitPreprocessing) -> Result<(), AbortError> {
        let current = self.positions();
        for kind in PreprocessingKind::ALL {
            if targets[kind] < current[kind] {
                return Err(CorrelatedStreamError::ResyncRewind {
                    current: current[kind] as u64,
                    target: targets[kind] as u64,
                }
                .into());
            }
        }
        ResyncHandle::join_all(
            PreprocessingKind::ALL
                .map(|kind| with_stream!(self, kind, resync(targets[kind] as u64))),
        )
        .await
    }

    /// The shared buffer config of every stream; itself a [`Buffer`] over the whole bundle.
    pub fn buffer_configs(&self) -> PerStream<SharedBufferConfig> {
        PerStream::from_fn(|kind| with_stream!(self, kind, config().clone()))
    }
}

// ──────────────────────── PreprocessingBundler impl ──────────────────────── //

impl<C: MpcConfig> Bundler for StreamBundler<C> {
    type Iterator = PreprocessingIterator<C>;

    fn fetch(
        &mut self,
        req: &CircuitPreprocessing,
    ) -> Result<PreprocessingIterator<C>, BundlerError> {
        Ok(build_iterator!(
            self,
            req,
            (binary_singlets, binary_singlet_stream, BitSinglets),
            (binary_triples, binary_triple_stream, BitTriples),
            (
                base_field_singlets,
                basefield_singlet_stream,
                BaseFieldSinglets
            ),
            (
                base_field_triples,
                basefield_triple_stream,
                BaseFieldTriples
            ),
            (base_field_dabits, basefield_dabit_stream, BaseFieldDaBits),
            (scalar_singlets, scalar_singlet_stream, ScalarSinglets),
            (scalar_triples, scalar_triple_stream, ScalarTriples),
            (scalar_dabits, scalar_dabit_stream, ScalarDaBits),
            (
                mpc_field_singlets,
                mpc_field_singlet_stream,
                MpcFieldSinglets
            ),
            (mpc_field_triples, mpc_field_triple_stream, MpcFieldTriples),
            (mpc_field_dabits, mpc_field_dabit_stream, MpcFieldDaBits),
        ))
    }

    /// Prefetches into every stream concurrently; resolves once all complete.
    fn prefetch(&self, req: &CircuitPreprocessing) -> PrefetchHandle<AbortError> {
        PrefetchHandle::join_all(
            PreprocessingKind::ALL.map(|kind| with_stream!(self, kind, prefetch_n(req[kind]))),
        )
    }
}