arcium-core-utils 0.8.6

//! Per-stream values addressed by [`PreprocessingKind`].

use std::{
    array,
    ops::{Add, AddAssign, Index, IndexMut, Sub, SubAssign},
};

use primitives::correlated_randomness::stream::{
    Buffer,
    BufferConfig,
    CorrelatedStreamError,
    SharedBufferConfig,
};
use serde::{Deserialize, Serialize};

use crate::circuit::FieldType;

/// Number of preprocessing streams a bundler holds.
pub const NUM_STREAMS: usize = 11;

/// A field's correlations, in [`PreprocessingKind`] group order.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[repr(usize)]
pub enum CorrelationKind {
    Singlets,
    Triples,
    DaBits,
}

/// One stream of a preprocessing bundle: one (field, correlation) pair. Discriminants index
/// [`ALL`](Self::ALL) and every `[T; NUM_STREAMS]` table below.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[repr(usize)]
pub enum PreprocessingKind {
    BitSinglets,
    BitTriples,
    BaseFieldSinglets,
    BaseFieldTriples,
    BaseFieldDaBits,
    ScalarSinglets,
    ScalarTriples,
    ScalarDaBits,
    MpcFieldSinglets,
    MpcFieldTriples,
    MpcFieldDaBits,
}

impl PreprocessingKind {
    /// Index of the first field-group kind.
    const FIELD_BASE: usize = 2;

    /// Every kind, in discriminant order.
    pub const ALL: [Self; NUM_STREAMS] = [
        Self::BitSinglets,
        Self::BitTriples,
        Self::BaseFieldSinglets,
        Self::BaseFieldTriples,
        Self::BaseFieldDaBits,
        Self::ScalarSinglets,
        Self::ScalarTriples,
        Self::ScalarDaBits,
        Self::MpcFieldSinglets,
        Self::MpcFieldTriples,
        Self::MpcFieldDaBits,
    ];

    const LABELS: [&'static str; NUM_STREAMS] = [
        "bit singlets",
        "bit triples",
        "base-field singlets",
        "base-field triples",
        "base-field daBits",
        "scalar singlets",
        "scalar triples",
        "scalar daBits",
        "MPC-field singlets",
        "MPC-field triples",
        "MPC-field daBits",
    ];

    /// The stream carrying `correlation` for `field`.
    pub const fn of(field: FieldType, correlation: CorrelationKind) -> Self {
        Self::ALL[Self::FIELD_BASE + 3 * field as usize + correlation as usize]
    }

    /// Human-readable name.
    pub const fn label(self) -> &'static str {
        Self::LABELS[self as usize]
    }
}

impl std::fmt::Display for PreprocessingKind {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(self.label())
    }
}

/// One `T` per preprocessing stream, indexed by [`PreprocessingKind`].
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[repr(transparent)]
pub struct PerStream<T>(pub [T; NUM_STREAMS]);

impl<T> PerStream<T> {
    /// Builds each slot from its kind.
    pub fn from_fn(mut f: impl FnMut(PreprocessingKind) -> T) -> Self {
        Self(array::from_fn(|i| f(PreprocessingKind::ALL[i])))
    }

    /// The same value for every stream.
    pub fn uniform(value: T) -> Self
    where
        T: Clone,
    {
        Self::from_fn(|_| value.clone())
    }

    /// Overrides one slot.
    pub fn with(mut self, kind: PreprocessingKind, value: T) -> Self {
        self[kind] = value;
        self
    }

    /// Every slot, in [`PreprocessingKind::ALL`] order.
    pub fn iter(&self) -> impl Iterator<Item = &T> {
        self.0.iter()
    }

    /// Slot-wise maximum.
    pub fn componentwise_max(self, other: Self) -> Self
    where
        T: Ord + Copy,
    {
        Self(array::from_fn(|i| self.0[i].max(other.0[i])))
    }
}

impl<T> Index<PreprocessingKind> for PerStream<T> {
    type Output = T;

    fn index(&self, kind: PreprocessingKind) -> &T {
        &self.0[kind as usize]
    }
}

