core_utils/preprocessing/
per_stream.rs

1//! Per-stream values addressed by [`PreprocessingKind`].
2
3use std::{
4    array,
5    ops::{Add, AddAssign, Index, IndexMut, Sub, SubAssign},
6};
7
8use primitives::correlated_randomness::stream::{
9    Buffer,
10    BufferConfig,
11    CorrelatedStreamError,
12    SharedBufferConfig,
13};
14use serde::{Deserialize, Serialize};
15
16use crate::circuit::FieldType;
17
18/// Number of preprocessing streams a bundler holds.
19pub const NUM_STREAMS: usize = 11;
20
21/// A field's correlations, in [`PreprocessingKind`] group order.
22#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
23#[repr(usize)]
24pub enum CorrelationKind {
25    Singlets,
26    Triples,
27    DaBits,
28}
29
30/// One stream of a preprocessing bundle: one (field, correlation) pair. Discriminants index
31/// [`ALL`](Self::ALL) and every `[T; NUM_STREAMS]` table below.
32#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
33#[repr(usize)]
34pub enum PreprocessingKind {
35    BitSinglets,
36    BitTriples,
37    BaseFieldSinglets,
38    BaseFieldTriples,
39    BaseFieldDaBits,
40    ScalarSinglets,
41    ScalarTriples,
42    ScalarDaBits,
43    MpcFieldSinglets,
44    MpcFieldTriples,
45    MpcFieldDaBits,
46}
47
48impl PreprocessingKind {
49    /// Index of the first field-group kind.
50    const FIELD_BASE: usize = 2;
51
52    /// Every kind, in discriminant order.
53    pub const ALL: [Self; NUM_STREAMS] = [
54        Self::BitSinglets,
55        Self::BitTriples,
56        Self::BaseFieldSinglets,
57        Self::BaseFieldTriples,
58        Self::BaseFieldDaBits,
59        Self::ScalarSinglets,
60        Self::ScalarTriples,
61        Self::ScalarDaBits,
62        Self::MpcFieldSinglets,
63        Self::MpcFieldTriples,
64        Self::MpcFieldDaBits,
65    ];
66
67    const LABELS: [&'static str; NUM_STREAMS] = [
68        "bit singlets",
69        "bit triples",
70        "base-field singlets",
71        "base-field triples",
72        "base-field daBits",
73        "scalar singlets",
74        "scalar triples",
75        "scalar daBits",
76        "MPC-field singlets",
77        "MPC-field triples",
78        "MPC-field daBits",
79    ];
80
81    /// The stream carrying `correlation` for `field`.
82    pub const fn of(field: FieldType, correlation: CorrelationKind) -> Self {
83        Self::ALL[Self::FIELD_BASE + 3 * field as usize + correlation as usize]
84    }
85
86    /// Human-readable name.
87    pub const fn label(self) -> &'static str {
88        Self::LABELS[self as usize]
89    }
90}
91
92impl std::fmt::Display for PreprocessingKind {
93    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
94        f.write_str(self.label())
95    }
96}
97
98/// One `T` per preprocessing stream, indexed by [`PreprocessingKind`].
99#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
100#[repr(transparent)]
101pub struct PerStream<T>(pub [T; NUM_STREAMS]);
102
103impl<T> PerStream<T> {
104    /// Builds each slot from its kind.
105    pub fn from_fn(mut f: impl FnMut(PreprocessingKind) -> T) -> Self {
106        Self(array::from_fn(|i| f(PreprocessingKind::ALL[i])))
107    }
108
109    /// The same value for every stream.
110    pub fn uniform(value: T) -> Self
111    where
112        T: Clone,
113    {
114        Self::from_fn(|_| value.clone())
115    }
116
117    /// Overrides one slot.
118    pub fn with(mut self, kind: PreprocessingKind, value: T) -> Self {
119        self[kind] = value;
120        self
121    }
122
123    /// Every slot, in [`PreprocessingKind::ALL`] order.
124    pub fn iter(&self) -> impl Iterator<Item = &T> {
125        self.0.iter()
126    }
127
128    /// Slot-wise maximum.
129    pub fn componentwise_max(self, other: Self) -> Self
130    where
131        T: Ord + Copy,
132    {
133        Self(array::from_fn(|i| self.0[i].max(other.0[i])))
134    }
135}
136
137impl<T> Index<PreprocessingKind> for PerStream<T> {
138    type Output = T;
139
140    fn index(&self, kind: PreprocessingKind) -> &T {
141        &self.0[kind as usize]
142    }
143}
144
145impl<T> IndexMut<PreprocessingKind> for PerStream<T> {
146    fn index_mut(&mut self, kind: PreprocessingKind) -> &mut T {
147        &mut self.0[kind as usize]
148    }
149}
150
151impl<T: AddAssign + Copy> AddAssign for PerStream<T> {
152    fn add_assign(&mut self, rhs: Self) {
153        self.0.iter_mut().zip(rhs.0).for_each(|(a, b)| *a += b);
154    }
155}
156
157impl<T: SubAssign + Copy> SubAssign for PerStream<T> {
158    fn sub_assign(&mut self, rhs: Self) {
159        self.0.iter_mut().zip(rhs.0).for_each(|(a, b)| *a -= b);
160    }
161}
162
163impl<T: AddAssign + Copy> Add for PerStream<T> {
164    type Output = Self;
165
166    fn add(mut self, rhs: Self) -> Self {
167        self += rhs;
168        self
169    }
170}
171
172impl<T: SubAssign + Copy> Sub for PerStream<T> {
173    type Output = Self;
174
175    fn sub(mut self, rhs: Self) -> Self {
176        self -= rhs;
177        self
178    }
179}
180
181/// A [`BufferConfig`] per stream, as taken by the bundler builders.
182pub type BufferConfigs = PerStream<BufferConfig>;
183
184impl From<BufferConfig> for BufferConfigs {
185    fn from(config: BufferConfig) -> Self {
186        Self::uniform(config)
187    }
188}
189
190/// Writes apply to every stream; reads return the first.
191impl Buffer for PerStream<SharedBufferConfig> {
192    fn config(&self) -> &SharedBufferConfig {
193        &self.0[0]
194    }
195    fn set_capacity(&self, capacity: usize) -> Result<(), CorrelatedStreamError> {
196        self.iter().try_for_each(|c| c.set_capacity(capacity))
197    }
198    fn set_max_request_size(&self, n: usize) -> Result<(), CorrelatedStreamError> {
199        self.iter().try_for_each(|c| c.set_max_request_size(n))
200    }
201    fn set_refill_threshold(&self, n: usize) -> Result<(), CorrelatedStreamError> {
202        self.iter().try_for_each(|c| c.set_refill_threshold(n))
203    }
204}
205
206#[cfg(test)]
207mod tests {
208    use super::*;
209
210    /// Discriminants must match positions in `ALL`.
211    #[test]
212    fn discriminants_index_all() {
213        for (i, kind) in PreprocessingKind::ALL.into_iter().enumerate() {
214            assert_eq!(kind as usize, i);
215            assert_eq!(PreprocessingKind::ALL[kind as usize], kind);
216        }
217    }
218
219    /// `of` agrees with the named variants.
220    #[test]
221    fn of_addresses_each_field_group() {
222        use CorrelationKind::*;
223        use FieldType::*;
224        let expected = [
225            (BaseField, Singlets, PreprocessingKind::BaseFieldSinglets),
226            (BaseField, Triples, PreprocessingKind::BaseFieldTriples),
227            (BaseField, DaBits, PreprocessingKind::BaseFieldDaBits),
228            (ScalarField, Singlets, PreprocessingKind::ScalarSinglets),
229            (ScalarField, Triples, PreprocessingKind::ScalarTriples),
230            (ScalarField, DaBits, PreprocessingKind::ScalarDaBits),
231            (MpcField, Singlets, PreprocessingKind::MpcFieldSinglets),
232            (MpcField, Triples, PreprocessingKind::MpcFieldTriples),
233            (MpcField, DaBits, PreprocessingKind::MpcFieldDaBits),
234        ];
235        for (field, correlation, kind) in expected {
236            assert_eq!(PreprocessingKind::of(field, correlation), kind);
237        }
238    }
239
240    #[test]
241    fn indexing_addresses_each_slot() {
242        let mut per_stream = PerStream::uniform(0usize);
243        for (n, kind) in PreprocessingKind::ALL.into_iter().enumerate() {
244            per_stream[kind] = n + 1;
245        }
246        assert_eq!(
247            per_stream.iter().copied().collect::<Vec<_>>(),
248            (1..=NUM_STREAMS).collect::<Vec<_>>()
249        );
250    }
251
252    #[test]
253    fn with_overrides_a_single_slot() {
254        let configs = BufferConfigs::from(BufferConfig::eager(16))
255            .with(PreprocessingKind::BitTriples, BufferConfig::eager(64));
256        assert_eq!(configs[PreprocessingKind::BitTriples].capacity(), 64);
257        assert_eq!(configs[PreprocessingKind::BitSinglets].capacity(), 16);
258    }
259}
core_utils/preprocessing/per_stream.rs

core_utils/preprocessing/
per_stream.rs
