openvm_recursion_circuit/
bus.rs

1use openvm_poseidon2_air::POSEIDON2_WIDTH;
2use openvm_recursion_circuit_derive::AlignedBorrow;
3use openvm_stark_backend::interaction::InteractionBuilder;
4use openvm_stark_sdk::config::baby_bear_poseidon2::{DIGEST_SIZE, D_EF};
5use p3_field::PrimeCharacteristicRing;
6
7#[macro_export]
8macro_rules! define_typed_lookup_bus {
9    ($Bus:ident, $Msg:ident) => {
10        #[derive(Copy, Clone, Debug)]
11        pub struct $Bus(openvm_stark_backend::interaction::LookupBus);
12
13        impl $Bus {
14            #[inline]
15            pub fn new(bus_index: openvm_stark_backend::interaction::BusIndex) -> Self {
16                Self(openvm_stark_backend::interaction::LookupBus::new(bus_index))
17            }
18
19            pub fn lookup_key<AB>(
20                &self,
21                builder: &mut AB,
22                key: $Msg<impl Into<AB::Expr> + Clone>,
23                enabled: impl Into<AB::Expr>,
24            ) where
25                AB: openvm_stark_backend::interaction::InteractionBuilder,
26            {
27                self.0.lookup_key(builder, key.to_vec(), enabled);
28            }
29
30            #[inline]
31            pub fn add_key_with_lookups<AB>(
32                &self,
33                builder: &mut AB,
34                key: $Msg<impl Into<AB::Expr> + Clone>,
35                num_lookups: impl Into<AB::Expr>,
36            ) where
37                AB: openvm_stark_backend::interaction::InteractionBuilder,
38            {
39                self.0
40                    .add_key_with_lookups(builder, key.to_vec(), num_lookups);
41            }
42        }
43    };
44}
45
46#[macro_export]
47macro_rules! define_typed_per_proof_lookup_bus {
48    ($Bus:ident, $Msg:ident) => {
49        #[derive(Copy, Clone, Debug)]
50        pub struct $Bus(openvm_stark_backend::interaction::LookupBus);
51
52        impl $Bus {
53            #[inline]
54            pub fn new(bus_index: openvm_stark_backend::interaction::BusIndex) -> Self {
55                Self(openvm_stark_backend::interaction::LookupBus::new(bus_index))
56            }
57
58            pub fn lookup_key<AB>(
59                &self,
60                builder: &mut AB,
61                proof_idx: impl Into<AB::Expr>,
62                key: $Msg<impl Into<AB::Expr> + Clone>,
63                enabled: impl Into<AB::Expr>,
64            ) where
65                AB: openvm_stark_backend::interaction::InteractionBuilder,
66            {
67                let key = core::iter::once(proof_idx.into())
68                    .chain(key.to_vec().into_iter().map(|x| x.into()))
69                    .collect::<Vec<_>>();
70                self.0.lookup_key(builder, key.to_vec(), enabled);
71            }
72
73            #[inline]
74            pub fn add_key_with_lookups<AB>(
75                &self,
76                builder: &mut AB,
77                proof_idx: impl Into<AB::Expr>,
78                key: $Msg<impl Into<AB::Expr> + Clone>,
79                num_lookups: impl Into<AB::Expr>,
80            ) where
81                AB: openvm_stark_backend::interaction::InteractionBuilder,
82            {
83                let key = core::iter::once(proof_idx.into())
84                    .chain(key.to_vec().into_iter().map(|x| x.into()))
85                    .collect::<Vec<_>>();
86                self.0
87                    .add_key_with_lookups(builder, key.to_vec(), num_lookups);
88            }
89        }
90    };
91}
92
93#[macro_export]
94macro_rules! define_typed_permutation_bus {
95    ($Bus:ident, $Msg:ident) => {
96        #[derive(Copy, Clone, Debug)]
97        pub struct $Bus(openvm_stark_backend::interaction::PermutationCheckBus);
98
99        impl $Bus {
100            #[inline]
101            pub fn new(bus_index: openvm_stark_backend::interaction::BusIndex) -> Self {
102                Self(openvm_stark_backend::interaction::PermutationCheckBus::new(
103                    bus_index,
104                ))
105            }
106
107            #[inline]
108            pub fn send<AB>(
109                &self,
110                builder: &mut AB,
111                message: $Msg<impl Into<AB::Expr> + Clone>,
112                enabled: impl Into<AB::Expr>,
113            ) where
114                AB: openvm_stark_backend::interaction::InteractionBuilder,
115            {
116                self.0.send(builder, message.to_vec(), enabled);
117            }
118
119            #[inline]
120            pub fn receive<AB>(
121                &self,
122                builder: &mut AB,
123                message: $Msg<impl Into<AB::Expr> + Clone>,
124                enabled: impl Into<AB::Expr>,
125            ) where
126                AB: openvm_stark_backend::interaction::InteractionBuilder,
127            {
128                self.0.receive(builder, message.to_vec(), enabled);
129            }
130        }
131    };
132}
133
134#[macro_export]
135macro_rules! define_typed_per_proof_permutation_bus {
136    ($Bus:ident, $Msg:ident) => {
137        #[derive(Copy, Clone, Debug)]
138        pub struct $Bus(openvm_stark_backend::interaction::PermutationCheckBus);
139
140        impl $Bus {
141            #[inline]
142            pub fn new(bus_index: openvm_stark_backend::interaction::BusIndex) -> Self {
143                Self(openvm_stark_backend::interaction::PermutationCheckBus::new(
144                    bus_index,
145                ))
146            }
147
148            #[inline]
149            pub fn send<AB>(
150                &self,
151                builder: &mut AB,
152                proof_idx: impl Into<AB::Expr>,
153                message: $Msg<impl Into<AB::Expr> + Clone>,
154                enabled: impl Into<AB::Expr>,
155            ) where
156                AB: openvm_stark_backend::interaction::InteractionBuilder,
157            {
158                let message = core::iter::once(proof_idx.into())
159                    .chain(message.to_vec().into_iter().map(|x| x.into()))
160                    .collect::<Vec<_>>();
161                self.0.send(builder, message, enabled);
162            }
163
164            #[inline]
165            pub fn receive<AB>(
166                &self,
167                builder: &mut AB,
168                proof_idx: impl Into<AB::Expr>,
169                message: $Msg<impl Into<AB::Expr> + Clone>,
170                enabled: impl Into<AB::Expr>,
171            ) where
172                AB: openvm_stark_backend::interaction::InteractionBuilder,
173            {
174                let message = core::iter::once(proof_idx.into())
175                    .chain(message.to_vec().into_iter().map(|x| x.into()))
176                    .collect::<Vec<_>>();
177                self.0.receive(builder, message, enabled);
178            }
179        }
180    };
181}
182
183#[repr(C)]
184#[derive(AlignedBorrow, Debug, Clone)]
185pub struct GkrModuleMessage<T> {
186    pub tidx: T,
187    pub n_logup: T,
188    pub n_max: T,
189    pub is_n_max_greater: T,
190}
191
192define_typed_per_proof_permutation_bus!(GkrModuleBus, GkrModuleMessage);
193
194#[repr(C)]
195#[derive(AlignedBorrow, Debug, Clone)]
196pub struct FractionFolderInputMessage<T> {
197    pub num_present_airs: T,
198}
199
200define_typed_per_proof_permutation_bus!(FractionFolderInputBus, FractionFolderInputMessage);
201
202#[repr(C)]
203#[derive(AlignedBorrow, Debug, Clone)]
204pub struct ExpressionClaimNMaxMessage<T> {
205    pub n_max: T,
206}
207
208define_typed_per_proof_permutation_bus!(ExpressionClaimNMaxBus, ExpressionClaimNMaxMessage);
209
210#[repr(C)]
211#[derive(AlignedBorrow, Debug, Clone)]
212pub struct NLiftMessage<T> {
213    pub air_idx: T,
214    pub n_lift: T,
215}
216
217define_typed_per_proof_permutation_bus!(NLiftBus, NLiftMessage);
218
219#[repr(C)]
220#[derive(AlignedBorrow, Debug, Clone)]
221pub struct BatchConstraintModuleMessage<T> {
222    pub tidx: T,
223    pub gkr_input_layer_claim: [[T; D_EF]; 2],
224}
225
226define_typed_per_proof_permutation_bus!(BatchConstraintModuleBus, BatchConstraintModuleMessage);
227
228#[repr(C)]
229#[derive(AlignedBorrow, Debug, Clone)]
230pub struct StackingModuleMessage<T> {
231    pub tidx: T,
232}
233
234define_typed_per_proof_permutation_bus!(StackingModuleBus, StackingModuleMessage);
235
236#[repr(C)]
237#[derive(AlignedBorrow, Debug, Clone)]
238pub struct WhirModuleMessage<T> {
239    /// The `tidx` _after_ batching randomness `mu` is sampled.
240    pub tidx: T,
241    /// The reduced opening claim after batching.
242    pub claim: [T; 4],
243}
244
245define_typed_per_proof_permutation_bus!(WhirModuleBus, WhirModuleMessage);
246
247#[repr(C)]
248#[derive(AlignedBorrow, Debug, Clone)]
249pub struct WhirMuMessage<T> {
250    /// The batching randomness to combine stacking claims.
251    pub mu: [T; D_EF],
252}
253
254define_typed_per_proof_permutation_bus!(WhirMuBus, WhirMuMessage);
255
256#[repr(C)]
257#[derive(AlignedBorrow, Debug, Clone)]
258pub struct TranscriptBusMessage<T> {
259    pub tidx: T,
260    pub value: T,
261    pub is_sample: T,
262}
263
264define_typed_per_proof_permutation_bus!(TranscriptBus, TranscriptBusMessage);
265
266impl TranscriptBus {
267    pub fn observe<AB: InteractionBuilder>(
268        &self,
269        builder: &mut AB,
270        proof_idx: AB::Var,
271        tidx: impl Into<AB::Expr>,
272        value: impl Into<AB::Expr>,
273        is_enabled: impl Into<AB::Expr>,
274    ) {
275        self.receive(
276            builder,
277            proof_idx,
278            TranscriptBusMessage {
279                tidx: tidx.into(),
280                value: value.into(),
281                is_sample: AB::Expr::ZERO,
282            },
283            is_enabled,
284        )
285    }
286
287    pub fn observe_ext<AB: InteractionBuilder>(
288        &self,
289        builder: &mut AB,
290        proof_idx: AB::Var,
291        tidx: impl Into<AB::Expr>,
292        value: [impl Into<AB::Expr>; D_EF],
293        is_enabled: impl Into<AB::Expr>,
294    ) {
295        let tidx = tidx.into();
296        let is_enabled = is_enabled.into();
297        for (i, x) in value.into_iter().enumerate() {
298            self.receive(
299                builder,
300                proof_idx,
301                TranscriptBusMessage {
302                    tidx: tidx.clone() + AB::Expr::from_usize(i),
303                    value: x.into(),
304                    is_sample: AB::Expr::ZERO,
305                },
306                is_enabled.clone(),
307            )
308        }
309    }
310
311    pub fn observe_commit<AB: InteractionBuilder>(
312        &self,
313        builder: &mut AB,
314        proof_idx: AB::Var,
315        tidx: impl Into<AB::Expr>,
316        commit: [impl Into<AB::Expr>; DIGEST_SIZE],
317        is_enabled: impl Into<AB::Expr>,
318    ) {
319        let tidx = tidx.into();
320        let is_enabled = is_enabled.into();
321        for (i, x) in commit.into_iter().enumerate() {
322            self.receive(
323                builder,
324                proof_idx,
325                TranscriptBusMessage {
326                    tidx: tidx.clone() + AB::Expr::from_usize(i),
327                    value: x.into(),
328                    is_sample: AB::Expr::ZERO,
329                },
330                is_enabled.clone(),
331            )
332        }
333    }
334
335    pub fn sample<AB: InteractionBuilder>(
336        &self,
337        builder: &mut AB,
338        proof_idx: AB::Var,
339        tidx: impl Into<AB::Expr>,
340        value: impl Into<AB::Expr>,
341        is_enabled: impl Into<AB::Expr>,
342    ) {
343        self.receive(
344            builder,
345            proof_idx,
346            TranscriptBusMessage {
347                tidx: tidx.into(),
348                value: value.into(),
349                is_sample: AB::Expr::ONE,
350            },
351            is_enabled,
352        )
353    }
354
355    pub fn sample_ext<AB: InteractionBuilder>(
356        &self,
357        builder: &mut AB,
358        proof_idx: AB::Var,
359        tidx: impl Into<AB::Expr>,
360        value: [impl Into<AB::Expr>; D_EF],
361        is_enabled: impl Into<AB::Expr>,
362    ) {
363        let tidx = tidx.into();
364        let is_enabled = is_enabled.into();
365        for (i, x) in value.into_iter().enumerate() {
366            self.receive(
367                builder,
368                proof_idx,
369                TranscriptBusMessage {
370                    tidx: tidx.clone() + AB::Expr::from_usize(i),
371                    value: x.into(),
372                    is_sample: AB::Expr::ONE,
373                },
374                is_enabled.clone(),
375            )
376        }
377    }
378}
379
380#[repr(C)]
381#[derive(AlignedBorrow, Debug, Clone)]
382pub struct Poseidon2PermuteMessage<T> {
383    pub input: [T; POSEIDON2_WIDTH],
384    pub output: [T; POSEIDON2_WIDTH],
385}
386
387define_typed_lookup_bus!(Poseidon2PermuteBus, Poseidon2PermuteMessage);
388
389#[repr(C)]
390#[derive(AlignedBorrow, Debug, Clone)]
391pub struct Poseidon2CompressMessage<T> {
392    pub input: [T; POSEIDON2_WIDTH],
393    pub output: [T; DIGEST_SIZE],
394}
395
396define_typed_lookup_bus!(Poseidon2CompressBus, Poseidon2CompressMessage);
397
398#[repr(u8)]
399#[derive(Debug, Copy, Clone)]
400pub(crate) enum AirShapeProperty {
401    AirId,
402    NumInteractions,
403    NeedRot,
404}
405
406impl AirShapeProperty {
407    pub fn to_field<T: PrimeCharacteristicRing>(self) -> T {
408        T::from_u8(self as u8)
409    }
410}
411
412#[repr(C)]
413#[derive(AlignedBorrow, Debug, Clone)]
414pub struct MerkleVerifyBusMessage<T> {
415    /// The idx of the merkle proof in the proof, might have additional bits (so not 0 at the root)
416    /// It will be the same for all the rows in the hashing leaves part.
417    pub merkle_idx_bit_src: T,
418    /// Merkle idx suffix after shifting right max(0, height - k) bits
419    pub current_idx_bit_src: T,
420    /// The total depth of the merkle proof including the leaves part, equal to merkle_proof.len()
421    /// + 1 + k
422    pub total_depth: T,
423    /// The height of this value, [0, k) are for the hashing leaves part, [k, total_depth) are for
424    /// the merkle proof part.
425    pub height: T,
426    /// Boolean value that indicates if this message is for the hashing leaves or Merkle proof part
427    pub is_leaf: T,
428    /// For the leaves, it will be 0 ~ 2^k - 1, for the next intermediate values, it will be 0 ~
429    /// 2^{k-1} - 1 0 for merkle proof part.
430    pub leaf_sub_idx: T,
431    /// Either the leaf hash, or the intermediate hash, or the sibling hash
432    pub value: [T; DIGEST_SIZE],
433
434    pub commit_major: T,
435    pub commit_minor: T,
436}
437
438define_typed_per_proof_permutation_bus!(MerkleVerifyBus, MerkleVerifyBusMessage);
439
440#[repr(C)]
441#[derive(AlignedBorrow, Debug, Clone)]
442pub struct AirShapeBusMessage<T> {
443    pub sort_idx: T,
444    /// The property this message encodes.
445    /// See associated enum `AirShapeProperty`.
446    pub property_idx: T,
447    /// The value of the corresponding property.
448    pub value: T,
449}
450
451define_typed_per_proof_lookup_bus!(AirShapeBus, AirShapeBusMessage);
452
453#[repr(C)]
454#[derive(AlignedBorrow, Debug, Clone)]
455pub struct AirPresenceBusMessage<T> {
456    pub air_idx: T,
457    pub is_present: T,
458}
459
460define_typed_per_proof_lookup_bus!(AirPresenceBus, AirPresenceBusMessage);
461
462#[repr(C)]
463#[derive(AlignedBorrow, Debug, Clone)]
464pub struct HyperdimBusMessage<T> {
465    pub sort_idx: T,
466    /// Sender constrains this is `abs(log_height - l_skip)`.
467    pub n_abs: T,
468    /// Sender constrains this is `n < 0 ? 1 : 0`.
469    pub n_sign_bit: T,
470}
471
472define_typed_per_proof_lookup_bus!(HyperdimBus, HyperdimBusMessage);
473
474#[repr(C)]
475#[derive(AlignedBorrow, Debug, Clone)]
476pub struct LiftedHeightsBusMessage<T> {
477    pub sort_idx: T,
478    pub part_idx: T,
479    pub commit_idx: T,
480    pub hypercube_dim: T,
481    /// Sender must constraint this equals `2^log_lifted_height`.
482    pub lifted_height: T,
483    /// Sender must constrain this equals `max(log_height, l_skip)`.
484    pub log_lifted_height: T,
485}
486
487define_typed_per_proof_lookup_bus!(LiftedHeightsBus, LiftedHeightsBusMessage);
488
489#[repr(C)]
490#[derive(AlignedBorrow, Debug, Clone)]
491pub struct StackingIndexMessage<T> {
492    pub commit_idx: T,
493    pub col_idx: T,
494}
495
496define_typed_per_proof_lookup_bus!(StackingIndicesBus, StackingIndexMessage);
497
498/// Carries all commitments in the proof.
499///
500/// The stacking commitments have `major_idx = 0` and `minor_idx =
501/// stacking_matrix_idx`. The WHIR commitments have `major_idx = whir_round + 1`
502/// and `minor_idx = 0`.
503#[repr(C)]
504#[derive(AlignedBorrow, Debug, Clone)]
505pub struct CommitmentsBusMessage<T> {
506    pub major_idx: T,
507    pub minor_idx: T,
508    pub commitment: [T; DIGEST_SIZE],
509}
510
511define_typed_per_proof_lookup_bus!(CommitmentsBus, CommitmentsBusMessage);
512
513#[repr(C)]
514#[derive(AlignedBorrow, Debug, Clone)]
515pub struct CachedCommitBusMessage<T> {
516    pub air_idx: T,
517    pub cached_idx: T,
518    pub global_cached_idx: T,
519    pub cached_commit: [T; DIGEST_SIZE],
520}
521
522define_typed_per_proof_permutation_bus!(CachedCommitBus, CachedCommitBusMessage);
523
524#[repr(C)]
525#[derive(AlignedBorrow, Debug, Clone)]
526pub struct FinalTranscriptStateMessage<T> {
527    pub state: [T; POSEIDON2_WIDTH],
528}
529
530define_typed_per_proof_permutation_bus!(FinalTranscriptStateBus, FinalTranscriptStateMessage);
531
532#[repr(C)]
533#[derive(AlignedBorrow, Debug, Clone)]
534pub struct PreHashMessage<T> {
535    pub vk_pre_hash: [T; DIGEST_SIZE],
536}
537
538define_typed_per_proof_permutation_bus!(PreHashBus, PreHashMessage);
539
540#[repr(C)]
541#[derive(AlignedBorrow, Debug, Clone)]
542pub struct XiRandomnessMessage<T> {
543    pub idx: T,
544    pub xi: [T; D_EF],
545}
546
547define_typed_per_proof_permutation_bus!(XiRandomnessBus, XiRandomnessMessage);
548
549#[repr(C)]
550#[derive(AlignedBorrow, Debug, Clone)]
551pub struct SelHypercubeBusMessage<T> {
552    pub n: T,
553    pub is_first: T,
554    pub value: [T; D_EF],
555}
556
557define_typed_per_proof_lookup_bus!(SelHypercubeBus, SelHypercubeBusMessage);
558
559#[repr(C)]
560#[derive(AlignedBorrow, Debug, Clone)]
561pub struct SelUniBusMessage<T> {
562    pub n: T,
563    pub is_first: T,
564    pub value: [T; D_EF],
565}
566
567define_typed_per_proof_lookup_bus!(SelUniBus, SelUniBusMessage);
568
569#[repr(C)]
570#[derive(AlignedBorrow, Debug, Clone)]
571pub struct ConstraintSumcheckRandomness<T> {
572    pub idx: T,
573    pub challenge: [T; D_EF],
574}
575
576define_typed_per_proof_permutation_bus!(
577    ConstraintSumcheckRandomnessBus,
578    ConstraintSumcheckRandomness
579);
580
581#[repr(C)]
582#[derive(AlignedBorrow, Debug, Clone)]
583pub struct ColumnClaimsMessage<T> {
584    // pub idx: T,
585    pub sort_idx: T,
586    pub part_idx: T,
587    pub col_idx: T,
588    pub claim: [T; D_EF],
589    pub is_rot: T,
590}
591
592define_typed_per_proof_permutation_bus!(ColumnClaimsBus, ColumnClaimsMessage);
593
594#[repr(C)]
595#[derive(AlignedBorrow, Debug, Clone)]
596pub struct WhirOpeningPointMessage<T> {
597    pub idx: T,
598    pub value: [T; D_EF],
599}
600
601// Permutation bus for opening points produced by the stacking sumcheck and eq_base AIRs.
602// Each point is sent exactly once by the producer and received exactly once by the consumer
603// (the first node in each layer of the MLE evaluation tree).
604define_typed_per_proof_permutation_bus!(WhirOpeningPointBus, WhirOpeningPointMessage);
605// Lookup bus for distributing opening points within the MLE evaluation tree. The first node
606// in each layer registers the point (received via WhirOpeningPointBus) as a lookup key,
607// and the remaining nodes in the layer look it up.
608define_typed_per_proof_lookup_bus!(WhirOpeningPointLookupBus, WhirOpeningPointMessage);
609
610#[repr(C)]
611#[derive(AlignedBorrow, Debug, Clone)]
612pub struct PublicValuesBusMessage<T> {
613    pub air_idx: T,
614    pub pv_idx: T,
615    pub value: T,
616}
617
618define_typed_per_proof_permutation_bus!(PublicValuesBus, PublicValuesBusMessage);
619
620#[repr(C)]
621#[derive(AlignedBorrow, Debug, Clone)]
622pub struct EqNegResultMessage<T> {
623    // hypercube dimension where n < 0
624    pub n: T,
625    // 2^{l_skip + n} * eq_n(u_0, r_0)
626    pub eq: [T; D_EF],
627    // 2^{l_skip + n} * k_rot_n(u_0, r_0)
628    pub k_rot: [T; D_EF],
629}
630
631define_typed_per_proof_permutation_bus!(EqNegResultBus, EqNegResultMessage);
632
633#[repr(C)]
634#[derive(AlignedBorrow, Debug, Clone)]
635pub struct EqNegBaseRandMessage<T> {
636    // sampled value u_0
637    pub u: [T; D_EF],
638    // sampled value r_0
639    pub r: [T; D_EF],
640}
641
642define_typed_per_proof_permutation_bus!(EqNegBaseRandBus, EqNegBaseRandMessage);
643
644#[repr(C)]
645#[derive(AlignedBorrow, Debug, Clone)]
646pub struct EqNsNLogupMaxMessage<T> {
647    pub n_logup: T,
648    pub n_max: T,
649}
650
651define_typed_per_proof_lookup_bus!(EqNsNLogupMaxBus, EqNsNLogupMaxMessage);
652
653#[repr(C)]
654#[derive(AlignedBorrow, Debug, Clone)]
655pub struct Eq3bShapeMessage<T> {
656    pub sort_idx: T,
657    pub n_lift: T,
658    pub n_logup: T,
659    pub num_interactions: T,
660}
661
662define_typed_per_proof_lookup_bus!(Eq3bShapeBus, Eq3bShapeMessage);
663
664#[repr(C)]
665#[derive(AlignedBorrow, Debug, Clone)]
666pub struct ConstraintsFoldingInputMessage<T> {
667    pub tidx: T,
668}
669
670define_typed_per_proof_permutation_bus!(ConstraintsFoldingInputBus, ConstraintsFoldingInputMessage);
671
672#[repr(C)]
673#[derive(AlignedBorrow, Debug, Clone)]
674pub struct InteractionsFoldingInputMessage<T> {
675    pub tidx: T,
676}
677
678define_typed_per_proof_permutation_bus!(
679    InteractionsFoldingInputBus,
680    InteractionsFoldingInputMessage
681);
openvm_recursion_circuit/bus.rs

openvm_recursion_circuit/
bus.rs