impl<T> IndexMut<PreprocessingKind> for PerStream<T> {
    fn index_mut(&mut self, kind: PreprocessingKind) -> &mut T {
        &mut self.0[kind as usize]
    }
}

impl<T: AddAssign + Copy> AddAssign for PerStream<T> {
    fn add_assign(&mut self, rhs: Self) {
        self.0.iter_mut().zip(rhs.0).for_each(|(a, b)| *a += b);
    }
}

impl<T: SubAssign + Copy> SubAssign for PerStream<T> {
    fn sub_assign(&mut self, rhs: Self) {
        self.0.iter_mut().zip(rhs.0).for_each(|(a, b)| *a -= b);
    }
}

impl<T: AddAssign + Copy> Add for PerStream<T> {
    type Output = Self;

    fn add(mut self, rhs: Self) -> Self {
        self += rhs;
        self
    }
}

impl<T: SubAssign + Copy> Sub for PerStream<T> {
    type Output = Self;

    fn sub(mut self, rhs: Self) -> Self {
        self -= rhs;
        self
    }
}

/// A [`BufferConfig`] per stream, as taken by the bundler builders.
pub type BufferConfigs = PerStream<BufferConfig>;

impl From<BufferConfig> for BufferConfigs {
    fn from(config: BufferConfig) -> Self {
        Self::uniform(config)
    }
}

/// Writes apply to every stream; reads return the first.
impl Buffer for PerStream<SharedBufferConfig> {
    fn config(&self) -> &SharedBufferConfig {
        &self.0[0]
    }
    fn set_capacity(&self, capacity: usize) -> Result<(), CorrelatedStreamError> {
        self.iter().try_for_each(|c| c.set_capacity(capacity))
    }
    fn set_max_request_size(&self, n: usize) -> Result<(), CorrelatedStreamError> {
        self.iter().try_for_each(|c| c.set_max_request_size(n))
    }
    fn set_refill_threshold(&self, n: usize) -> Result<(), CorrelatedStreamError> {
        self.iter().try_for_each(|c| c.set_refill_threshold(n))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Discriminants must match positions in `ALL`.
    #[test]
    fn discriminants_index_all() {
        for (i, kind) in PreprocessingKind::ALL.into_iter().enumerate() {
            assert_eq!(kind as usize, i);
            assert_eq!(PreprocessingKind::ALL[kind as usize], kind);
        }
    }

    /// `of` agrees with the named variants.
    #[test]
    fn of_addresses_each_field_group() {
        use CorrelationKind::*;
        use FieldType::*;
        let expected = [
            (BaseField, Singlets, PreprocessingKind::BaseFieldSinglets),
            (BaseField, Triples, PreprocessingKind::BaseFieldTriples),
            (BaseField, DaBits, PreprocessingKind::BaseFieldDaBits),
            (ScalarField, Singlets, PreprocessingKind::ScalarSinglets),
            (ScalarField, Triples, PreprocessingKind::ScalarTriples),
            (ScalarField, DaBits, PreprocessingKind::ScalarDaBits),
            (MpcField, Singlets, PreprocessingKind::MpcFieldSinglets),
            (MpcField, Triples, PreprocessingKind::MpcFieldTriples),
            (MpcField, DaBits, PreprocessingKind::MpcFieldDaBits),
        ];
        for (field, correlation, kind) in expected {
            assert_eq!(PreprocessingKind::of(field, correlation), kind);
        }
    }

    #[test]
    fn indexing_addresses_each_slot() {
        let mut per_stream = PerStream::uniform(0usize);
        for (n, kind) in PreprocessingKind::ALL.into_iter().enumerate() {
            per_stream[kind] = n + 1;
        }
        assert_eq!(
            per_stream.iter().copied().collect::<Vec<_>>(),
            (1..=NUM_STREAMS).collect::<Vec<_>>()
        );
    }

    #[test]
    fn with_overrides_a_single_slot() {
        let configs = BufferConfigs::from(BufferConfig::eager(16))
            .with(PreprocessingKind::BitTriples, BufferConfig::eager(64));
        assert_eq!(configs[PreprocessingKind::BitTriples].capacity(), 64);
        assert_eq!(configs[PreprocessingKind::BitSinglets].capacity(), 16);
    }
}