core_utils/preprocessing/
bundler.rs

1use primitives::{
2    algebra::field::binary::Gf2_128,
3    correlated_randomness::{
4        bundler::{errors::BundlerError, Bundler},
5        dabits::DaBit,
6        singlets::Singlet,
7        stream::{
8            Buffer,
9            CorrelatedStream,
10            CorrelatedStreamError,
11            PrefetchHandle,
12            ResyncHandle,
13            SharedBufferConfig,
14        },
15        triples::Triple,
16        Correlation,
17    },
18};
19
20use crate::{
21    circuit::preprocessing::CircuitPreprocessing,
22    config::{BaseFieldOf, MpcConfig, MpcFieldOf, ScalarFieldOf},
23    errors::AbortError,
24    preprocessing::{iterator::PreprocessingIterator, PerStream, PreprocessingKind},
25};
26
27// ── Fan-out over the bundler's 11 streams ───────────────────────────────────────────────────── //
28
29/// Fetches all 11 streams into a `PreprocessingIterator`.
30macro_rules! build_iterator {
31    ($self:ident, $req:ident, $(($field:ident, $stream:ident, $kind:ident)),+ $(,)?) => {
32        PreprocessingIterator {
33            $( $field: $self.$stream.next_n($req[PreprocessingKind::$kind])?.into_iter() ),+
34        }
35    };
36}
37
38/// Applies a method call to the stream serving `kind`, e.g. `with_stream!(self, kind, position())`.
39/// The call's result type must not vary with the stream.
40macro_rules! with_stream {
41    ($self:ident, $kind:expr, $($call:tt)+) => {
42        match $kind {
43            PreprocessingKind::BitSinglets       => $self.binary_singlet_stream.$($call)+,
44            PreprocessingKind::BitTriples        => $self.binary_triple_stream.$($call)+,
45            PreprocessingKind::BaseFieldSinglets => $self.basefield_singlet_stream.$($call)+,
46            PreprocessingKind::BaseFieldTriples  => $self.basefield_triple_stream.$($call)+,
47            PreprocessingKind::BaseFieldDaBits   => $self.basefield_dabit_stream.$($call)+,
48            PreprocessingKind::ScalarSinglets    => $self.scalar_singlet_stream.$($call)+,
49            PreprocessingKind::ScalarTriples     => $self.scalar_triple_stream.$($call)+,
50            PreprocessingKind::ScalarDaBits      => $self.scalar_dabit_stream.$($call)+,
51            PreprocessingKind::MpcFieldSinglets  => $self.mpc_field_singlet_stream.$($call)+,
52            PreprocessingKind::MpcFieldTriples   => $self.mpc_field_triple_stream.$($call)+,
53            PreprocessingKind::MpcFieldDaBits    => $self.mpc_field_dabit_stream.$($call)+,
54        }
55    };
56}
57
58/// A stream that can fill a [`StreamBundler`] slot.
59pub trait BundledStream<P: Correlation>: CorrelatedStream<P, Error = AbortError> + Buffer {}
60
61impl<P: Correlation, T: CorrelatedStream<P, Error = AbortError> + Buffer> BundledStream<P> for T {}
62
63/// One type-erased bundler slot.
64type Slot<P> = Box<dyn BundledStream<P>>;
65
66/// Stream bundler, holding one stream per preprocessing type, to provide all preprocessing for a
67/// circuit via its streams.
68pub struct StreamBundler<C: MpcConfig> {
69    // Base field
70    pub basefield_dabit_stream: Slot<DaBit<BaseFieldOf<C>>>,
71    pub basefield_singlet_stream: Slot<Singlet<BaseFieldOf<C>>>,
72    pub basefield_triple_stream: Slot<Triple<BaseFieldOf<C>>>,
73    // Binary (Gf2_128)
74    pub binary_singlet_stream: Slot<Singlet<Gf2_128>>,
75    pub binary_triple_stream: Slot<Triple<Gf2_128>>,
76    // MPC field
77    pub mpc_field_dabit_stream: Slot<DaBit<MpcFieldOf<C>>>,
78    pub mpc_field_singlet_stream: Slot<Singlet<MpcFieldOf<C>>>,
79    pub mpc_field_triple_stream: Slot<Triple<MpcFieldOf<C>>>,
80    // Scalar
81    pub scalar_dabit_stream: Slot<DaBit<ScalarFieldOf<C>>>,
82    pub scalar_singlet_stream: Slot<Singlet<ScalarFieldOf<C>>>,
83    pub scalar_triple_stream: Slot<Triple<ScalarFieldOf<C>>>,
84}
85
86impl<C: MpcConfig> StreamBundler<C> {
87    /// Creates a new bundler from the given streams.
88    #[allow(clippy::too_many_arguments)]
89    pub fn new(
90        basefield_dabit_stream: impl BundledStream<DaBit<BaseFieldOf<C>>> + 'static,
91        basefield_singlet_stream: impl BundledStream<Singlet<BaseFieldOf<C>>> + 'static,
92        basefield_triple_stream: impl BundledStream<Triple<BaseFieldOf<C>>> + 'static,
93        binary_singlet_stream: impl BundledStream<Singlet<Gf2_128>> + 'static,
94        binary_triple_stream: impl BundledStream<Triple<Gf2_128>> + 'static,
95        mpc_field_dabit_stream: impl BundledStream<DaBit<MpcFieldOf<C>>> + 'static,
96        mpc_field_singlet_stream: impl BundledStream<Singlet<MpcFieldOf<C>>> + 'static,
97        mpc_field_triple_stream: impl BundledStream<Triple<MpcFieldOf<C>>> + 'static,
98        scalar_dabit_stream: impl BundledStream<DaBit<ScalarFieldOf<C>>> + 'static,
99        scalar_singlet_stream: impl BundledStream<Singlet<ScalarFieldOf<C>>> + 'static,
100        scalar_triple_stream: impl BundledStream<Triple<ScalarFieldOf<C>>> + 'static,
101    ) -> Self {
102        Self {
103            basefield_dabit_stream: Box::new(basefield_dabit_stream),
104            basefield_singlet_stream: Box::new(basefield_singlet_stream),
105            basefield_triple_stream: Box::new(basefield_triple_stream),
106            binary_singlet_stream: Box::new(binary_singlet_stream),
107            binary_triple_stream: Box::new(binary_triple_stream),
108            mpc_field_dabit_stream: Box::new(mpc_field_dabit_stream),
109            mpc_field_singlet_stream: Box::new(mpc_field_singlet_stream),
110            mpc_field_triple_stream: Box::new(mpc_field_triple_stream),
111            scalar_dabit_stream: Box::new(scalar_dabit_stream),
112            scalar_singlet_stream: Box::new(scalar_singlet_stream),
113            scalar_triple_stream: Box::new(scalar_triple_stream),
114        }
115    }
116
117    /// Elements delivered by every stream, per type.
118    pub fn positions(&self) -> CircuitPreprocessing {
119        let mut pos = CircuitPreprocessing::default();
120        for kind in PreprocessingKind::ALL {
121            pos[kind] = with_stream!(self, kind, position()) as usize;
122        }
123        pos
124    }
125
126    /// Elements currently buffered in every stream, per type.
127    pub fn buffered(&self) -> CircuitPreprocessing {
128        let mut buf = CircuitPreprocessing::default();
129        for kind in PreprocessingKind::ALL {
130            buf[kind] = with_stream!(self, kind, buffered()) as usize;
131        }
132        buf
133    }
134
135    /// Advances every stream to its per-type `target`. Rejects the whole call if any target is
136    /// behind its stream's position; otherwise dispatches all resyncs and joins them.
137    pub async fn resync(&self, targets: &CircuitPreprocessing) -> Result<(), AbortError> {
138        let current = self.positions();
139        for kind in PreprocessingKind::ALL {
140            if targets[kind] < current[kind] {
141                return Err(CorrelatedStreamError::ResyncRewind {
142                    current: current[kind] as u64,
143                    target: targets[kind] as u64,
144                }
145                .into());
146            }
147        }
148        ResyncHandle::join_all(
149            PreprocessingKind::ALL
150                .map(|kind| with_stream!(self, kind, resync(targets[kind] as u64))),
151        )
152        .await
153    }
154
155    /// The shared buffer config of every stream; itself a [`Buffer`] over the whole bundle.
156    pub fn buffer_configs(&self) -> PerStream<SharedBufferConfig> {
157        PerStream::from_fn(|kind| with_stream!(self, kind, config().clone()))
158    }
159}
160
161// ──────────────────────── PreprocessingBundler impl ──────────────────────── //
162
163impl<C: MpcConfig> Bundler for StreamBundler<C> {
164    type Iterator = PreprocessingIterator<C>;
165
166    fn fetch(
167        &mut self,
168        req: &CircuitPreprocessing,
169    ) -> Result<PreprocessingIterator<C>, BundlerError> {
170        Ok(build_iterator!(
171            self,
172            req,
173            (binary_singlets, binary_singlet_stream, BitSinglets),
174            (binary_triples, binary_triple_stream, BitTriples),
175            (
176                base_field_singlets,
177                basefield_singlet_stream,
178                BaseFieldSinglets
179            ),
180            (
181                base_field_triples,
182                basefield_triple_stream,
183                BaseFieldTriples
184            ),
185            (base_field_dabits, basefield_dabit_stream, BaseFieldDaBits),
186            (scalar_singlets, scalar_singlet_stream, ScalarSinglets),
187            (scalar_triples, scalar_triple_stream, ScalarTriples),
188            (scalar_dabits, scalar_dabit_stream, ScalarDaBits),
189            (
190                mpc_field_singlets,
191                mpc_field_singlet_stream,
192                MpcFieldSinglets
193            ),
194            (mpc_field_triples, mpc_field_triple_stream, MpcFieldTriples),
195            (mpc_field_dabits, mpc_field_dabit_stream, MpcFieldDaBits),
196        ))
197    }
198
199    /// Prefetches into every stream concurrently; resolves once all complete.
200    fn prefetch(&self, req: &CircuitPreprocessing) -> PrefetchHandle<AbortError> {
201        PrefetchHandle::join_all(
202            PreprocessingKind::ALL.map(|kind| with_stream!(self, kind, prefetch_n(req[kind]))),
203        )
204    }
205}
core_utils/preprocessing/bundler.rs

core_utils/preprocessing/
bundler.rs
