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,
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! check_no_rewind {
59    ($cur:ident, $tgt:ident, $(($it:ident, $s:ident, $($p:tt)+)),+ $(,)?) => {
60        $( if $tgt.$($p)+ < $cur.$($p)+ {
61            return Err(CorrelatedStreamError::ResyncRewind {
62                current: $cur.$($p)+ as u64,
63                target: $tgt.$($p)+ as u64,
64            }.into());
65        } )+
66    };
67}
68macro_rules! resync_handles {
69    ($self:ident, $tgt:ident, $(($it:ident, $s:ident, $($p:tt)+)),+ $(,)?) => {
70        [ $( $self.$s.resync($tgt.$($p)+ as u64) ),+ ]
71    };
72}
73macro_rules! prefetch_handles {
74    ($self:ident, $req:ident, $(($it:ident, $s:ident, $($p:tt)+)),+ $(,)?) => {
75        [ $( $self.$s.prefetch_n($req.$($p)+) ),+ ]
76    };
77}
78macro_rules! collect_configs {
79    ($self:ident, $(($it:ident, $s:ident, $($p:tt)+)),+ $(,)?) => {
80        [ $( $self.$s.config().clone() ),+ ]
81    };
82}
83
84/// Stream bundler, holding one stream per preprocessing type, to provide all preprocessing for a
85/// circuit via its streams.
86pub struct StreamBundler<
87    C: MpcConfig,
88    BFDS: CorrelatedStream<DaBit<BaseFieldOf<C>>>,
89    BFSS: CorrelatedStream<Singlet<BaseFieldOf<C>>>,
90    BFTS: CorrelatedStream<Triple<BaseFieldOf<C>>>,
91    BSS: CorrelatedStream<Singlet<Gf2_128>>,
92    BTS: CorrelatedStream<Triple<Gf2_128>>,
93    MDS: CorrelatedStream<DaBit<MpcFieldOf<C>>>,
94    MSS: CorrelatedStream<Singlet<MpcFieldOf<C>>>,
95    MTS: CorrelatedStream<Triple<MpcFieldOf<C>>>,
96    SDS: CorrelatedStream<DaBit<ScalarFieldOf<C>>>,
97    SSS: CorrelatedStream<Singlet<ScalarFieldOf<C>>>,
98    STS: CorrelatedStream<Triple<ScalarFieldOf<C>>>,
99> {
100    // Base field
101    pub basefield_dabit_stream: BFDS,
102    pub basefield_singlet_stream: BFSS,
103    pub basefield_triple_stream: BFTS,
104    // Binary (Gf2_128)
105    pub binary_singlet_stream: BSS,
106    pub binary_triple_stream: BTS,
107    // MPC field
108    pub mpc_field_dabit_stream: MDS,
109    pub mpc_field_singlet_stream: MSS,
110    pub mpc_field_triple_stream: MTS,
111    // Scalar
112    pub scalar_dabit_stream: SDS,
113    pub scalar_singlet_stream: SSS,
114    pub scalar_triple_stream: STS,
115
116    pub _c: PhantomData<C>,
117}
118
119impl<C, BFDS, BFSS, BFTS, BSS, BTS, MDS, MSS, MTS, SDS, SSS, STS>
120    StreamBundler<C, BFDS, BFSS, BFTS, BSS, BTS, MDS, MSS, MTS, SDS, SSS, STS>
121where
122    C: MpcConfig,
123    BFDS: CorrelatedStream<DaBit<BaseFieldOf<C>>>,
124    BFSS: CorrelatedStream<Singlet<BaseFieldOf<C>>>,
125    BFTS: CorrelatedStream<Triple<BaseFieldOf<C>>>,
126    BSS: CorrelatedStream<Singlet<Gf2_128>>,
127    BTS: CorrelatedStream<Triple<Gf2_128>>,
128    MDS: CorrelatedStream<DaBit<MpcFieldOf<C>>>,
129    MSS: CorrelatedStream<Singlet<MpcFieldOf<C>>>,
130    MTS: CorrelatedStream<Triple<MpcFieldOf<C>>>,
131    SDS: CorrelatedStream<DaBit<ScalarFieldOf<C>>>,
132    SSS: CorrelatedStream<Singlet<ScalarFieldOf<C>>>,
133    STS: CorrelatedStream<Triple<ScalarFieldOf<C>>>,
134{
135    /// Creates a new bundler from the given streams.
136    #[allow(clippy::too_many_arguments)]
137    pub fn new(
138        basefield_dabit_stream: BFDS,
139        basefield_singlet_stream: BFSS,
140        basefield_triple_stream: BFTS,
141        binary_singlet_stream: BSS,
142        binary_triple_stream: BTS,
143        mpc_field_dabit_stream: MDS,
144        mpc_field_singlet_stream: MSS,
145        mpc_field_triple_stream: MTS,
146        scalar_dabit_stream: SDS,
147        scalar_singlet_stream: SSS,
148        scalar_triple_stream: STS,
149    ) -> Self {
150        Self {
151            basefield_dabit_stream,
152            basefield_singlet_stream,
153            basefield_triple_stream,
154            binary_singlet_stream,
155            binary_triple_stream,
156            mpc_field_dabit_stream,
157            mpc_field_singlet_stream,
158            mpc_field_triple_stream,
159            scalar_dabit_stream,
160            scalar_singlet_stream,
161            scalar_triple_stream,
162            _c: PhantomData,
163        }
164    }
165}
166
167// ──────────────────────── PreprocessingBundler impl ──────────────────────── //
168
169impl<C, BFDS, BFSS, BFTS, BSS, BTS, MDS, MSS, MTS, SDS, SSS, STS> Bundler
170    for StreamBundler<C, BFDS, BFSS, BFTS, BSS, BTS, MDS, MSS, MTS, SDS, SSS, STS>
171where
172    C: MpcConfig,
173    BFDS: CorrelatedStream<DaBit<BaseFieldOf<C>>, Error = AbortError>,
174    BFSS: CorrelatedStream<Singlet<BaseFieldOf<C>>, Error = AbortError>,
175    BFTS: CorrelatedStream<Triple<BaseFieldOf<C>>, Error = AbortError>,
176    BSS: CorrelatedStream<Singlet<Gf2_128>, Error = AbortError>,
177    BTS: CorrelatedStream<Triple<Gf2_128>, Error = AbortError>,
178    MDS: CorrelatedStream<DaBit<MpcFieldOf<C>>, Error = AbortError>,
179    MSS: CorrelatedStream<Singlet<MpcFieldOf<C>>, Error = AbortError>,
180    MTS: CorrelatedStream<Triple<MpcFieldOf<C>>, Error = AbortError>,
181    SDS: CorrelatedStream<DaBit<ScalarFieldOf<C>>, Error = AbortError>,
182    SSS: CorrelatedStream<Singlet<ScalarFieldOf<C>>, Error = AbortError>,
183    STS: CorrelatedStream<Triple<ScalarFieldOf<C>>, Error = AbortError>,
184{
185    type Iterator = PreprocessingIterator<C>;
186    fn fetch(
187        &mut self,
188        req: &CircuitPreprocessing,
189    ) -> Result<PreprocessingIterator<C>, BundlerError> {
190        Ok(for_each_stream!(build_iterator!(self, req,)))
191    }
192}
193
194// ──────────────────────── Resynchronization ──────────────────────── //
195
196impl<C, BFDS, BFSS, BFTS, BSS, BTS, MDS, MSS, MTS, SDS, SSS, STS>
197    StreamBundler<C, BFDS, BFSS, BFTS, BSS, BTS, MDS, MSS, MTS, SDS, SSS, STS>
198where
199    C: MpcConfig,
200    BFDS: CorrelatedStream<DaBit<BaseFieldOf<C>>, Error = AbortError>,
201    BFSS: CorrelatedStream<Singlet<BaseFieldOf<C>>, Error = AbortError>,
202    BFTS: CorrelatedStream<Triple<BaseFieldOf<C>>, Error = AbortError>,
203    BSS: CorrelatedStream<Singlet<Gf2_128>, Error = AbortError>,
204    BTS: CorrelatedStream<Triple<Gf2_128>, Error = AbortError>,
205    MDS: CorrelatedStream<DaBit<MpcFieldOf<C>>, Error = AbortError>,
206    MSS: CorrelatedStream<Singlet<MpcFieldOf<C>>, Error = AbortError>,
207    MTS: CorrelatedStream<Triple<MpcFieldOf<C>>, Error = AbortError>,
208    SDS: CorrelatedStream<DaBit<ScalarFieldOf<C>>, Error = AbortError>,
209    SSS: CorrelatedStream<Singlet<ScalarFieldOf<C>>, Error = AbortError>,
210    STS: CorrelatedStream<Triple<ScalarFieldOf<C>>, Error = AbortError>,
211{
212    /// Prefetches the given amount into every stream and returns a single [`PrefetchHandle`] that
213    /// resolves once all of them complete (first error wins). The prefetches run concurrently in
214    /// the background; the handle can be awaited or dropped.
215    pub fn prefetch(&self, req: &CircuitPreprocessing) -> PrefetchHandle<AbortError> {
216        let handles = for_each_stream!(prefetch_handles!(self, req,));
217        PrefetchHandle::from_future(async move {
218            let mut first_err = None;
219            for h in handles {
220                if let Err(e) = h.await {
221                    first_err.get_or_insert(e);
222                }
223            }
224            first_err.map_or(Ok(()), Err)
225        })
226    }
227
228    /// The logical position (elements delivered) of every stream, per type. Stays in sync across
229    /// parties; take the per-type maximum to agree on a resync target, then pass it to
230    /// [`resync`](Self::resync).
231    pub fn positions(&self) -> CircuitPreprocessing {
232        let mut pos = CircuitPreprocessing::default();
233        for_each_stream!(assign_positions!(self, pos,));
234        pos
235    }
236
237    /// Advances every stream to its per-type `target`, realigning all parties on the same prefix.
238    ///
239    /// All-or-nothing on rewinds: any target behind a stream's current position is rejected before
240    /// touching any stream. Otherwise all resyncs are dispatched together, every handle awaited,
241    /// and the first error (if any) returned — no short-circuiting mid-fan-out.
242    pub async fn resync(&self, targets: &CircuitPreprocessing) -> Result<(), AbortError> {
243        let current = self.positions();
244        for_each_stream!(check_no_rewind!(current, targets,));
245        let handles = for_each_stream!(resync_handles!(self, targets,));
246        // Await every handle (no early `?`) so all streams advance, then surface the first error.
247        let mut first_err = None;
248        for handle in handles {
249            if let Err(e) = handle.await {
250                first_err.get_or_insert(e);
251            }
252        }
253        first_err.map_or(Ok(()), Err)
254    }
255}
256
257// ──────────────────────── Buffer configuration ──────────────────────── //
258
259impl<C, BFDS, BFSS, BFTS, BSS, BTS, MDS, MSS, MTS, SDS, SSS, STS>
260    StreamBundler<C, BFDS, BFSS, BFTS, BSS, BTS, MDS, MSS, MTS, SDS, SSS, STS>
261where
262    C: MpcConfig,
263    BFDS: CorrelatedStream<DaBit<BaseFieldOf<C>>> + Buffer,
264    BFSS: CorrelatedStream<Singlet<BaseFieldOf<C>>> + Buffer,
265    BFTS: CorrelatedStream<Triple<BaseFieldOf<C>>> + Buffer,
266    BSS: CorrelatedStream<Singlet<Gf2_128>> + Buffer,
267    BTS: CorrelatedStream<Triple<Gf2_128>> + Buffer,
268    MDS: CorrelatedStream<DaBit<MpcFieldOf<C>>> + Buffer,
269    MSS: CorrelatedStream<Singlet<MpcFieldOf<C>>> + Buffer,
270    MTS: CorrelatedStream<Triple<MpcFieldOf<C>>> + Buffer,
271    SDS: CorrelatedStream<DaBit<ScalarFieldOf<C>>> + Buffer,
272    SSS: CorrelatedStream<Singlet<ScalarFieldOf<C>>> + Buffer,
273    STS: CorrelatedStream<Triple<ScalarFieldOf<C>>> + Buffer,
274{
275    /// The shared buffer config of every stream. The returned array is a [`Buffer`] slice, so
276    /// the whole bundle can be tuned in one call: `bundler.buffer_configs().set_capacity(n)`
277    /// writes all streams, while `.capacity()` etc. read the first.
278    pub fn buffer_configs(&self) -> [SharedBufferConfig; 11] {
279        for_each_stream!(collect_configs!(self,))
280    }
281}