openvm_static_verifier/transcript/
mod.rs

1use core::{array, iter};
2use std::sync::OnceLock;
3
4use halo2_base::{
5    gates::{GateInstructions, RangeInstructions},
6    halo2_proofs::arithmetic::Field,
7    utils::{biguint_to_fe, fe_to_biguint},
8    AssignedValue, Context, QuantumCell,
9};
10use num_bigint::BigUint;
11use openvm_stark_sdk::{
12    config::baby_bear_bn254_poseidon2::{Bn254Scalar, Digest as RootDigest},
13    openvm_stark_backend::p3_field::{Field as P3Field, PrimeField},
14};
15
16use crate::{
17    field::baby_bear::{
18        BabyBearChip, BabyBearExt4Wire, BabyBearExtWire, BabyBearWire, ReducedBabyBearExtWire,
19        ReducedBabyBearWire, BABY_BEAR_BITS, BABY_BEAR_MODULUS_U64,
20    },
21    hash::{
22        poseidon2::{pack_base_2_31_cells, Poseidon2State, DIGEST_WIDTH, POSEIDON2_RATE},
23        POSEIDON2_PARAMS, POSEIDON2_WIDTH,
24    },
25    Fr,
26};
27
28/// Number of BabyBear values bit-packed into one BN254 word (base-2^31).
29/// = floor(254 / 31) = 8
30const NUM_OBS_PER_WORD: usize = 8;
31
32/// Number of BabyBear samples extracted from one BN254 word (base-BabyBear decomposition).
33/// Must match `compute_num_samples_per_elem` in `MultiFieldTranscript`.
34const NUM_SAMPLES_PER_WORD: usize = 5;
35
36/// Precomputed bounds for the base-BabyBear hint decomposition.
37///
38/// Below `k = NUM_SAMPLES_PER_WORD`.
39struct BaseBabyBearDecompBounds {
40    /// floor((q-1) / p^k) as a field element, for the boundary equality check.
41    top_quotient_max_fe: Fr,
42    /// floor((q-1) / p^k) + 1, for the range check.
43    top_quotient_max_plus_one: BigUint,
44    /// One more than the maximum lower part when top quotient is at its max.
45    lower_max_plus_one: BigUint,
46    /// p^k as a BigUint.
47    pow_k: BigUint,
48}
49
50fn base_baby_bear_decomp_bounds() -> &'static BaseBabyBearDecompBounds {
51    static BOUNDS: OnceLock<BaseBabyBearDecompBounds> = OnceLock::new();
52    BOUNDS.get_or_init(|| {
53        let p = BigUint::from(BABY_BEAR_MODULUS_U64);
54        let one = BigUint::from(1u64);
55        let modulus = <Bn254Scalar as P3Field>::order();
56        let modulus_minus_one = &modulus - &one;
57        let pow_k = p.pow(NUM_SAMPLES_PER_WORD as u32);
58        let q_k_max = &modulus_minus_one / &pow_k;
59        let lower_max = modulus_minus_one - &q_k_max * &pow_k;
60        BaseBabyBearDecompBounds {
61            top_quotient_max_fe: biguint_to_fe(&q_k_max),
62            top_quotient_max_plus_one: &q_k_max + &one,
63            lower_max_plus_one: lower_max + one,
64            pow_k,
65        }
66    })
67}
68
69/// Decompose a BN254 field element into base-BabyBear digits using the
70/// witness-and-verify (hint) approach.
71///
72/// This helper only computes the decomposition off-circuit and loads the resulting
73/// raw digit cells plus top quotient as witnesses. Constraints must be enforced
74/// separately before the digits are wrapped as `BabyBearWire`s.
75fn load_base_baby_bear_decomposition_witness(
76    ctx: &mut Context<Fr>,
77    packed: AssignedValue<Fr>,
78) -> ([AssignedValue<Fr>; NUM_SAMPLES_PER_WORD], AssignedValue<Fr>) {
79    let p = BigUint::from(BABY_BEAR_MODULUS_U64);
80
81    // Witness: compute digits and top quotient out-of-circuit.
82    let mut value = fe_to_biguint(packed.value());
83    let digit_witnesses_big: [BigUint; NUM_SAMPLES_PER_WORD] = array::from_fn(|_| {
84        let digit = &value % &p;
85        value /= &p;
86        digit
87    });
88    let top_quotient_big = value;
89
90    // Load each digit witness.
91    let digit_witnesses =
92        array::from_fn(|idx| ctx.load_witness(biguint_to_fe(&digit_witnesses_big[idx])));
93
94    // Load top quotient witness.
95    let top_quotient = ctx.load_witness(biguint_to_fe(&top_quotient_big));
96    (digit_witnesses, top_quotient)
97}
98
99fn constrain_base_baby_bear_decomposition(
100    ctx: &mut Context<Fr>,
101    gate: &impl GateInstructions<Fr>,
102    range: &impl RangeInstructions<Fr>,
103    packed: AssignedValue<Fr>,
104    digit_witnesses: [AssignedValue<Fr>; NUM_SAMPLES_PER_WORD],
105    top_quotient: AssignedValue<Fr>,
106    bounds: &BaseBabyBearDecompBounds,
107) -> [BabyBearWire; NUM_SAMPLES_PER_WORD] {
108    let p = BigUint::from(BABY_BEAR_MODULUS_U64);
109    let one = BigUint::from(1u64);
110
111    for &digit in &digit_witnesses {
112        range.check_less_than_safe(ctx, digit, BABY_BEAR_MODULUS_U64);
113    }
114    let top_quotient_valid =
115        range.is_big_less_than_safe(ctx, top_quotient, bounds.top_quotient_max_plus_one.clone());
116    gate.assert_is_const(ctx, &top_quotient_valid, &Fr::ONE);
117
118    // Verify recomposition: lower = sum(digit_i * p^i)
119    let lower = gate.inner_product(
120        ctx,
121        digit_witnesses.iter().copied(),
122        iter::successors(Some(one), |power| Some(power * &p))
123            .take(NUM_SAMPLES_PER_WORD)
124            .map(|power| QuantumCell::Constant(biguint_to_fe(&power))),
125    );
126
127    // packed == top_quotient * p^k + lower
128    let recomposed = gate.mul_add(
129        ctx,
130        top_quotient,
131        QuantumCell::Constant(biguint_to_fe(&bounds.pow_k)),
132        lower,
133    );
134    ctx.constrain_equal(&packed, &recomposed);
135
136    // Boundary check: when top_quotient is at max, lower must not exceed the remainder.
137    let at_top_boundary = gate.is_equal(
138        ctx,
139        top_quotient,
140        QuantumCell::Constant(bounds.top_quotient_max_fe),
141    );
142    // Range-check lower against p^k (not lower_max_plus_one) because lower can be
143    // any value in [0, p^k - 1] and p^k may need more bits than lower_max_plus_one.
144    // Using is_big_less_than_safe with lower_max_plus_one would derive an insufficient
145    // range_bits from lower_max_plus_one.bits() (e.g. 152 vs the 155 bits needed).
146    let lower_range_bits =
147        (bounds.pow_k.bits() as usize).div_ceil(range.lookup_bits()) * range.lookup_bits();
148    range.range_check(ctx, lower, lower_range_bits);
149    let lower_is_valid = range.is_less_than(
150        ctx,
151        lower,
152        QuantumCell::Constant(biguint_to_fe(&bounds.lower_max_plus_one)),
153        lower_range_bits,
154    );
155    let lower_is_invalid = gate.not(ctx, lower_is_valid);
156    let lower_violation = gate.mul(ctx, at_top_boundary, lower_is_invalid);
157    gate.assert_is_const(ctx, &lower_violation, &Fr::ZERO);
158
159    digit_witnesses.map(|value| BabyBearWire {
160        value,
161        max_bits: BABY_BEAR_BITS,
162    })
163}
164
165fn decompose_bn254_to_base_baby_bear_digits(
166    ctx: &mut Context<Fr>,
167    baby_bear: &BabyBearChip,
168    packed: AssignedValue<Fr>,
169) -> [BabyBearWire; NUM_SAMPLES_PER_WORD] {
170    let bounds = base_baby_bear_decomp_bounds();
171    let range = baby_bear.range();
172    let gate = range.gate();
173    let (digit_witnesses, top_quotient) = load_base_baby_bear_decomposition_witness(ctx, packed);
174    constrain_base_baby_bear_decomposition(
175        ctx,
176        gate,
177        range,
178        packed,
179        digit_witnesses,
180        top_quotient,
181        bounds,
182    )
183}
184
185#[derive(Clone, Debug)]
186pub struct DigestWire {
187    pub elems: [AssignedValue<Fr>; DIGEST_WIDTH],
188}
189
190pub fn digest_wire_from_root(root: AssignedValue<Fr>) -> DigestWire {
191    DigestWire {
192        elems: array::from_fn(|_| root),
193    }
194}
195
196fn bn254_to_halo2(value: Bn254Scalar) -> Fr {
197    biguint_to_fe(&value.as_canonical_biguint())
198}
199
200/// Circuit gadget that mirrors `MultiFieldTranscript<BabyBear, Bn254Scalar, _, 3, 2>`.
201///
202/// Uses absorb_idx/sample_idx sponge tracking, base-2^31 observe packing,
203/// base-BabyBear sample decomposition, and direct digest absorption.
204#[derive(Clone, Debug)]
205pub struct TranscriptChip {
206    baby_bear: BabyBearChip,
207    sponge_state: [AssignedValue<Fr>; POSEIDON2_WIDTH],
208    absorb_idx: usize,
209    sample_idx: usize,
210    observe_buf: Vec<ReducedBabyBearWire>,
211    sample_buf: Vec<BabyBearWire>,
212}
213
214impl TranscriptChip {
215    pub fn baby_bear(&self) -> &BabyBearChip {
216        &self.baby_bear
217    }
218
219    pub fn new(ctx: &mut Context<Fr>, baby_bear: BabyBearChip) -> Self {
220        let zero = ctx.load_zero();
221        Self {
222            baby_bear,
223            sponge_state: array::from_fn(|_| zero),
224            absorb_idx: 0,
225            sample_idx: 0,
226            observe_buf: Vec::with_capacity(NUM_OBS_PER_WORD),
227            sample_buf: Vec::with_capacity(NUM_SAMPLES_PER_WORD),
228        }
229    }
230
231    pub fn load_digest_witness(ctx: &mut Context<Fr>, digest: RootDigest) -> DigestWire {
232        DigestWire {
233            elems: array::from_fn(|i| ctx.load_witness(bn254_to_halo2(digest[i]))),
234        }
235    }
236
237    // --- Low-level sponge (matches DuplexSponge::absorb/squeeze) ---
238
239    fn sponge_absorb(&mut self, ctx: &mut Context<Fr>, value: AssignedValue<Fr>) {
240        self.sponge_state[self.absorb_idx] = value;
241        self.absorb_idx += 1;
242        if self.absorb_idx == POSEIDON2_RATE {
243            self.permute_state(ctx);
244            self.absorb_idx = 0;
245            self.sample_idx = POSEIDON2_RATE;
246        }
247    }
248
249    fn sponge_squeeze(&mut self, ctx: &mut Context<Fr>) -> AssignedValue<Fr> {
250        if self.absorb_idx != 0 || self.sample_idx == 0 {
251            self.permute_state(ctx);
252            self.absorb_idx = 0;
253            self.sample_idx = POSEIDON2_RATE;
254        }
255        self.sample_idx -= 1;
256        self.sponge_state[self.sample_idx]
257    }
258
259    fn permute_state(&mut self, ctx: &mut Context<Fr>) {
260        let gate = self.baby_bear.range().gate();
261        let mut state = Poseidon2State::new(self.sponge_state);
262        state.permutation(ctx, gate, &POSEIDON2_PARAMS);
263        self.sponge_state = state.s;
264    }
265
266    // --- Observe/sample buffer management ---
267
268    // Rule: observe-side operations call `invalidate_samples`,
269    //       sample-side operations call `flush_observe_buf`,
270    //       cross-layer operations call both.
271
272    fn invalidate_samples(&mut self) {
273        self.sample_buf.clear();
274    }
275
276    fn flush_observe_buf(&mut self, ctx: &mut Context<Fr>) {
277        if !self.observe_buf.is_empty() {
278            let gate = self.baby_bear.range().gate();
279            let packed = pack_base_2_31_cells(ctx, gate, &self.observe_buf);
280            self.sponge_absorb(ctx, packed);
281            self.observe_buf.clear();
282        }
283    }
284
285    fn absorb_digest(&mut self, ctx: &mut Context<Fr>, digest: &DigestWire) {
286        self.invalidate_samples();
287        self.flush_observe_buf(ctx);
288        for &elem in &digest.elems {
289            self.sponge_absorb(ctx, elem);
290        }
291    }
292
293    pub fn observe(&mut self, ctx: &mut Context<Fr>, value: &ReducedBabyBearWire) {
294        self.invalidate_samples();
295        self.observe_buf.push(*value);
296        if self.observe_buf.len() == NUM_OBS_PER_WORD {
297            self.flush_observe_buf(ctx);
298        }
299    }
300
301    pub fn observe_ext(&mut self, ctx: &mut Context<Fr>, value: &ReducedBabyBearExtWire) {
302        for coeff in value.coeffs() {
303            self.observe(ctx, coeff);
304        }
305    }
306
307    /// Absorb digest words directly into the sponge (lossless).
308    pub fn observe_commit(&mut self, ctx: &mut Context<Fr>, digest: &DigestWire) {
309        self.absorb_digest(ctx, digest);
310    }
311
312    pub fn sample(&mut self, ctx: &mut Context<Fr>) -> BabyBearWire {
313        if let Some(val) = self.sample_buf.pop() {
314            return val;
315        }
316        self.flush_observe_buf(ctx);
317        let squeezed = self.sponge_squeeze(ctx);
318        self.sample_buf = Vec::from(decompose_bn254_to_base_baby_bear_digits(
319            ctx,
320            &self.baby_bear,
321            squeezed,
322        ));
323        // Reverse so pop() returns digits in order (b_0 first).
324        self.sample_buf.reverse();
325        self.sample_buf
326            .pop()
327            .expect("sample_buf should be non-empty")
328    }
329
330    pub fn sample_ext(&mut self, ctx: &mut Context<Fr>) -> BabyBearExtWire {
331        let coeffs = array::from_fn(|_| self.sample(ctx));
332        BabyBearExt4Wire(coeffs)
333    }
334
335    pub fn sample_bits(&mut self, ctx: &mut Context<Fr>, bits: usize) -> AssignedValue<Fr> {
336        assert!(
337            bits < (u32::BITS as usize),
338            "sample_bits requires bits < 32: {bits}"
339        );
340        assert!(
341            (1u64 << bits) < BABY_BEAR_MODULUS_U64,
342            "sample_bits requires (1 << bits) < modulus: bits={bits}"
343        );
344
345        let sampled = self.sample(ctx);
346        if bits == 0 {
347            return ctx.load_zero();
348        }
349        // PERF[jpw]: we could optimize this since the divisor is a power of 2
350        let range = self.baby_bear.range();
351        let divisor = BigUint::from(1u64) << bits;
352        let (_, rem) = range.div_mod(ctx, sampled.value, divisor, BABY_BEAR_BITS);
353        rem
354    }
355
356    /// Asserts that the PoW witness must pass.
357    pub fn check_witness(
358        &mut self,
359        ctx: &mut Context<Fr>,
360        bits: usize,
361        witness: &ReducedBabyBearWire,
362    ) {
363        if bits == 0 {
364            return;
365        }
366
367        self.observe(ctx, witness);
368        let sampled_bits = self.sample_bits(ctx, bits);
369        self.baby_bear
370            .range()
371            .gate()
372            .assert_is_const(ctx, &sampled_bits, &Fr::ZERO);
373    }
374}
375
376#[cfg(test)]
377pub(crate) mod tests;
openvm_static_verifier/transcript/mod.rs

openvm_static_verifier/transcript/
mod.rs
