openvm_static_verifier/tracegen/
ir_builder.rs

1//! Graph-IR generation backend for the [`chip_traits`](crate::chip_traits) traits.
2//!
3//! [`Halo2IRBuilder`] implements the same chip traits as
4//! [`Halo2Backend`](crate::backend::Halo2Backend) but records a dataflow graph
5//! of [`Halo2GraphNode`]s instead of assigning halo2 cells. Each node writes a
6//! statically-sized slice of the advice + range tapes; executing them in tape
7//! order reproduces the halo2 backend's stream byte-for-byte.
8//!
9//! To keep replay stateless (parallel-safe), the builder mirrors chip
10//! bookkeeping exactly:
11//! - `max_bits` tracking + explicit [`Halo2Opcode::BBReduce`] nodes for pre-op reduces; atomic ops
12//!   (`div`, ext `mul`/`div`) re-derive interleaved reduce decisions from operand bit bounds at
13//!   replay.
14//! - Constant caching (`zero_cell` + BabyBear const cache) matches `Context::load_zero` /
15//!   `BabyBearChip::const_cache`. Atomic ops expose internally-loaded constants (`BBDiv`→ONE,
16//!   `ExtMul`→W, `ExtDiv`→ONE/ZERO/W) as extra outputs; whether each materializes is decided at
17//!   build time and recorded in [`NodeMeta::constant_skip_inds`].
18//! - Transcript sponge/buffer state mirrors `TranscriptChip`.
19//!
20//! Fixed-column `QuantumCell::Constant`s (select branches, IP coefficients,
21//! `Const` witness values) are [`GraphCell::Const`] operands, not nodes.
22//!
23//! Non-arithmetic opcodes cover the rest of the populate pipeline:
24//! `LoadWitness`, `InnerProduct`, `DecomposeBn254ToBabyBear`, `RangeDiv`.
25
26use core::{array, iter};
27use std::collections::HashMap;
28
29use halo2_base::{
30    halo2_proofs::{
31        arithmetic::Field,
32        halo2curves::{bn256::Fr, ff::PrimeField as _},
33    },
34    utils::{bit_length, fe_to_biguint},
35};
36use itertools::Itertools;
37use openvm_stark_sdk::{
38    openvm_stark_backend::p3_field::{
39        extension::BinomiallyExtendable, BasedVectorSpace, PrimeField64,
40    },
41    p3_baby_bear::BabyBear,
42};
43use serde::{Deserialize, Serialize};
44
45use crate::{
46    chip_traits::{
47        BabyBearExt4Inst, BabyBearInst, ChipBase, GateInst, PopulateInputs, Poseidon2Inst,
48        TranscriptInst,
49    },
50    field::baby_bear::{
51        BabyBearExt4, BabyBearExt4Wire, BabyBearWire, ReducedBabyBearExt4Wire, ReducedBabyBearWire,
52        BABYBEAR_MAX_BITS, BABY_BEAR_MODULUS_U64, RESERVED_HIGH_BITS,
53    },
54    hash::{
55        poseidon2::{MULTI_FIELD32_NUM_F_ELMS, MULTI_FIELD32_RATE, POSEIDON2_RATE},
56        POSEIDON2_WIDTH,
57    },
58    tracegen::opcode_impl::{derive_opcode_metadata, UNMATERIALIZED},
59    transcript::{DigestWire, NUM_OBS_PER_WORD, NUM_SAMPLES_PER_WORD},
60};
61
62/// Unique id of an IR node; also its index in [`Halo2IRBuilder::nodes`].
63pub type NodeId = u32;
64
65/// Reduce when a bound would exceed this, mirroring `BabyBearChip`.
66const REDUCE_THRESHOLD: usize = Fr::CAPACITY as usize - RESERVED_HIGH_BITS;
67/// Bit bound for raw Bn254 cells (digests, sponge state, packed words).
68const RAW_MAX_BITS: usize = Fr::NUM_BITS as usize;
69
70const _: () = assert!(POSEIDON2_WIDTH == 3);
71
72/// IR operand/result: a node-produced value or a fixed-column constant.
73#[derive(Copy, Clone, Debug, PartialEq, Eq)]
74pub enum GraphCell {
75    /// `(node id, absolute advice offset, bits bound)`. The bits bound is the
76    /// [`BabyBearWire`] `max_bits` invariant; executors read it to re-derive
77    /// internal reduce decisions inside atomic ops.
78    Cell(NodeId, usize, u16),
79    /// Fixed-column `QuantumCell::Constant`; writes no advice cells.
80    Const(Fr),
81}
82
83impl GraphCell {
84    pub fn bits(&self) -> usize {
85        match self {
86            Self::Cell(_, _, bits) => *bits as usize,
87            Self::Const(value) => fr_bits(value),
88        }
89    }
90}
91
92/// IR opcodes. Each op writes a statically-known slice of the advice + range
93/// tapes when executed.
94#[derive(Copy, Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
95pub enum Halo2Opcode {
96    // -- gate ops --
97    /// One constant advice cell; single [`GraphCell::Const`] operand.
98    Const,
99    /// `if cond { a } else { b }`; operands `[a, b, cond]` (`a`/`b` may be const).
100    Select,
101    /// Little-endian bit decomposition into `n` bit cells.
102    Num2Bits(u16),
103    // -- babybear ops --
104    /// Barrett-style signed reduction to `[0, p)`. Operand bit bound drives the
105    /// quotient range-check width.
106    BBReduce,
107    BBAdd,
108    BBNeg,
109    BBSub,
110    BBMul,
111    BBMulAdd,
112    /// Atomic `BabyBearChip::div`. Outputs: 0=quotient, 1=internally-loaded
113    /// ONE cell (materialized only on cache miss).
114    BBDiv,
115    /// Asserts operand `≡ 0 (mod p)`.
116    BBAssertZero,
117    // -- extension ops --
118    /// Atomic `BabyBearExt4Chip::mul`; operands `[a0..a3, b0..b3]`. Outputs:
119    /// 0-3=product coefficients, 4=internally-loaded W cell (cache miss only).
120    ExtMul,
121    /// Atomic `BabyBearExt4Chip::div`; operands `[a0..a3, b0..b3]`. Outputs:
122    /// 0-3=quotient coefficients, 4-6=internally-loaded ONE/ZERO/W cells
123    /// (each cache miss only).
124    ExtDiv,
125    // -- poseidon ops --
126    /// Width-2 Poseidon2 (digest compression).
127    PoseidonPermute2T2,
128    /// Width-3 Poseidon2 (transcript sponge, digest hashing).
129    PoseidonPermute2T3,
130    // -- populate-pipeline ops --
131    /// One proof-input advice cell.
132    LoadWitness,
133    /// Proof-input cell constrained to `[0, p)` (`check_less_than_safe`).
134    CheckLessThanSafe,
135    /// Inner product of `n` `(value, coefficient)` pairs, interleaved as
136    /// `[v0, c0, v1, c1, ..]`; coefficients are [`GraphCell::Const`] in practice.
137    InnerProduct(u16),
138    /// Base-BabyBear decomposition of one Bn254 word into
139    /// `NUM_SAMPLES_PER_WORD` digit cells + boundary checks. Mirrors
140    /// `decompose_bn254_to_base_baby_bear_digits`.
141    DecomposeBn254ToBabyBear,
142    /// `rem` of `range.div_mod(operand, 2^n)` (`TranscriptChip::sample_bits`).
143    RangeDiv(u16),
144}
145
146impl Halo2Opcode {
147    pub fn num_operands(&self) -> usize {
148        match self {
149            Self::LoadWitness => 0,
150            Self::CheckLessThanSafe => 1,
151            Self::Const
152            | Self::Num2Bits(_)
153            | Self::BBReduce
154            | Self::BBNeg
155            | Self::BBAssertZero
156            | Self::DecomposeBn254ToBabyBear
157            | Self::RangeDiv(_) => 1,
158            Self::BBAdd | Self::BBSub | Self::BBMul | Self::BBDiv | Self::PoseidonPermute2T2 => 2,
159            Self::Select | Self::BBMulAdd | Self::PoseidonPermute2T3 => 3,
160            Self::ExtMul | Self::ExtDiv => 8,
161            Self::InnerProduct(n) => 2 * *n as usize,
162        }
163    }
164
165    pub fn num_results(&self) -> usize {
166        match self {
167            Self::BBAssertZero | Self::CheckLessThanSafe => 0,
168            Self::Const
169            | Self::Select
170            | Self::BBReduce
171            | Self::BBAdd
172            | Self::BBNeg
173            | Self::BBSub
174            | Self::BBMul
175            | Self::BBMulAdd
176            | Self::LoadWitness
177            | Self::InnerProduct(_)
178            | Self::RangeDiv(_) => 1,
179            Self::BBDiv | Self::PoseidonPermute2T2 => 2,
180            Self::PoseidonPermute2T3 => 3,
181            Self::ExtMul => 5,
182            Self::DecomposeBn254ToBabyBear => NUM_SAMPLES_PER_WORD,
183            Self::ExtDiv => 7,
184            Self::Num2Bits(n) => *n as usize,
185        }
186    }
187
188    pub fn name(&self) -> &'static str {
189        match self {
190            Self::Const => "Const",
191            Self::Select => "Select",
192            Self::Num2Bits(_) => "Num2Bits",
193            Self::BBReduce => "BBReduce",
194            Self::BBAdd => "BBAdd",
195            Self::BBNeg => "BBNeg",
196            Self::BBSub => "BBSub",
197            Self::BBMul => "BBMul",
198            Self::BBMulAdd => "BBMulAdd",
199            Self::BBDiv => "BBDiv",
200            Self::BBAssertZero => "BBAssertZero",
201            Self::ExtMul => "ExtMul",
202            Self::ExtDiv => "ExtDiv",
203            Self::PoseidonPermute2T2 => "PoseidonPermute2T2",
204            Self::PoseidonPermute2T3 => "PoseidonPermute2T3",
205            Self::LoadWitness => "LoadWitness",
206            Self::CheckLessThanSafe => "CheckLessThanSafe",
207            Self::InnerProduct(_) => "InnerProduct",
208            Self::DecomposeBn254ToBabyBear => "DecomposeBn254ToBabyBear",
209            Self::RangeDiv(_) => "RangeDiv",
210        }
211    }
212}
213
214#[derive(Clone, Debug)]
215pub struct Halo2GraphNode {
216    pub opcode: Halo2Opcode,
217    pub operands: Vec<GraphCell>,
218    pub id: NodeId,
219}
220
221/// Per-node replay metadata deduced at build time; together with operand
222/// values it is everything an executor needs to replay a node in isolation
223/// (no runtime cache), enabling parallel replay.
224#[derive(Clone, Debug)]
225pub struct NodeMeta {
226    /// Advice slots written by the node.
227    pub ctx_len: usize,
228    /// Absolute advice-tape offset the node begins writing at.
229    pub ctx_offset: usize,
230    /// Range slots written by the node.
231    pub lookups_len: usize,
232    /// Absolute range-tape offset the node begins writing at.
233    pub lookup_offset: usize,
234    /// Indices of the node's `load_constant` calls that write a cell (cache
235    /// misses), in call order. Drives `WitnessTape` at replay.
236    pub constant_skip_inds: Vec<u32>,
237    /// Absolute advice offset of each operand (`UNMATERIALIZED` for
238    /// [`GraphCell::Const`] operands).
239    pub arg_offsets: Vec<usize>,
240    /// Dataflow depth: `1 + max(level of operand nodes)`, 0 for source nodes.
241    pub level: u32,
242}
243
244/// Transcript sponge/buffer state, mirroring `TranscriptChip`.
245#[derive(Clone, Debug)]
246struct IrTranscript {
247    sponge_state: [GraphCell; POSEIDON2_WIDTH],
248    absorb_idx: usize,
249    sample_idx: usize,
250    observe_buf: Vec<ReducedBabyBearWire<GraphCell>>,
251    sample_buf: Vec<BabyBearWire<GraphCell>>,
252}
253
254/// Backend that records the circuit-population trace as a graph IR.
255pub struct Halo2IRBuilder {
256    /// Nodes, in tape-emission order.
257    pub nodes: Vec<Halo2GraphNode>,
258    /// Replay metadata per node, indexed by [`NodeId`].
259    pub node_meta: Vec<NodeMeta>,
260    /// Proof-input witnesses, one per `LoadWitness` node, in emission order.
261    pub input_values: Vec<Fr>,
262    /// Range-check lookup bits (drives limb decompositions).
263    lookup_bits: usize,
264    /// Next node's `ctx_offset` (advice write cursor).
265    ctx_offset: usize,
266    /// Next node's `lookup_offset` (range write cursor).
267    lookup_offset: usize,
268    /// Mirrors `Context::zero_cell`.
269    zero_cell: Option<GraphCell>,
270    /// Mirrors `BabyBearChip::const_cache` (keyed by canonical u64).
271    bb_const_cache: HashMap<Fr, BabyBearWire<GraphCell>>,
272    /// Separate cache mirroring `TranscriptChip`'s own baby-bear constant
273    /// dedup, which is distinct from `BabyBearChip::const_cache`. Preserving
274    /// this split is required to keep the verifying key unchanged.
275    transcript_bb_constant_cache: HashMap<Fr, BabyBearWire<GraphCell>>,
276    transcript: Option<IrTranscript>,
277}
278
279fn fr_bits(value: &Fr) -> usize {
280    fe_to_biguint(value).bits() as usize
281}
282
283fn bb_wire(value: GraphCell, max_bits: usize) -> BabyBearWire<GraphCell> {
284    BabyBearWire { value, max_bits }
285}
286
287enum ConstCacheType {
288    Transcript,
289    BabyBear,
290    Gate,
291}
292impl Halo2IRBuilder {
293    pub fn new(lookup_bits: usize) -> Self {
294        Halo2IRBuilder {
295            nodes: Vec::new(),
296            node_meta: Vec::new(),
297            input_values: Vec::new(),
298            lookup_bits,
299            ctx_offset: 0,
300            lookup_offset: 0,
301            zero_cell: None,
302            bb_const_cache: HashMap::new(),
303            transcript_bb_constant_cache: HashMap::new(),
304            transcript: None,
305        }
306    }
307
308    /// Total number of context-tape (advice) cells written by all nodes.
309    pub fn total_ctx_len(&self) -> usize {
310        self.ctx_offset
311    }
312
313    /// Total number of range-tape (lookup) cells written by all nodes.
314    pub fn total_lookups_len(&self) -> usize {
315        self.lookup_offset
316    }
317
318    /// Range-check lookup bits the IR was built for.
319    pub fn lookup_bits(&self) -> usize {
320        self.lookup_bits
321    }
322
323    /// Emits a node and deduces its [`NodeMeta`] by replaying the op against the
324    /// current cache state. Returns the node id and the **absolute** context-tape
325    /// offset of each logical output ([`UNMATERIALIZED`] for constants that hit a
326    /// cache and assigned no cell).
327    fn emit(
328        &mut self,
329        opcode: Halo2Opcode,
330        operands: Vec<GraphCell>,
331        cache: ConstCacheType,
332    ) -> (NodeId, Vec<usize>) {
333        debug_assert_eq!(operands.len(), opcode.num_operands());
334        let id = NodeId::try_from(self.nodes.len()).expect("node count exceeds NodeId range");
335
336        // Tape shape depends only on constant operand values, bit bounds,
337        // lookup_bits, and cache state — a nonzero dummy stands in for cell values.
338        let mut level = 0u32;
339        let mut args: Vec<Fr> = Vec::with_capacity(operands.len());
340        let mut bits: Vec<u16> = Vec::with_capacity(operands.len());
341        let mut arg_offsets: Vec<usize> = Vec::with_capacity(operands.len());
342        for cell in &operands {
343            match cell {
344                GraphCell::Cell(node, offset, _) => {
345                    level = level.max(self.node_meta[*node as usize].level + 1);
346                    args.push(Fr::ONE);
347                    arg_offsets.push(*offset);
348                }
349                GraphCell::Const(value) => {
350                    args.push(*value);
351                    arg_offsets.push(UNMATERIALIZED);
352                }
353            }
354            bits.push(cell.bits() as u16);
355        }
356
357        let meta = match cache {
358            ConstCacheType::Transcript => derive_opcode_metadata(
359                &opcode,
360                &args,
361                &bits,
362                self.lookup_bits,
363                self.transcript_bb_constant_cache.keys(),
364            ),
365            ConstCacheType::BabyBear => derive_opcode_metadata(
366                &opcode,
367                &args,
368                &bits,
369                self.lookup_bits,
370                self.bb_const_cache.keys(),
371            ),
372            // `Context::load_zero` is the only gate-level constant cache.
373            ConstCacheType::Gate => {
374                let zero = self.zero_cell.map(|_| Fr::ZERO);
375                derive_opcode_metadata(&opcode, &args, &bits, self.lookup_bits, zero.iter())
376            }
377        };
378
379        let ctx_offset = self.ctx_offset;
380        let output_offsets: Vec<usize> = meta
381            .output_offsets
382            .iter()
383            .map(|&rel| {
384                if rel == UNMATERIALIZED {
385                    UNMATERIALIZED
386                } else {
387                    ctx_offset + rel
388                }
389            })
390            .collect();
391        self.node_meta.push(NodeMeta {
392            ctx_len: meta.ctx_len,
393            ctx_offset,
394            lookups_len: meta.lookups_len,
395            lookup_offset: self.lookup_offset,
396            constant_skip_inds: meta.constant_skip_inds,
397            arg_offsets,
398            level,
399        });
400        self.ctx_offset += meta.ctx_len;
401        self.lookup_offset += meta.lookups_len;
402
403        self.nodes.push(Halo2GraphNode {
404            opcode,
405            operands,
406            id,
407        });
408        (id, output_offsets)
409    }
410
411    fn emit1(
412        &mut self,
413        opcode: Halo2Opcode,
414        operands: Vec<GraphCell>,
415        bits: usize,
416        cache: ConstCacheType,
417    ) -> GraphCell {
418        let (id, offsets) = self.emit(opcode, operands, cache);
419        GraphCell::Cell(id, offsets[0], bits as u16)
420    }
421
422    /// Mirrors `Context::load_zero`.
423    fn load_zero(&mut self) -> GraphCell {
424        if let Some(zero) = self.zero_cell {
425            return zero;
426        }
427        let zero = self.emit1(
428            Halo2Opcode::Const,
429            vec![GraphCell::Const(Fr::ZERO)],
430            0,
431            ConstCacheType::BabyBear,
432        );
433        self.zero_cell = Some(zero);
434        self.bb_const_cache.insert(Fr::ZERO, bb_wire(zero, 0));
435        self.transcript_bb_constant_cache
436            .insert(Fr::ZERO, bb_wire(zero, 0));
437        zero
438    }
439
440    /// Registers a BabyBear constant loaded *inside* an atomic node
441    /// (`BBDiv`/`ExtMul`/`ExtDiv`) at absolute advice `offset`. Mirrors
442    /// `BabyBearChip::load_constant` cache/zero-cell behavior.
443    fn note_internal_bb_const(&mut self, id: NodeId, offset: usize, value: BabyBear) {
444        let key_u64 = value.as_canonical_u64();
445        let key = Fr::from(key_u64);
446        if self.bb_const_cache.contains_key(&key) {
447            return;
448        }
449        let max_bits = bit_length(key_u64);
450        let (cell, wire_bits) = if key_u64 == 0 {
451            let zero = match self.zero_cell {
452                Some(zero) => zero,
453                None => {
454                    debug_assert_ne!(offset, UNMATERIALIZED);
455                    let zero = GraphCell::Cell(id, offset, 0);
456                    self.zero_cell = Some(zero);
457                    zero
458                }
459            };
460            (zero, 0)
461        } else {
462            debug_assert_ne!(offset, UNMATERIALIZED);
463            (GraphCell::Cell(id, offset, max_bits as u16), max_bits)
464        };
465        self.bb_const_cache.insert(key, bb_wire(cell, wire_bits));
466    }
467
468    fn bb_reduce_wire(&mut self, a: BabyBearWire<GraphCell>) -> BabyBearWire<GraphCell> {
469        assert!(a.max_bits <= REDUCE_THRESHOLD);
470        let cell = self.emit1(
471            Halo2Opcode::BBReduce,
472            vec![a.value],
473            BABYBEAR_MAX_BITS,
474            ConstCacheType::BabyBear,
475        );
476        bb_wire(cell, BABYBEAR_MAX_BITS)
477    }
478
479    fn inner_product(&mut self, values: &[GraphCell], coeffs: &[Fr]) -> GraphCell {
480        assert_eq!(values.len(), coeffs.len());
481        let operands = values
482            .iter()
483            .zip(coeffs)
484            .flat_map(|(&value, &coeff)| [value, GraphCell::Const(coeff)])
485            .collect_vec();
486        self.emit1(
487            Halo2Opcode::InnerProduct(values.len() as u16),
488            operands,
489            RAW_MAX_BITS,
490            ConstCacheType::BabyBear,
491        )
492    }
493
494    /// Base-2^31 packing of reduced BabyBear wires, mirroring `pack_base_2_31_cells`.
495    fn pack_base_2_31(&mut self, values: &[ReducedBabyBearWire<GraphCell>]) -> GraphCell {
496        assert!(values.len() <= MULTI_FIELD32_NUM_F_ELMS);
497        let base = Fr::from(1u64 << 31);
498        let coeffs = iter::successors(Some(Fr::ONE), |power| Some(*power * base))
499            .take(values.len())
500            .collect_vec();
501        let operands = values.iter().map(|value| value.value()).collect_vec();
502        self.inner_product(&operands, &coeffs)
503    }
504
505    fn permute_t3(&mut self, state: &mut [GraphCell; POSEIDON2_WIDTH]) {
506        let (id, offsets) = self.emit(
507            Halo2Opcode::PoseidonPermute2T3,
508            state.to_vec(),
509            ConstCacheType::BabyBear,
510        );
511        *state = array::from_fn(|i| GraphCell::Cell(id, offsets[i], RAW_MAX_BITS as u16));
512    }
513
514    // --- transcript internals mirroring `TranscriptChip` ---
515
516    fn take_transcript(&mut self) -> IrTranscript {
517        self.transcript
518            .take()
519            .expect("transcript not initialized; call init_transcript first")
520    }
521
522    fn sponge_absorb(&mut self, t: &mut IrTranscript, value: GraphCell) {
523        t.sponge_state[t.absorb_idx] = value;
524        t.absorb_idx += 1;
525        if t.absorb_idx == POSEIDON2_RATE {
526            self.permute_t3(&mut t.sponge_state);
527            t.absorb_idx = 0;
528            t.sample_idx = POSEIDON2_RATE;
529        }
530    }
531
532    fn sponge_squeeze(&mut self, t: &mut IrTranscript) -> GraphCell {
533        if t.absorb_idx != 0 || t.sample_idx == 0 {
534            self.permute_t3(&mut t.sponge_state);
535            t.absorb_idx = 0;
536            t.sample_idx = POSEIDON2_RATE;
537        }
538        t.sample_idx -= 1;
539        t.sponge_state[t.sample_idx]
540    }
541
542    fn flush_observe_buf(&mut self, t: &mut IrTranscript) {
543        if !t.observe_buf.is_empty() {
544            let packed = self.pack_base_2_31(&t.observe_buf);
545            self.sponge_absorb(t, packed);
546            t.observe_buf.clear();
547        }
548    }
549
550    fn observe_inner(&mut self, t: &mut IrTranscript, value: &ReducedBabyBearWire<GraphCell>) {
551        t.sample_buf.clear();
552        t.observe_buf.push(*value);
553        if t.observe_buf.len() == NUM_OBS_PER_WORD {
554            self.flush_observe_buf(t);
555        }
556    }
557
558    fn sample_inner(&mut self, t: &mut IrTranscript) -> BabyBearWire<GraphCell> {
559        if let Some(val) = t.sample_buf.pop() {
560            return val;
561        }
562        self.flush_observe_buf(t);
563        let squeezed = self.sponge_squeeze(t);
564        let (id, offsets) = self.emit(
565            Halo2Opcode::DecomposeBn254ToBabyBear,
566            vec![squeezed],
567            ConstCacheType::BabyBear,
568        );
569        t.sample_buf = (0..NUM_SAMPLES_PER_WORD)
570            .map(|i| {
571                bb_wire(
572                    GraphCell::Cell(id, offsets[i], BABYBEAR_MAX_BITS as u16),
573                    BABYBEAR_MAX_BITS,
574                )
575            })
576            .collect();
577        // Reverse so pop() returns digits in order (b_0 first).
578        t.sample_buf.reverse();
579        t.sample_buf.pop().expect("sample_buf should be non-empty")
580    }
581
582    fn sample_bits_inner(&mut self, t: &mut IrTranscript, bits: usize) -> GraphCell {
583        assert!(
584            bits < (u32::BITS as usize),
585            "sample_bits requires bits < 32: {bits}"
586        );
587        assert!(
588            (1u64 << bits) < BABY_BEAR_MODULUS_U64,
589            "sample_bits requires (1 << bits) < modulus: bits={bits}"
590        );
591        let sampled = self.sample_inner(t);
592        if bits == 0 {
593            return self.load_zero();
594        }
595        self.emit1(
596            Halo2Opcode::RangeDiv(bits as u16),
597            vec![sampled.value],
598            bits,
599            ConstCacheType::BabyBear,
600        )
601    }
602
603    // --- graph statistics (dev diagnostics only) ---
604
605    #[cfg(test)]
606    pub fn stats(&self) -> IrStats {
607        let mut level_widths: Vec<usize> = Vec::new();
608        for level in self.node_meta.iter().map(|meta| meta.level) {
609            let level = level as usize;
610            if level >= level_widths.len() {
611                level_widths.resize(level + 1, 0);
612            }
613            level_widths[level] += 1;
614        }
615        let mut counts: HashMap<&'static str, usize> = HashMap::new();
616        for node in &self.nodes {
617            *counts.entry(node.opcode.name()).or_default() += 1;
618        }
619        let mut opcode_counts = counts.into_iter().collect_vec();
620        opcode_counts.sort_by(|a, b| b.1.cmp(&a.1).then(a.0.cmp(b.0)));
621        let num_levels = level_widths.len();
622        IrStats {
623            num_nodes: self.nodes.len(),
624            num_inputs: self.input_values.len(),
625            num_levels,
626            max_width: level_widths.iter().copied().max().unwrap_or(0),
627            avg_width: self.nodes.len() as f64 / num_levels.max(1) as f64,
628            opcode_counts,
629            level_widths,
630        }
631    }
632}
633
634/// Structural statistics of the generated IR graph.
635#[cfg(test)]
636#[derive(Clone, Debug)]
637pub struct IrStats {
638    pub num_nodes: usize,
639    pub num_inputs: usize,
640    /// Maximum dataflow depth (number of levels).
641    pub num_levels: usize,
642    /// Maximum number of nodes on any single level.
643    pub max_width: usize,
644    /// Average number of nodes per level.
645    pub avg_width: f64,
646    /// Opcode distribution, sorted by descending count.
647    pub opcode_counts: Vec<(&'static str, usize)>,
648    /// Number of nodes at each level.
649    pub level_widths: Vec<usize>,
650}
651
652#[cfg(test)]
653impl core::fmt::Display for IrStats {
654    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
655        writeln!(f, "nodes: {}", self.num_nodes)?;
656        writeln!(f, "proof-input witnesses: {}", self.num_inputs)?;
657        writeln!(f, "max depth (levels): {}", self.num_levels)?;
658        writeln!(f, "max width: {}", self.max_width)?;
659        writeln!(f, "avg width per level: {:.1}", self.avg_width)?;
660        writeln!(f, "opcode distribution:")?;
661        for (name, count) in &self.opcode_counts {
662            writeln!(
663                f,
664                "  {name:<22} {count:>10} ({:.2}%)",
665                100.0 * *count as f64 / self.num_nodes as f64
666            )?;
667        }
668        Ok(())
669    }
670}
671
672// --- static bit-bound simulation of atomic extension ops ---
673
674fn w_bits() -> usize {
675    bit_length(<BabyBear as BinomiallyExtendable<4>>::W.as_canonical_u64())
676}
677
678/// Mirrors the `max_bits` and reduce logic of `BabyBearChip::mul_add`.
679fn mul_add_result_bits(mut a: usize, mut b: usize, mut c: usize) -> usize {
680    if a < b {
681        core::mem::swap(&mut a, &mut b);
682    }
683    if a + b + 1 > REDUCE_THRESHOLD {
684        a = BABYBEAR_MAX_BITS;
685        if a + b + 1 > REDUCE_THRESHOLD {
686            b = BABYBEAR_MAX_BITS;
687        }
688    }
689    if c + 1 > REDUCE_THRESHOLD {
690        c = BABYBEAR_MAX_BITS;
691    }
692    c.max(a + b) + 1
693}
694
695/// Mirrors the `max_bits` and reduce logic of `BabyBearChip::special_inner_product`,
696/// mutating the bit bounds like the chip mutates the wires.
697fn special_inner_product_bits(a: &mut [usize; 4], b: &mut [usize; 4], s: usize) -> usize {
698    let mut max_bits = 0;
699    let lb = s.saturating_sub(3);
700    let ub = 4.min(s + 1);
701    let len = if s < 3 { s + 1 } else { 7 - s };
702    for (i, (ci, di)) in (lb..ub).zip(((s + 1 - ub)..(s + 1 - lb)).rev()).enumerate() {
703        let limit = REDUCE_THRESHOLD - len + i;
704        if a[ci] + b[di] > limit {
705            if a[ci] >= b[di] {
706                a[ci] = BABYBEAR_MAX_BITS;
707                if a[ci] + b[di] > limit {
708                    b[di] = BABYBEAR_MAX_BITS;
709                }
710            } else {
711                b[di] = BABYBEAR_MAX_BITS;
712                if a[ci] + b[di] > limit {
713                    a[ci] = BABYBEAR_MAX_BITS;
714                }
715            }
716        }
717        max_bits = if i == 0 {
718            a[ci] + b[di]
719        } else {
720            max_bits.max(a[ci] + b[di]) + 1
721        };
722    }
723    max_bits
724}
725
726/// Mirrors the `max_bits` bookkeeping of `BabyBearExt4Chip::mul`.
727fn ext_mul_result_bits(mut a: [usize; 4], mut b: [usize; 4]) -> [usize; 4] {
728    let mut coeffs = Vec::with_capacity(7);
729    for s in 0..7 {
730        coeffs.push(special_inner_product_bits(&mut a, &mut b, s));
731    }
732    let w = w_bits();
733    for i in 4..7 {
734        coeffs[i - 4] = mul_add_result_bits(coeffs[i], w, coeffs[i - 4]);
735    }
736    [coeffs[0], coeffs[1], coeffs[2], coeffs[3]]
737}
738
739impl ChipBase for Halo2IRBuilder {
740    type F = GraphCell;
741}
742
743impl PopulateInputs for Halo2IRBuilder {
744    fn load_witness(&mut self, value: Fr) -> GraphCell {
745        self.input_values.push(value);
746        self.emit1(
747            Halo2Opcode::LoadWitness,
748            vec![],
749            RAW_MAX_BITS,
750            ConstCacheType::BabyBear,
751        )
752    }
753
754    fn bb_load_reduced_witness(&mut self, value: BabyBear) -> ReducedBabyBearWire<GraphCell> {
755        self.input_values.push(Fr::from(value.as_canonical_u64()));
756        let cell = self.emit1(
757            Halo2Opcode::LoadWitness,
758            vec![],
759            BABYBEAR_MAX_BITS,
760            ConstCacheType::BabyBear,
761        );
762        self.emit(
763            Halo2Opcode::CheckLessThanSafe,
764            vec![cell],
765            ConstCacheType::BabyBear,
766        );
767        ReducedBabyBearWire::assume_reduced(bb_wire(cell, BABYBEAR_MAX_BITS))
768    }
769
770    fn ext_load_reduced_witness(
771        &mut self,
772        value: BabyBearExt4,
773    ) -> ReducedBabyBearExt4Wire<GraphCell> {
774        let coeffs = value.as_basis_coefficients_slice();
775        ReducedBabyBearExt4Wire::assume_reduced(array::from_fn(|i| {
776            self.bb_load_reduced_witness(coeffs[i])
777        }))
778    }
779}
780
781impl GateInst for Halo2IRBuilder {
782    /// Mirrors `ctx.load_constant`: always assigns a fresh cell (no cache).
783    fn load_constant(&mut self, value: Fr) -> GraphCell {
784        let bits = fr_bits(&value);
785        self.emit1(
786            Halo2Opcode::Const,
787            vec![GraphCell::Const(value)],
788            bits,
789            ConstCacheType::Gate,
790        )
791    }
792
793    /// Copy constraint only; assigns no advice cells, so no node is emitted.
794    fn constrain_equal(&mut self, _a: GraphCell, _b: GraphCell) {}
795
796    fn select(
797        &mut self,
798        when_true: GraphCell,
799        when_false: GraphCell,
800        cond: GraphCell,
801    ) -> GraphCell {
802        let bits = when_true.bits().max(when_false.bits());
803        self.emit1(
804            Halo2Opcode::Select,
805            vec![when_true, when_false, cond],
806            bits,
807            ConstCacheType::Gate,
808        )
809    }
810
811    fn select_const(&mut self, when_true: Fr, when_false: Fr, cond: GraphCell) -> GraphCell {
812        self.select(
813            GraphCell::Const(when_true),
814            GraphCell::Const(when_false),
815            cond,
816        )
817    }
818
819    fn num_to_bits(&mut self, a: GraphCell, range_bits: usize) -> Vec<GraphCell> {
820        let (id, offsets) = self.emit(
821            Halo2Opcode::Num2Bits(range_bits as u16),
822            vec![a],
823            ConstCacheType::Gate,
824        );
825        (0..range_bits)
826            .map(|i| GraphCell::Cell(id, offsets[i], 1))
827            .collect()
828    }
829
830    fn inner_product_const(&mut self, values: &[GraphCell], coeffs: &[Fr]) -> GraphCell {
831        self.inner_product(values, coeffs)
832    }
833
834    fn cell_count(&self) -> usize {
835        self.nodes.len()
836    }
837}
838
839impl BabyBearInst for Halo2IRBuilder {
840    /// Mirrors `BabyBearChip::load_constant` (const cache + `load_zero` for zero).
841    fn bb_load_constant(&mut self, value: BabyBear) -> BabyBearWire<GraphCell> {
842        let key_u64 = value.as_canonical_u64();
843        let key = Fr::from(key_u64);
844        if let Some(&cached) = self.bb_const_cache.get(&key) {
845            return cached;
846        }
847        let max_bits = bit_length(key_u64);
848        let cell = if key_u64 == 0 {
849            self.load_zero()
850        } else {
851            self.emit1(
852                Halo2Opcode::Const,
853                vec![GraphCell::Const(key)],
854                max_bits,
855                ConstCacheType::BabyBear,
856            )
857        };
858        let wire = bb_wire(cell, max_bits);
859        self.bb_const_cache.insert(key, wire);
860        wire
861    }
862
863    fn bb_load_reduced_constant(&mut self, value: BabyBear) -> ReducedBabyBearWire<GraphCell> {
864        // Constants are canonical by construction.
865        ReducedBabyBearWire::assume_reduced(self.bb_load_constant(value))
866    }
867
868    fn bb_reduce(&mut self, a: BabyBearWire<GraphCell>) -> BabyBearWire<GraphCell> {
869        self.bb_reduce_wire(a)
870    }
871
872    fn bb_reduce_max_bits(&mut self, a: BabyBearWire<GraphCell>) -> BabyBearWire<GraphCell> {
873        if a.max_bits > BABYBEAR_MAX_BITS {
874            self.bb_reduce_wire(a)
875        } else {
876            a
877        }
878    }
879
880    fn bb_add(
881        &mut self,
882        mut a: BabyBearWire<GraphCell>,
883        mut b: BabyBearWire<GraphCell>,
884    ) -> BabyBearWire<GraphCell> {
885        if a.max_bits + 1 > REDUCE_THRESHOLD {
886            a = self.bb_reduce_wire(a);
887        }
888        if b.max_bits + 1 > REDUCE_THRESHOLD {
889            b = self.bb_reduce_wire(b);
890        }
891        let max_bits = a.max_bits.max(b.max_bits) + 1;
892        let cell = self.emit1(
893            Halo2Opcode::BBAdd,
894            vec![a.value, b.value],
895            max_bits,
896            ConstCacheType::BabyBear,
897        );
898        bb_wire(cell, max_bits)
899    }
900
901    fn bb_neg(&mut self, a: BabyBearWire<GraphCell>) -> BabyBearWire<GraphCell> {
902        let cell = self.emit1(
903            Halo2Opcode::BBNeg,
904            vec![a.value],
905            a.max_bits,
906            ConstCacheType::BabyBear,
907        );
908        bb_wire(cell, a.max_bits)
909    }
910
911    fn bb_sub(
912        &mut self,
913        mut a: BabyBearWire<GraphCell>,
914        mut b: BabyBearWire<GraphCell>,
915    ) -> BabyBearWire<GraphCell> {
916        if a.max_bits + 1 > REDUCE_THRESHOLD {
917            a = self.bb_reduce_wire(a);
918        }
919        if b.max_bits + 1 > REDUCE_THRESHOLD {
920            b = self.bb_reduce_wire(b);
921        }
922        let max_bits = a.max_bits.max(b.max_bits) + 1;
923        let cell = self.emit1(
924            Halo2Opcode::BBSub,
925            vec![a.value, b.value],
926            max_bits,
927            ConstCacheType::BabyBear,
928        );
929        bb_wire(cell, max_bits)
930    }
931
932    fn bb_mul(
933        &mut self,
934        mut a: BabyBearWire<GraphCell>,
935        mut b: BabyBearWire<GraphCell>,
936    ) -> BabyBearWire<GraphCell> {
937        if a.max_bits < b.max_bits {
938            core::mem::swap(&mut a, &mut b);
939        }
940        if a.max_bits + b.max_bits > REDUCE_THRESHOLD {
941            a = self.bb_reduce_wire(a);
942            if a.max_bits + b.max_bits > REDUCE_THRESHOLD {
943                b = self.bb_reduce_wire(b);
944            }
945        }
946        let max_bits = a.max_bits + b.max_bits;
947        let cell = self.emit1(
948            Halo2Opcode::BBMul,
949            vec![a.value, b.value],
950            max_bits,
951            ConstCacheType::BabyBear,
952        );
953        bb_wire(cell, max_bits)
954    }
955
956    fn bb_mul_add(
957        &mut self,
958        mut a: BabyBearWire<GraphCell>,
959        mut b: BabyBearWire<GraphCell>,
960        mut c: BabyBearWire<GraphCell>,
961    ) -> BabyBearWire<GraphCell> {
962        if a.max_bits < b.max_bits {
963            core::mem::swap(&mut a, &mut b);
964        }
965        if a.max_bits + b.max_bits + 1 > REDUCE_THRESHOLD {
966            a = self.bb_reduce_wire(a);
967            if a.max_bits + b.max_bits + 1 > REDUCE_THRESHOLD {
968                b = self.bb_reduce_wire(b);
969            }
970        }
971        if c.max_bits + 1 > REDUCE_THRESHOLD {
972            c = self.bb_reduce_wire(c);
973        }
974        let max_bits = c.max_bits.max(a.max_bits + b.max_bits) + 1;
975        let cell = self.emit1(
976            Halo2Opcode::BBMulAdd,
977            vec![a.value, b.value, c.value],
978            max_bits,
979            ConstCacheType::BabyBear,
980        );
981        bb_wire(cell, max_bits)
982    }
983
984    fn bb_div(
985        &mut self,
986        a: BabyBearWire<GraphCell>,
987        b: BabyBearWire<GraphCell>,
988    ) -> BabyBearWire<GraphCell> {
989        let (id, offsets) = self.emit(
990            Halo2Opcode::BBDiv,
991            vec![a.value, b.value],
992            ConstCacheType::BabyBear,
993        );
994        // `BabyBearChip::div` internally loads the ONE constant (output 1 when uncached).
995        self.note_internal_bb_const(id, offsets[1], BabyBear::new(1));
996        bb_wire(
997            GraphCell::Cell(id, offsets[0], BABYBEAR_MAX_BITS as u16),
998            BABYBEAR_MAX_BITS,
999        )
1000    }
1001
1002    fn bb_assert_zero(&mut self, a: BabyBearWire<GraphCell>) {
1003        assert!(a.max_bits <= REDUCE_THRESHOLD);
1004        self.emit(
1005            Halo2Opcode::BBAssertZero,
1006            vec![a.value],
1007            ConstCacheType::BabyBear,
1008        );
1009    }
1010
1011    fn bb_assert_equal(&mut self, a: BabyBearWire<GraphCell>, b: BabyBearWire<GraphCell>) {
1012        let diff = self.bb_sub(a, b);
1013        self.bb_assert_zero(diff);
1014    }
1015
1016    fn bb_zero(&mut self) -> BabyBearWire<GraphCell> {
1017        self.bb_load_constant(BabyBear::new(0))
1018    }
1019
1020    fn bb_one(&mut self) -> BabyBearWire<GraphCell> {
1021        self.bb_load_constant(BabyBear::new(1))
1022    }
1023
1024    fn bb_mul_const(&mut self, a: BabyBearWire<GraphCell>, c: BabyBear) -> BabyBearWire<GraphCell> {
1025        let c_wire = self.bb_load_constant(c);
1026        self.bb_mul(a, c_wire)
1027    }
1028
1029    fn bb_square(&mut self, a: BabyBearWire<GraphCell>) -> BabyBearWire<GraphCell> {
1030        self.bb_mul(a, a)
1031    }
1032
1033    fn bb_pow_power_of_two(
1034        &mut self,
1035        a: BabyBearWire<GraphCell>,
1036        n: usize,
1037    ) -> BabyBearWire<GraphCell> {
1038        let mut result = a;
1039        for _ in 0..n {
1040            result = self.bb_square(result);
1041        }
1042        result
1043    }
1044}
1045
1046impl BabyBearExt4Inst for Halo2IRBuilder {
1047    fn ext_load_constant(&mut self, value: BabyBearExt4) -> BabyBearExt4Wire<GraphCell> {
1048        let coeffs = value.as_basis_coefficients_slice();
1049        BabyBearExt4Wire(array::from_fn(|i| self.bb_load_constant(coeffs[i])))
1050    }
1051
1052    fn ext_load_reduced_constant(
1053        &mut self,
1054        value: BabyBearExt4,
1055    ) -> ReducedBabyBearExt4Wire<GraphCell> {
1056        let coeffs = value.as_basis_coefficients_slice();
1057        ReducedBabyBearExt4Wire::assume_reduced(array::from_fn(|i| {
1058            self.bb_load_reduced_constant(coeffs[i])
1059        }))
1060    }
1061
1062    fn ext_add(
1063        &mut self,
1064        a: BabyBearExt4Wire<GraphCell>,
1065        b: BabyBearExt4Wire<GraphCell>,
1066    ) -> BabyBearExt4Wire<GraphCell> {
1067        BabyBearExt4Wire(array::from_fn(|i| self.bb_add(a.0[i], b.0[i])))
1068    }
1069
1070    fn ext_neg(&mut self, a: BabyBearExt4Wire<GraphCell>) -> BabyBearExt4Wire<GraphCell> {
1071        BabyBearExt4Wire(array::from_fn(|i| self.bb_neg(a.0[i])))
1072    }
1073
1074    fn ext_sub(
1075        &mut self,
1076        a: BabyBearExt4Wire<GraphCell>,
1077        b: BabyBearExt4Wire<GraphCell>,
1078    ) -> BabyBearExt4Wire<GraphCell> {
1079        BabyBearExt4Wire(array::from_fn(|i| self.bb_sub(a.0[i], b.0[i])))
1080    }
1081
1082    fn ext_scalar_mul(
1083        &mut self,
1084        a: BabyBearExt4Wire<GraphCell>,
1085        b: BabyBearWire<GraphCell>,
1086    ) -> BabyBearExt4Wire<GraphCell> {
1087        BabyBearExt4Wire(array::from_fn(|i| self.bb_mul(a.0[i], b)))
1088    }
1089
1090    fn ext_scalar_mul_add(
1091        &mut self,
1092        a: BabyBearExt4Wire<GraphCell>,
1093        b: BabyBearWire<GraphCell>,
1094        c: BabyBearExt4Wire<GraphCell>,
1095    ) -> BabyBearExt4Wire<GraphCell> {
1096        BabyBearExt4Wire(array::from_fn(|i| self.bb_mul_add(a.0[i], b, c.0[i])))
1097    }
1098
1099    fn ext_assert_zero(&mut self, a: BabyBearExt4Wire<GraphCell>) {
1100        for x in a.0 {
1101            self.bb_assert_zero(x);
1102        }
1103    }
1104
1105    fn ext_assert_equal(&mut self, a: BabyBearExt4Wire<GraphCell>, b: BabyBearExt4Wire<GraphCell>) {
1106        for (a, b) in a.0.into_iter().zip(b.0) {
1107            self.bb_assert_equal(a, b);
1108        }
1109    }
1110
1111    fn ext_mul(
1112        &mut self,
1113        a: BabyBearExt4Wire<GraphCell>,
1114        b: BabyBearExt4Wire<GraphCell>,
1115    ) -> BabyBearExt4Wire<GraphCell> {
1116        let operands = a.0.iter().chain(b.0.iter()).map(|w| w.value).collect_vec();
1117        let (id, offsets) = self.emit(Halo2Opcode::ExtMul, operands, ConstCacheType::BabyBear);
1118        // `BabyBearExt4Chip::mul` internally loads the W constant (output 4 when uncached).
1119        self.note_internal_bb_const(id, offsets[4], <BabyBear as BinomiallyExtendable<4>>::W);
1120        let bits = ext_mul_result_bits(a.0.map(|w| w.max_bits), b.0.map(|w| w.max_bits));
1121        BabyBearExt4Wire(array::from_fn(|i| {
1122            bb_wire(GraphCell::Cell(id, offsets[i], bits[i] as u16), bits[i])
1123        }))
1124    }
1125
1126    fn ext_div(
1127        &mut self,
1128        a: BabyBearExt4Wire<GraphCell>,
1129        b: BabyBearExt4Wire<GraphCell>,
1130    ) -> BabyBearExt4Wire<GraphCell> {
1131        let operands = a.0.iter().chain(b.0.iter()).map(|w| w.value).collect_vec();
1132        let (id, offsets) = self.emit(Halo2Opcode::ExtDiv, operands, ConstCacheType::BabyBear);
1133        // `BabyBearExt4Chip::div` internally loads ext ONE = bb [1, 0, 0, 0] (outputs 4/5
1134        // when uncached) and, via the internal ext mul, the W constant (output 6).
1135        self.note_internal_bb_const(id, offsets[4], BabyBear::new(1));
1136        self.note_internal_bb_const(id, offsets[5], BabyBear::new(0));
1137        self.note_internal_bb_const(id, offsets[6], <BabyBear as BinomiallyExtendable<4>>::W);
1138        // The quotient is loaded as an ext witness: each coefficient is 31 bits.
1139        BabyBearExt4Wire(array::from_fn(|i| {
1140            bb_wire(
1141                GraphCell::Cell(id, offsets[i], BABYBEAR_MAX_BITS as u16),
1142                BABYBEAR_MAX_BITS,
1143            )
1144        }))
1145    }
1146
1147    fn ext_reduce_max_bits(
1148        &mut self,
1149        a: BabyBearExt4Wire<GraphCell>,
1150    ) -> BabyBearExt4Wire<GraphCell> {
1151        BabyBearExt4Wire(array::from_fn(|i| self.bb_reduce_max_bits(a.0[i])))
1152    }
1153
1154    fn ext_zero(&mut self) -> BabyBearExt4Wire<GraphCell> {
1155        self.ext_from_base_const(BabyBear::new(0))
1156    }
1157
1158    fn ext_from_base_const(&mut self, value: BabyBear) -> BabyBearExt4Wire<GraphCell> {
1159        let base_val = self.bb_load_constant(value);
1160        let z = self.bb_load_constant(BabyBear::new(0));
1161        BabyBearExt4Wire([base_val, z, z, z])
1162    }
1163
1164    fn ext_from_base_var(&mut self, value: BabyBearWire<GraphCell>) -> BabyBearExt4Wire<GraphCell> {
1165        let z = self.bb_load_constant(BabyBear::new(0));
1166        BabyBearExt4Wire([value, z, z, z])
1167    }
1168
1169    fn ext_mul_base_const(
1170        &mut self,
1171        a: BabyBearExt4Wire<GraphCell>,
1172        c: BabyBear,
1173    ) -> BabyBearExt4Wire<GraphCell> {
1174        let c_wire = self.bb_load_constant(c);
1175        self.ext_scalar_mul(a, c_wire)
1176    }
1177
1178    fn ext_square(&mut self, a: BabyBearExt4Wire<GraphCell>) -> BabyBearExt4Wire<GraphCell> {
1179        self.ext_mul(a, a)
1180    }
1181
1182    fn ext_pow_power_of_two(
1183        &mut self,
1184        a: BabyBearExt4Wire<GraphCell>,
1185        n: usize,
1186    ) -> BabyBearExt4Wire<GraphCell> {
1187        let mut result = a;
1188        for _ in 0..n {
1189            result = self.ext_square(result);
1190        }
1191        result
1192    }
1193}
1194
1195impl Poseidon2Inst for Halo2IRBuilder {
1196    /// Mirrors `hash::poseidon2::hash_babybear_slice_to_digest`.
1197    fn hash_babybear_slice_to_digest(
1198        &mut self,
1199        values: &[ReducedBabyBearWire<GraphCell>],
1200    ) -> GraphCell {
1201        let zero = self.load_zero();
1202        let mut state = [zero; POSEIDON2_WIDTH];
1203        for block_chunk in values.chunks(MULTI_FIELD32_RATE) {
1204            for (chunk_id, chunk) in block_chunk.chunks(MULTI_FIELD32_NUM_F_ELMS).enumerate() {
1205                state[chunk_id] = self.pack_base_2_31(chunk);
1206            }
1207            self.permute_t3(&mut state);
1208        }
1209        state[0]
1210    }
1211
1212    fn compress_digests(&mut self, left: GraphCell, right: GraphCell) -> GraphCell {
1213        let (id, offsets) = self.emit(
1214            Halo2Opcode::PoseidonPermute2T2,
1215            vec![left, right],
1216            ConstCacheType::BabyBear,
1217        );
1218        GraphCell::Cell(id, offsets[0], RAW_MAX_BITS as u16)
1219    }
1220}
1221
1222impl TranscriptInst for Halo2IRBuilder {
1223    fn init_transcript(&mut self) {
1224        let zero = self.load_zero();
1225        self.transcript = Some(IrTranscript {
1226            sponge_state: [zero; POSEIDON2_WIDTH],
1227            absorb_idx: 0,
1228            sample_idx: 0,
1229            observe_buf: Vec::with_capacity(NUM_OBS_PER_WORD),
1230            sample_buf: Vec::with_capacity(NUM_SAMPLES_PER_WORD),
1231        });
1232    }
1233
1234    fn observe(&mut self, value: &ReducedBabyBearWire<GraphCell>) {
1235        let mut t = self.take_transcript();
1236        self.observe_inner(&mut t, value);
1237        self.transcript = Some(t);
1238    }
1239
1240    fn observe_ext(&mut self, value: &ReducedBabyBearExt4Wire<GraphCell>) {
1241        let mut t = self.take_transcript();
1242        for coeff in value.coeffs() {
1243            self.observe_inner(&mut t, coeff);
1244        }
1245        self.transcript = Some(t);
1246    }
1247
1248    fn observe_commit(&mut self, digest: &DigestWire<GraphCell>) {
1249        let mut t = self.take_transcript();
1250        t.sample_buf.clear();
1251        self.flush_observe_buf(&mut t);
1252        for &elem in &digest.elems {
1253            self.sponge_absorb(&mut t, elem);
1254        }
1255        self.transcript = Some(t);
1256    }
1257
1258    fn sample(&mut self) -> BabyBearWire<GraphCell> {
1259        let mut t = self.take_transcript();
1260        let out = self.sample_inner(&mut t);
1261        self.transcript = Some(t);
1262        out
1263    }
1264
1265    fn sample_ext(&mut self) -> BabyBearExt4Wire<GraphCell> {
1266        let mut t = self.take_transcript();
1267        let coeffs = array::from_fn(|_| self.sample_inner(&mut t));
1268        self.transcript = Some(t);
1269        BabyBearExt4Wire(coeffs)
1270    }
1271
1272    fn sample_bits(&mut self, bits: usize) -> GraphCell {
1273        let mut t = self.take_transcript();
1274        let out = self.sample_bits_inner(&mut t, bits);
1275        self.transcript = Some(t);
1276        out
1277    }
1278
1279    fn check_witness(&mut self, bits: usize, witness: &ReducedBabyBearWire<GraphCell>) {
1280        if bits == 0 {
1281            return;
1282        }
1283        let mut t = self.take_transcript();
1284        self.observe_inner(&mut t, witness);
1285        // `assert_is_const(sampled_bits, 0)` is a copy constraint to a fixed cell and
1286        // assigns no advice cells, so only the sample_bits cells appear in the tape.
1287        let _sampled_bits = self.sample_bits_inner(&mut t, bits);
1288        self.transcript = Some(t);
1289    }
1290
1291    fn transcript_load_reduced_constant(
1292        &mut self,
1293        value: BabyBear,
1294    ) -> ReducedBabyBearWire<Self::F> {
1295        let key_u64 = value.as_canonical_u64();
1296        let key = Fr::from(key_u64);
1297        if let Some(&cached) = self.transcript_bb_constant_cache.get(&key) {
1298            return ReducedBabyBearWire::assume_reduced(cached);
1299        }
1300        let max_bits = bit_length(key_u64);
1301        let cell = if key_u64 == 0 {
1302            self.load_zero()
1303        } else {
1304            self.emit1(
1305                Halo2Opcode::Const,
1306                vec![GraphCell::Const(key)],
1307                max_bits,
1308                ConstCacheType::Transcript,
1309            )
1310        };
1311        let wire = bb_wire(cell, max_bits);
1312        self.transcript_bb_constant_cache.insert(key, wire);
1313        ReducedBabyBearWire::assume_reduced(wire)
1314    }
1315}
1316
1317#[cfg(test)]
1318mod tests {
1319    use super::*;
1320    use crate::test_fixtures::{fixture_circuit_and_proof, FIXTURE_K};
1321
1322    #[test]
1323    fn ir_generation_stats_for_fixture_proof() {
1324        let (circuit, proof) = fixture_circuit_and_proof();
1325
1326        let mut builder = Halo2IRBuilder::new(FIXTURE_K - 1);
1327        circuit.populate_verify_stark_constraints(&mut builder, &proof);
1328
1329        let stats = builder.stats();
1330        println!("=== static verifier IR stats ===");
1331        print!("{stats}");
1332        println!("ctx cells: {}", builder.total_ctx_len());
1333        println!("lookup cells: {}", builder.total_lookups_len());
1334        assert!(stats.num_nodes > 0);
1335        assert_eq!(builder.nodes.len(), builder.node_meta.len());
1336    }
1337}
openvm_static_verifier/tracegen/ir_builder.rs

openvm_static_verifier/tracegen/
ir_builder.rs
