openvm_recursion_circuit/gkr/xi_sampler/
air.rs

1use core::borrow::Borrow;
2use std::convert::Into;
3
4use openvm_circuit_primitives::{ColumnsAir, StructReflection, StructReflectionHelper, SubAir};
5use openvm_recursion_circuit_derive::AlignedBorrow;
6use openvm_stark_backend::{
7    interaction::InteractionBuilder, BaseAirWithPublicValues, PartitionedBaseAir,
8};
9use openvm_stark_sdk::config::baby_bear_poseidon2::D_EF;
10use p3_air::{Air, AirBuilder, BaseAir};
11use p3_field::{extension::BinomiallyExtendable, Field, PrimeCharacteristicRing};
12use p3_matrix::Matrix;
13
14use crate::{
15    bus::{TranscriptBus, XiRandomnessBus, XiRandomnessMessage},
16    gkr::bus::{GkrXiSamplerBus, GkrXiSamplerMessage},
17    subairs::nested_for_loop::{NestedForLoopIoCols, NestedForLoopSubAir},
18};
19
20// perf(ayush): can probably get rid of this whole air if challenges -> transcript
21// interactions are constrained in batch constraint module
22#[repr(C)]
23#[derive(AlignedBorrow, Debug, StructReflection)]
24pub struct GkrXiSamplerCols<T> {
25    /// Whether the current row is enabled (i.e. not padding)
26    pub is_enabled: T,
27    pub proof_idx: T,
28    pub is_first_challenge: T,
29
30    /// An enabled row which is not involved in any interactions
31    /// but should satisfy air constraints
32    pub is_dummy: T,
33
34    /// Challenge index
35    // perf(ayush): can probably remove idx if XiRandomnessMessage takes tidx instead
36    pub idx: T,
37
38    /// Sampled challenge
39    pub xi: [T; D_EF],
40    /// Transcript index
41    pub tidx: T,
42}
43
44#[derive(ColumnsAir)]
45#[columns_via(GkrXiSamplerCols<u8>)]
46pub struct GkrXiSamplerAir {
47    pub xi_randomness_bus: XiRandomnessBus,
48    pub transcript_bus: TranscriptBus,
49    pub xi_sampler_bus: GkrXiSamplerBus,
50}
51
52impl<F: Field> BaseAir<F> for GkrXiSamplerAir {
53    fn width(&self) -> usize {
54        GkrXiSamplerCols::<F>::width()
55    }
56}
57
58impl<F: Field> BaseAirWithPublicValues<F> for GkrXiSamplerAir {}
59impl<F: Field> PartitionedBaseAir<F> for GkrXiSamplerAir {}
60
61impl<AB: AirBuilder + InteractionBuilder> Air<AB> for GkrXiSamplerAir
62where
63    <AB::Expr as PrimeCharacteristicRing>::PrimeSubfield: BinomiallyExtendable<{ D_EF }>,
64{
65    fn eval(&self, builder: &mut AB) {
66        let main = builder.main();
67        let (local, next) = (
68            main.row_slice(0).expect("window should have two elements"),
69            main.row_slice(1).expect("window should have two elements"),
70        );
71        let local: &GkrXiSamplerCols<AB::Var> = (*local).borrow();
72        let next: &GkrXiSamplerCols<AB::Var> = (*next).borrow();
73
74        ///////////////////////////////////////////////////////////////////////
75        // Boolean Constraints
76        ///////////////////////////////////////////////////////////////////////
77
78        builder.assert_bool(local.is_dummy);
79
80        ///////////////////////////////////////////////////////////////////////
81        // Proof Index and Loop Constraints
82        ///////////////////////////////////////////////////////////////////////
83
84        type LoopSubAir = NestedForLoopSubAir<1>;
85        LoopSubAir {}.eval(
86            builder,
87            (
88                NestedForLoopIoCols {
89                    is_enabled: local.is_enabled,
90                    counter: [local.proof_idx],
91                    is_first: [local.is_first_challenge],
92                }
93                .map_into(),
94                NestedForLoopIoCols {
95                    is_enabled: next.is_enabled,
96                    counter: [next.proof_idx],
97                    is_first: [next.is_first_challenge],
98                }
99                .map_into(),
100            ),
101        );
102
103        let is_transition_challenge =
104            LoopSubAir::local_is_transition(next.is_enabled, next.is_first_challenge);
105        let is_last_challenge =
106            LoopSubAir::local_is_last(local.is_enabled, next.is_enabled, next.is_first_challenge);
107
108        // Challenge index increments by 1
109        builder
110            .when(is_transition_challenge.clone())
111            .assert_eq(next.idx, local.idx + AB::Expr::ONE);
112
113        ///////////////////////////////////////////////////////////////////////
114        // Dummy Row Constraints
115        ///////////////////////////////////////////////////////////////////////
116
117        // A proof can't contribute both dummy and non-dummy rows
118        builder
119            .when(is_transition_challenge.clone())
120            .assert_eq(next.is_dummy, local.is_dummy);
121        // Any proof segment with more than one challenge row must be non-dummy
122        builder
123            .when(is_transition_challenge.clone())
124            .assert_zero(local.is_dummy);
125
126        // Dummy rows are only allowed as a singleton placeholder row
127        builder
128            .when(local.is_dummy)
129            .assert_one(local.is_first_challenge);
130        builder
131            .when(local.is_dummy)
132            .assert_one(is_last_challenge.clone());
133
134        ///////////////////////////////////////////////////////////////////////
135        // Transition Constraints
136        ///////////////////////////////////////////////////////////////////////
137
138        builder
139            .when(is_transition_challenge.clone())
140            .assert_eq(next.tidx, local.tidx + AB::Expr::from_usize(D_EF));
141
142        ///////////////////////////////////////////////////////////////////////
143        // Module Interactions
144        ///////////////////////////////////////////////////////////////////////
145
146        let is_not_dummy = AB::Expr::ONE - local.is_dummy;
147
148        // 1. GkrXiSamplerBus
149        // 1a. Receive input from GkrInputAir
150        self.xi_sampler_bus.receive(
151            builder,
152            local.proof_idx,
153            GkrXiSamplerMessage {
154                idx: local.idx.into(),
155                tidx: local.tidx.into(),
156            },
157            local.is_first_challenge * is_not_dummy.clone(),
158        );
159        // 1b. Send output to GkrInputAir
160        let tidx_end = local.tidx + AB::Expr::from_usize(D_EF);
161        self.xi_sampler_bus.send(
162            builder,
163            local.proof_idx,
164            GkrXiSamplerMessage {
165                idx: local.idx.into(),
166                tidx: tidx_end,
167            },
168            is_last_challenge.clone() * is_not_dummy.clone(),
169        );
170
171        ///////////////////////////////////////////////////////////////////////
172        // External Interactions
173        ///////////////////////////////////////////////////////////////////////
174
175        // 1. TranscriptBus
176        // 1a. Sample challenge from transcript
177        self.transcript_bus.sample_ext(
178            builder,
179            local.proof_idx,
180            local.tidx,
181            local.xi,
182            local.is_enabled * is_not_dummy.clone(),
183        );
184
185        // 2. XiRandomnessBus
186        // 2a. Send shared randomness
187        self.xi_randomness_bus.send(
188            builder,
189            local.proof_idx,
190            XiRandomnessMessage {
191                idx: local.idx.into(),
192                xi: local.xi.map(Into::into),
193            },
194            local.is_enabled * is_not_dummy,
195        );
196    }
197}
openvm_recursion_circuit/gkr/xi_sampler/air.rs

openvm_recursion_circuit/gkr/xi_sampler/
air.rs
