openvm_recursion_circuit/stacking/eq_bits/
air.rs

1use std::borrow::Borrow;
2
3use openvm_circuit_primitives::{
4    utils::{and, not},
5    ColumnsAir, StructReflection, StructReflectionHelper, SubAir,
6};
7use openvm_recursion_circuit_derive::AlignedBorrow;
8use openvm_stark_backend::{
9    interaction::InteractionBuilder, BaseAirWithPublicValues, PartitionedBaseAir,
10};
11use openvm_stark_sdk::config::baby_bear_poseidon2::{D_EF, F};
12use p3_air::{Air, AirBuilder, BaseAir};
13use p3_field::{
14    extension::BinomiallyExtendable, Field, PrimeCharacteristicRing, PrimeField32, TwoAdicField,
15};
16use p3_matrix::Matrix;
17
18use crate::{
19    stacking::bus::{
20        EqBitsInternalBus, EqBitsInternalMessage, EqBitsLookupBus, EqBitsLookupMessage,
21        EqRandValuesLookupBus, EqRandValuesLookupMessage,
22    },
23    subairs::nested_for_loop::{NestedForLoopIoCols, NestedForLoopSubAir},
24    utils::{
25        assert_one_ext, assert_zeros, ext_field_add_scalar, ext_field_multiply,
26        ext_field_multiply_scalar, ext_field_subtract, ext_field_subtract_scalar,
27    },
28};
29
30#[repr(C)]
31#[derive(AlignedBorrow, StructReflection)]
32pub struct EqBitsCols<F> {
33    // Proof index columns for continuations
34    pub proof_idx: F,
35    pub is_valid: F,
36    pub is_first: F,
37
38    // Multiplicities of internal and external lookups
39    pub internal_child_flag: F,
40    pub external_mult: F,
41
42    // Parent row's b_value and evaluation
43    pub sub_b_value: F,
44    pub sub_eval: [F; D_EF],
45
46    // This row's b_value's LSB, number of bits, and u_{n_stack - num_bits + 1}
47    pub b_lsb: F,
48    pub num_bits: F,
49    pub u_val: [F; D_EF],
50}
51
52#[derive(ColumnsAir)]
53#[columns_via(EqBitsCols<u8>)]
54pub struct EqBitsAir {
55    // Internal buses
56    pub eq_bits_internal_bus: EqBitsInternalBus,
57    pub eq_bits_lookup_bus: EqBitsLookupBus,
58    pub eq_rand_values_bus: EqRandValuesLookupBus,
59
60    // Other fields
61    pub n_stack: usize,
62    pub l_skip: usize,
63}
64
65impl BaseAirWithPublicValues<F> for EqBitsAir {}
66impl PartitionedBaseAir<F> for EqBitsAir {}
67
68impl<F> BaseAir<F> for EqBitsAir {
69    fn width(&self) -> usize {
70        EqBitsCols::<F>::width()
71    }
72}
73
74impl<AB: AirBuilder + InteractionBuilder> Air<AB> for EqBitsAir
75where
76    AB::F: PrimeField32 + TwoAdicField,
77    <AB::Expr as PrimeCharacteristicRing>::PrimeSubfield: BinomiallyExtendable<{ D_EF }>,
78{
79    fn eval(&self, builder: &mut AB) {
80        let main = builder.main();
81        let (local, next) = (
82            main.row_slice(0).expect("window should have two elements"),
83            main.row_slice(1).expect("window should have two elements"),
84        );
85
86        let local: &EqBitsCols<AB::Var> = (*local).borrow();
87        let next: &EqBitsCols<AB::Var> = (*next).borrow();
88
89        NestedForLoopSubAir::<1> {}.eval(
90            builder,
91            (
92                NestedForLoopIoCols {
93                    is_enabled: local.is_valid,
94                    counter: [local.proof_idx],
95                    is_first: [local.is_first],
96                }
97                .map_into(),
98                NestedForLoopIoCols {
99                    is_enabled: next.is_valid,
100                    counter: [next.proof_idx],
101                    is_first: [next.is_first],
102                }
103                .map_into(),
104            ),
105        );
106        /*
107         * Constrain that the root evaluation is correct, i.e. eq_bits([], []) = 1.
108         */
109        let mut when_first = builder.when(local.is_first);
110
111        when_first.assert_zero(local.sub_b_value);
112        when_first.assert_zero(local.num_bits);
113        when_first.assert_zero(local.b_lsb);
114
115        assert_zeros(&mut when_first, local.u_val);
116        assert_one_ext(&mut when_first, local.sub_eval);
117
118        /*
119         * Receive the parent b_value and eq_bits(u[0..k - 1], b[0..k - 1]) eval, as
120         * well as the value of u_{n_stack - num_bits + 1}. Note that this forces
121         * each num_bits to be in [0, n_stack], as lookups are only provided for this
122         * range and idx = 0 lookups are all consumed by EqBaseAir.
123         */
124        self.eq_rand_values_bus.lookup_key(
125            builder,
126            local.proof_idx,
127            EqRandValuesLookupMessage {
128                idx: AB::Expr::from_usize(self.n_stack + 1) - local.num_bits,
129                u: local.u_val.map(Into::into),
130            },
131            and(not(local.is_first), local.is_valid),
132        );
133
134        self.eq_bits_internal_bus.receive(
135            builder,
136            local.proof_idx,
137            EqBitsInternalMessage {
138                b_value: local.sub_b_value.into(),
139                num_bits: local.num_bits.into() - AB::Expr::ONE,
140                eval: local.sub_eval.map(Into::into),
141                child_lsb: local.b_lsb.into(),
142            },
143            and(not(local.is_first), local.is_valid),
144        );
145
146        let b_value = AB::Expr::TWO * local.sub_b_value + local.b_lsb;
147        builder.assert_bool(local.b_lsb);
148
149        /*
150         * Compute eq_bits(u, b) and send it to the appropriate internal and external
151         * buses. Field internal_child_flag is odd iff it has a child with lsb 0, and
152         * is >= 2 iff it has a child with new bit 1. Because internal message field
153         * num_bits must be strictly increasing, each (b_value, num_bits) pair must be
154         * unique by induction.
155         */
156        let eval = ext_field_multiply(
157            local.sub_eval,
158            ext_field_subtract_scalar::<AB::Expr>(
159                ext_field_subtract::<AB::Expr>(
160                    ext_field_add_scalar::<AB::Expr>(
161                        ext_field_multiply_scalar(local.u_val, AB::Expr::TWO * local.b_lsb),
162                        AB::Expr::ONE,
163                    ),
164                    local.u_val,
165                ),
166                local.b_lsb,
167            ),
168        );
169
170        let three = AB::F::from_u8(3);
171        builder
172            .when(not(local.is_valid))
173            .assert_zero(local.internal_child_flag);
174        builder
175            .when_ne(local.internal_child_flag, three)
176            .when_ne(local.internal_child_flag, AB::Expr::TWO)
177            .assert_bool(local.internal_child_flag);
178
179        // Multiplicity is f(x) = 1/3 * x * (x - 2) * (2x - 5), which is such that
180        // f(0), f(2) = 0 and f(1), f(3) = 1.
181        self.eq_bits_internal_bus.send(
182            builder,
183            local.proof_idx,
184            EqBitsInternalMessage {
185                b_value: b_value.clone(),
186                num_bits: local.num_bits.into(),
187                eval: eval.clone(),
188                child_lsb: AB::Expr::ZERO,
189            },
190            local.internal_child_flag
191                * (local.internal_child_flag - AB::Expr::TWO)
192                * (local.internal_child_flag * AB::Expr::TWO - AB::Expr::from_u8(5))
193                * three.inverse(),
194        );
195
196        // Multiplicity is f(x) = -1/6 * x * (x - 1) * (2x - 7), which is such that
197        // f(0), f(1) = 0 and f(2), f(3) = 1.
198        self.eq_bits_internal_bus.send(
199            builder,
200            local.proof_idx,
201            EqBitsInternalMessage {
202                b_value: b_value.clone(),
203                num_bits: local.num_bits.into(),
204                eval: eval.clone(),
205                child_lsb: AB::Expr::ONE,
206            },
207            local.internal_child_flag
208                * (AB::Expr::ONE - local.internal_child_flag)
209                * (local.internal_child_flag * AB::Expr::TWO - AB::Expr::from_u8(7))
210                * AB::F::from_u8(6).inverse(),
211        );
212
213        self.eq_bits_lookup_bus.add_key_with_lookups(
214            builder,
215            local.proof_idx,
216            EqBitsLookupMessage {
217                b_value: b_value * AB::Expr::from_usize(1 << self.l_skip),
218                num_bits: local.num_bits.into(),
219                eval,
220            },
221            local.is_valid * local.external_mult,
222        );
223    }
224}
openvm_recursion_circuit/stacking/eq_bits/air.rs

openvm_recursion_circuit/stacking/eq_bits/
air.rs
