Skip to main content

core_utils/preprocessing/
bundler.rs

1use std::marker::PhantomData;
2
3use primitives::{
4    algebra::field::binary::Gf2_128,
5    correlated_randomness::{
6        bundler::{errors::BundlerError, Bundler},
7        dabits::DaBit,
8        singlets::Singlet,
9        stream::{
10            Buffer,
11            CorrelatedStream,
12            CorrelatedStreamError,
13            PrefetchHandle,
14            SharedBufferConfig,
15        },
16        triples::Triple,
17    },
18};
19
20use crate::{
21    circuit::preprocessing::CircuitPreprocessing,
22    config::{BaseFieldOf, MpcConfig, MpcFieldOf, ScalarFieldOf},
23    errors::AbortError,
24    preprocessing::{iterator::PreprocessingIterator, BundlerPositions},
25};
26
27// ── Fan-out over the bundler's 11 streams ───────────────────────────────────────────────────── //
28// Single source of truth pairing each stream with its iterator field and `CircuitPreprocessing`
29// size path. `for_each_stream!(cb!(args,))` forwards `args` plus the 11 triples
30// `(iterator_field, stream_field, size_path)` to the callback macro `cb!`.
31macro_rules! for_each_stream {
32    ($cb:ident ! ( $($pre:tt)* )) => {
33        $cb!($($pre)*
34            (base_field_dabits,    basefield_dabit_stream,     base_field.dabits),
35            (base_field_singlets,  basefield_singlet_stream,   base_field.singlets),
36            (base_field_triples,   basefield_triple_stream,    base_field.triples),
37            (binary_singlets,      binary_singlet_stream,      bit_singlets),
38            (binary_triples,       binary_triple_stream,       bit_triples),
39            (mpc_field_dabits,   mpc_field_dabit_stream,   mpc_field.dabits),
40            (mpc_field_singlets, mpc_field_singlet_stream, mpc_field.singlets),
41            (mpc_field_triples,  mpc_field_triple_stream,  mpc_field.triples),
42            (scalar_dabits,        scalar_dabit_stream,        scalar.dabits),
43            (scalar_singlets,      scalar_singlet_stream,      scalar.singlets),
44            (scalar_triples,       scalar_triple_stream,       scalar.triples),
45        )
46    };
47}
48macro_rules! build_iterator {
49    ($self:ident, $req:ident, $(($it:ident, $s:ident, $($p:tt)+)),+ $(,)?) => {
50        PreprocessingIterator { $( $it: $self.$s.next_n($req.$($p)+)?.into_iter() ),+ }
51    };
52}
53macro_rules! assign_positions {
54    ($self:ident, $pos:ident, $(($it:ident, $s:ident, $($p:tt)+)),+ $(,)?) => {
55        $( $pos.$($p)+ = $self.$s.position() as usize; )+
56    };
57}
58macro_rules! assign_buffered {
59    ($self:ident, $buf:ident, $(($it:ident, $s:ident, $($p:tt)+)),+ $(,)?) => {
60        $( $buf.$($p)+ = $self.$s.buffered() as usize; )+
61    };
62}
63macro_rules! check_no_rewind {
64    ($cur:ident, $tgt:ident, $(($it:ident, $s:ident, $($p:tt)+)),+ $(,)?) => {
65        $( if $tgt.$($p)+ < $cur.$($p)+ {
66            return Err(CorrelatedStreamError::ResyncRewind {
67                current: $cur.$($p)+ as u64,
68                target: $tgt.$($p)+ as u64,
69            }.into());
70        } )+
71    };
72}
73macro_rules! resync_handles {
74    ($self:ident, $tgt:ident, $(($it:ident, $s:ident, $($p:tt)+)),+ $(,)?) => {
75        [ $( $self.$s.resync($tgt.$($p)+ as u64) ),+ ]
76    };
77}
78macro_rules! prefetch_handles {
79    ($self:ident, $req:ident, $(($it:ident, $s:ident, $($p:tt)+)),+ $(,)?) => {
80        [ $( $self.$s.prefetch_n($req.$($p)+) ),+ ]
81    };
82}
83macro_rules! collect_configs {
84    ($self:ident, $(($it:ident, $s:ident, $($p:tt)+)),+ $(,)?) => {
85        [ $( $self.$s.config().clone() ),+ ]
86    };
87}
88
89/// Stream bundler, holding one stream per preprocessing type, to provide all preprocessing for a
90/// circuit via its streams.
91pub struct StreamBundler<
92    C: MpcConfig,
93    BFDS: CorrelatedStream<DaBit<BaseFieldOf<C>>>,
94    BFSS: CorrelatedStream<Singlet<BaseFieldOf<C>>>,
95    BFTS: CorrelatedStream<Triple<BaseFieldOf<C>>>,
96    BSS: CorrelatedStream<Singlet<Gf2_128>>,
97    BTS: CorrelatedStream<Triple<Gf2_128>>,
98    MDS: CorrelatedStream<DaBit<MpcFieldOf<C>>>,
99    MSS: CorrelatedStream<Singlet<MpcFieldOf<C>>>,
100    MTS: CorrelatedStream<Triple<MpcFieldOf<C>>>,
101    SDS: CorrelatedStream<DaBit<ScalarFieldOf<C>>>,
102    SSS: CorrelatedStream<Singlet<ScalarFieldOf<C>>>,
103    STS: CorrelatedStream<Triple<ScalarFieldOf<C>>>,
104> {
105    // Base field
106    pub basefield_dabit_stream: BFDS,
107    pub basefield_singlet_stream: BFSS,
108    pub basefield_triple_stream: BFTS,
109    // Binary (Gf2_128)
110    pub binary_singlet_stream: BSS,
111    pub binary_triple_stream: BTS,
112    // MPC field
113    pub mpc_field_dabit_stream: MDS,
114    pub mpc_field_singlet_stream: MSS,
115    pub mpc_field_triple_stream: MTS,
116    // Scalar
117    pub scalar_dabit_stream: SDS,
118    pub scalar_singlet_stream: SSS,
119    pub scalar_triple_stream: STS,
120
121    pub _c: PhantomData<C>,
122}
123
124impl<C, BFDS, BFSS, BFTS, BSS, BTS, MDS, MSS, MTS, SDS, SSS, STS>
125    StreamBundler<C, BFDS, BFSS, BFTS, BSS, BTS, MDS, MSS, MTS, SDS, SSS, STS>
126where
127    C: MpcConfig,
128    BFDS: CorrelatedStream<DaBit<BaseFieldOf<C>>>,
129    BFSS: CorrelatedStream<Singlet<BaseFieldOf<C>>>,
130    BFTS: CorrelatedStream<Triple<BaseFieldOf<C>>>,
131    BSS: CorrelatedStream<Singlet<Gf2_128>>,
132    BTS: CorrelatedStream<Triple<Gf2_128>>,
133    MDS: CorrelatedStream<DaBit<MpcFieldOf<C>>>,
134    MSS: CorrelatedStream<Singlet<MpcFieldOf<C>>>,
135    MTS: CorrelatedStream<Triple<MpcFieldOf<C>>>,
136    SDS: CorrelatedStream<DaBit<ScalarFieldOf<C>>>,
137    SSS: CorrelatedStream<Singlet<ScalarFieldOf<C>>>,
138    STS: CorrelatedStream<Triple<ScalarFieldOf<C>>>,
139{
140    /// Creates a new bundler from the given streams.
141    #[allow(clippy::too_many_arguments)]
142    pub fn new(
143        basefield_dabit_stream: BFDS,
144        basefield_singlet_stream: BFSS,
145        basefield_triple_stream: BFTS,
146        binary_singlet_stream: BSS,
147        binary_triple_stream: BTS,
148        mpc_field_dabit_stream: MDS,
149        mpc_field_singlet_stream: MSS,
150        mpc_field_triple_stream: MTS,
151        scalar_dabit_stream: SDS,
152        scalar_singlet_stream: SSS,
153        scalar_triple_stream: STS,
154    ) -> Self {
155        Self {
156            basefield_dabit_stream,
157            basefield_singlet_stream,
158            basefield_triple_stream,
159            binary_singlet_stream,
160            binary_triple_stream,
161            mpc_field_dabit_stream,
162            mpc_field_singlet_stream,
163            mpc_field_triple_stream,
164            scalar_dabit_stream,
165            scalar_singlet_stream,
166            scalar_triple_stream,
167            _c: PhantomData,
168        }
169    }
170}
171
172// ──────────────────────── PreprocessingBundler impl ──────────────────────── //
173
174impl<C, BFDS, BFSS, BFTS, BSS, BTS, MDS, MSS, MTS, SDS, SSS, STS> Bundler
175    for StreamBundler<C, BFDS, BFSS, BFTS, BSS, BTS, MDS, MSS, MTS, SDS, SSS, STS>
176where
177    C: MpcConfig,
178    BFDS: CorrelatedStream<DaBit<BaseFieldOf<C>>, Error = AbortError>,
179    BFSS: CorrelatedStream<Singlet<BaseFieldOf<C>>, Error = AbortError>,
180    BFTS: CorrelatedStream<Triple<BaseFieldOf<C>>, Error = AbortError>,
181    BSS: CorrelatedStream<Singlet<Gf2_128>, Error = AbortError>,
182    BTS: CorrelatedStream<Triple<Gf2_128>, Error = AbortError>,
183    MDS: CorrelatedStream<DaBit<MpcFieldOf<C>>, Error = AbortError>,
184    MSS: CorrelatedStream<Singlet<MpcFieldOf<C>>, Error = AbortError>,
185    MTS: CorrelatedStream<Triple<MpcFieldOf<C>>, Error = AbortError>,
186    SDS: CorrelatedStream<DaBit<ScalarFieldOf<C>>, Error = AbortError>,
187    SSS: CorrelatedStream<Singlet<ScalarFieldOf<C>>, Error = AbortError>,
188    STS: CorrelatedStream<Triple<ScalarFieldOf<C>>, Error = AbortError>,
189{
190    type Iterator = PreprocessingIterator<C>;
191    fn fetch(
192        &mut self,
193        req: &CircuitPreprocessing,
194    ) -> Result<PreprocessingIterator<C>, BundlerError> {
195        Ok(for_each_stream!(build_iterator!(self, req,)))
196    }
197}
198
199// ──────────────────────── Resynchronization ──────────────────────── //
200
201impl<C, BFDS, BFSS, BFTS, BSS, BTS, MDS, MSS, MTS, SDS, SSS, STS>
202    StreamBundler<C, BFDS, BFSS, BFTS, BSS, BTS, MDS, MSS, MTS, SDS, SSS, STS>
203where
204    C: MpcConfig,
205    BFDS: CorrelatedStream<DaBit<BaseFieldOf<C>>, Error = AbortError>,
206    BFSS: CorrelatedStream<Singlet<BaseFieldOf<C>>, Error = AbortError>,
207    BFTS: CorrelatedStream<Triple<BaseFieldOf<C>>, Error = AbortError>,
208    BSS: CorrelatedStream<Singlet<Gf2_128>, Error = AbortError>,
209    BTS: CorrelatedStream<Triple<Gf2_128>, Error = AbortError>,
210    MDS: CorrelatedStream<DaBit<MpcFieldOf<C>>, Error = AbortError>,
211    MSS: CorrelatedStream<Singlet<MpcFieldOf<C>>, Error = AbortError>,
212    MTS: CorrelatedStream<Triple<MpcFieldOf<C>>, Error = AbortError>,
213    SDS: CorrelatedStream<DaBit<ScalarFieldOf<C>>, Error = AbortError>,
214    SSS: CorrelatedStream<Singlet<ScalarFieldOf<C>>, Error = AbortError>,
215    STS: CorrelatedStream<Triple<ScalarFieldOf<C>>, Error = AbortError>,
216{
217    /// Prefetches the given amount into every stream and returns a single [`PrefetchHandle`] that
218    /// resolves once all of them complete (first error wins). The prefetches run concurrently in
219    /// the background; the handle can be awaited or dropped.
220    pub fn prefetch(&self, req: &CircuitPreprocessing) -> PrefetchHandle<AbortError> {
221        let handles = for_each_stream!(prefetch_handles!(self, req,));
222        PrefetchHandle::from_future(async move {
223            let mut first_err = None;
224            for h in handles {
225                if let Err(e) = h.await {
226                    first_err.get_or_insert(e);
227                }
228            }
229            first_err.map_or(Ok(()), Err)
230        })
231    }
232
233    /// Advances every stream to its per-type `target`, realigning all parties on the same prefix.
234    ///
235    /// All-or-nothing on rewinds: any target behind a stream's current position is rejected before
236    /// touching any stream. Otherwise all resyncs are dispatched together, every handle awaited,
237    /// and the first error (if any) returned — no short-circuiting mid-fan-out.
238    pub async fn resync(&self, targets: &CircuitPreprocessing) -> Result<(), AbortError> {
239        let current = self.positions();
240        for_each_stream!(check_no_rewind!(current, targets,));
241        let handles = for_each_stream!(resync_handles!(self, targets,));
242        // Await every handle (no early `?`) so all streams advance, then surface the first error.
243        let mut first_err = None;
244        for handle in handles {
245            if let Err(e) = handle.await {
246                first_err.get_or_insert(e);
247            }
248        }
249        first_err.map_or(Ok(()), Err)
250    }
251}
252
253// ──────────────────────── Position / occupancy accessors ──────────────────────── //
254
255impl<C, BFDS, BFSS, BFTS, BSS, BTS, MDS, MSS, MTS, SDS, SSS, STS> BundlerPositions
256    for StreamBundler<C, BFDS, BFSS, BFTS, BSS, BTS, MDS, MSS, MTS, SDS, SSS, STS>
257where
258    C: MpcConfig,
259    BFDS: CorrelatedStream<DaBit<BaseFieldOf<C>>, Error = AbortError>,
260    BFSS: CorrelatedStream<Singlet<BaseFieldOf<C>>, Error = AbortError>,
261    BFTS: CorrelatedStream<Triple<BaseFieldOf<C>>, Error = AbortError>,
262    BSS: CorrelatedStream<Singlet<Gf2_128>, Error = AbortError>,
263    BTS: CorrelatedStream<Triple<Gf2_128>, Error = AbortError>,
264    MDS: CorrelatedStream<DaBit<MpcFieldOf<C>>, Error = AbortError>,
265    MSS: CorrelatedStream<Singlet<MpcFieldOf<C>>, Error = AbortError>,
266    MTS: CorrelatedStream<Triple<MpcFieldOf<C>>, Error = AbortError>,
267    SDS: CorrelatedStream<DaBit<ScalarFieldOf<C>>, Error = AbortError>,
268    SSS: CorrelatedStream<Singlet<ScalarFieldOf<C>>, Error = AbortError>,
269    STS: CorrelatedStream<Triple<ScalarFieldOf<C>>, Error = AbortError>,
270{
271    fn positions(&self) -> CircuitPreprocessing {
272        let mut pos = CircuitPreprocessing::default();
273        for_each_stream!(assign_positions!(self, pos,));
274        pos
275    }
276
277    fn buffered(&self) -> CircuitPreprocessing {
278        let mut buf = CircuitPreprocessing::default();
279        for_each_stream!(assign_buffered!(self, buf,));
280        buf
281    }
282
283    // Inherent `resync` (defined above) takes priority over this trait method in method-call
284    // resolution, so this just forwards to it — needed only to stay reachable through the opaque
285    // `impl PreprocessingBundler` returned by bundler builders.
286    async fn resync(&self, targets: &CircuitPreprocessing) -> Result<(), AbortError> {
287        self.resync(targets).await
288    }
289}
290
291// ──────────────────────── Buffer configuration ──────────────────────── //
292
293impl<C, BFDS, BFSS, BFTS, BSS, BTS, MDS, MSS, MTS, SDS, SSS, STS>
294    StreamBundler<C, BFDS, BFSS, BFTS, BSS, BTS, MDS, MSS, MTS, SDS, SSS, STS>
295where
296    C: MpcConfig,
297    BFDS: CorrelatedStream<DaBit<BaseFieldOf<C>>> + Buffer,
298    BFSS: CorrelatedStream<Singlet<BaseFieldOf<C>>> + Buffer,
299    BFTS: CorrelatedStream<Triple<BaseFieldOf<C>>> + Buffer,
300    BSS: CorrelatedStream<Singlet<Gf2_128>> + Buffer,
301    BTS: CorrelatedStream<Triple<Gf2_128>> + Buffer,
302    MDS: CorrelatedStream<DaBit<MpcFieldOf<C>>> + Buffer,
303    MSS: CorrelatedStream<Singlet<MpcFieldOf<C>>> + Buffer,
304    MTS: CorrelatedStream<Triple<MpcFieldOf<C>>> + Buffer,
305    SDS: CorrelatedStream<DaBit<ScalarFieldOf<C>>> + Buffer,
306    SSS: CorrelatedStream<Singlet<ScalarFieldOf<C>>> + Buffer,
307    STS: CorrelatedStream<Triple<ScalarFieldOf<C>>> + Buffer,
308{
309    /// The shared buffer config of every stream. The returned array is a [`Buffer`] slice, so
310    /// the whole bundle can be tuned in one call: `bundler.buffer_configs().set_capacity(n)`
311    /// writes all streams, while `.capacity()` etc. read the first.
312    pub fn buffer_configs(&self) -> [SharedBufferConfig; 11] {
313        for_each_stream!(collect_configs!(self,))
314    }
315}