Skip to main content

hekate_program/
outer.rs

1// SPDX-FileCopyrightText: 2026 Andrei Kochergin <andrei@oumuamua.dev>
2// SPDX-FileCopyrightText: 2026 Oumuamua Labs <info@oumuamua.dev>
3// SPDX-License-Identifier: AGPL-3.0-only
4
5use crate::chiplet::ChipletDef;
6use crate::constraint::{ConstraintAst, ConstraintExpr, ExprId};
7use crate::expander::{RING_BLIND_BITS, RingSwitchPlan, claim_weights, eq_tensor_b};
8use crate::linearized::{self, RingGadget, linearized_coeffs};
9use crate::permutation::{BusKind, Source, eval_row_idx_byte_mle, eval_row_idx_le_mle};
10use crate::predicate::{AffineRow, ClaimLayout, Form, PredicateRows, Unknown, WireRole, compile};
11use crate::{Air, ProgramInstance};
12use alloc::collections::BTreeMap;
13use alloc::string::String;
14use alloc::vec;
15use alloc::vec::Vec;
16use core::mem;
17use hekate_core::errors;
18use hekate_core::outer::{OUTER_MASK_ROWS, OuterGeometry};
19use hekate_core::poly::univariate::{MAX_POINTS, UnivariatePoly};
20use hekate_core::tensor::TensorProduct;
21use hekate_math::{Block128, Flat, HardwareField, TowerField};
22
23pub struct LinearBatch<F> {
24    pub weights: Vec<Vec<Flat<F>>>,
25    pub rows: Vec<usize>,
26    pub target: Flat<F>,
27}
28
29pub enum BusSource<F> {
30    Claim(u32),
31    Public(Flat<F>),
32}
33
34pub struct BusRow<F> {
35    pub h_pad: u32,
36    pub h_masked: Flat<F>,
37    pub h_wire: u32,
38    pub sources: Vec<BusSource<F>>,
39    pub selector: Option<u32>,
40    pub recv_selector: Option<u32>,
41    pub eq_lookup: Flat<F>,
42}
43
44pub struct ConsistencyInputs<F> {
45    /// Pad index of claim 0, matching `ClaimLayout`.
46    pub pad_first: u32,
47
48    pub eq_zc: Flat<F>,
49    pub alpha: Flat<F>,
50    pub gamma: Flat<F>,
51    pub beta: Flat<F>,
52
53    /// `(claim index, public value, eq(r_final, row))`.
54    pub boundary: Vec<(u32, Flat<F>, Flat<F>)>,
55
56    /// `(b_k claim index, b_k_next claim index)`.
57    pub telescope: Vec<(u32, u32)>,
58
59    pub buses: Vec<BusRow<F>>,
60
61    /// Mask form of the ZeroCheck's final claim,
62    /// and the masked value the replay reaches.
63    pub val_final_form: Vec<(Unknown, Flat<F>)>,
64    pub val_final_masked: Flat<F>,
65}
66
67#[derive(Clone, Copy, Debug, Eq, PartialEq)]
68pub struct TableShape {
69    pub num_vars: usize,
70
71    /// Virtual columns, after expansion.
72    pub num_columns: usize,
73
74    pub sumcheck_degree: usize,
75    pub num_buses: usize,
76    pub mul_nodes: usize,
77    pub ring_units: bool,
78}
79
80impl TableShape {
81    pub fn from_air<F: TowerField>(
82        air: &impl Air<F>,
83        num_vars: usize,
84        ast: &ConstraintAst<F>,
85    ) -> errors::Result<Self> {
86        let num_buses = air.permutation_checks().len();
87
88        let mut sumcheck_degree = ast.max_degree() + 1;
89
90        if !air.boundary_constraints().is_empty() {
91            sumcheck_degree = sumcheck_degree.max(2);
92        }
93
94        if num_buses > 0 {
95            sumcheck_degree = sumcheck_degree.max(3);
96        }
97
98        if sumcheck_degree + 1 > MAX_POINTS {
99            return Err(errors::Error::Protocol {
100                protocol: "air",
101                message: "sumcheck degree exceeds the interpolation stack bound",
102            });
103        }
104
105        if air.num_columns() != air.virtual_column_layout().len() {
106            return Err(errors::Error::Protocol {
107                protocol: "air",
108                message: "num_columns does not match the virtual column layout",
109            });
110        }
111
112        let entries = air.virtual_expander().map(|e| e.expansion_entries());
113        let plan = RingSwitchPlan::new(air.column_layout(), entries.as_deref(), 0, 0)?;
114
115        Ok(Self {
116            num_vars,
117            num_columns: air.num_columns(),
118            sumcheck_degree,
119            num_buses,
120            mul_nodes: mul_node_count(ast),
121            ring_units: plan.has_ring(),
122        })
123    }
124
125    /// Eval claims per half: virtual columns, whole blind
126    /// units, the ring blind unit's planes, one `h` per bus.
127    pub fn eval_claims(&self, blinding_columns: usize) -> usize {
128        let ring_blind = match self.ring_units && blinding_columns > 0 {
129            true => RING_BLIND_BITS,
130            false => 0,
131        };
132
133        self.num_columns + blinding_columns + ring_blind + self.num_buses
134    }
135
136    /// Final sumcheck values are the masked running
137    /// claims themselves and consume no pad entry.
138    pub fn masked_scalars(&self, blinding_columns: usize) -> usize {
139        let zerocheck = self.num_vars * self.sumcheck_degree;
140        let claims = 2 * self.eval_claims(blinding_columns);
141        let trace_eval = 2 * self.num_vars;
142
143        2 * self.num_buses + zerocheck + claims + trace_eval
144    }
145
146    /// Each bus contributes one `h · key` product;
147    /// ring units add the gadget's squaring chain.
148    pub fn mul_wires(&self) -> usize {
149        let gadget = match self.ring_units {
150            true => linearized::BITS - 1,
151            false => 0,
152        };
153
154        self.mul_nodes + self.num_buses + gadget
155    }
156}
157
158#[derive(Clone, Copy, Debug, Eq, PartialEq)]
159pub struct OuterStatement {
160    pub masked_scalars: usize,
161    pub mul_wires: usize,
162}
163
164impl OuterStatement {
165    pub fn new(shapes: &[TableShape], blinding_columns: usize) -> Self {
166        Self {
167            masked_scalars: shapes
168                .iter()
169                .map(|s| s.masked_scalars(blinding_columns))
170                .sum(),
171            mul_wires: shapes.iter().map(TableShape::mul_wires).sum(),
172        }
173    }
174
175    /// Statement size for a whole proof.
176    pub fn for_tables<F: TowerField, A: Air<F>>(
177        main: &A,
178        main_num_vars: usize,
179        chiplets: &[ChipletDef<F>],
180        chiplet_num_vars: &[usize],
181        blinding_columns: usize,
182    ) -> errors::Result<Self> {
183        if chiplets.len() != chiplet_num_vars.len() {
184            return Err(errors::Error::Protocol {
185                protocol: "outer",
186                message: "one height per chiplet is required",
187            });
188        }
189
190        let mut shapes = vec![TableShape::from_air(
191            main,
192            main_num_vars,
193            &main.constraint_ast(),
194        )?];
195
196        for (def, &num_vars) in chiplets.iter().zip(chiplet_num_vars) {
197            shapes.push(TableShape::from_air(def, num_vars, &def.constraint_ast())?);
198        }
199
200        Ok(Self::new(&shapes, blinding_columns))
201    }
202}
203
204/// Oracle row order: pad, then the `Lhs`, `Rhs` and
205/// `Product` wire blocks, then the five mask rows.
206#[derive(Clone, Copy, Debug, Eq, PartialEq)]
207pub struct OuterLayout {
208    pub message_len: usize,
209    pub pad_rows: usize,
210    pub aux_rows: usize,
211    pub domain_len: usize,
212}
213
214impl OuterLayout {
215    pub fn new(
216        geom: &OuterGeometry,
217        masked_scalars: usize,
218        mul_wires: usize,
219    ) -> errors::Result<Self> {
220        if geom.message_len == 0 {
221            return Err(errors::Error::Protocol {
222                protocol: "outer",
223                message: "message length must be positive",
224            });
225        }
226
227        Ok(Self {
228            message_len: geom.message_len,
229            pad_rows: masked_scalars.div_ceil(geom.message_len),
230            aux_rows: mul_wires.div_ceil(geom.message_len).max(1),
231            domain_len: geom.domain_len,
232        })
233    }
234
235    pub fn slot(&self, unknown: Unknown) -> (usize, usize) {
236        match unknown {
237            Unknown::Pad(i) => {
238                let i = i as usize;
239
240                (i / self.message_len, i % self.message_len)
241            }
242            Unknown::Wire { mul, role } => {
243                let mul = mul as usize;
244                let block = role as usize;
245
246                (
247                    self.pad_rows + block * self.aux_rows + mul / self.message_len,
248                    mul % self.message_len,
249                )
250            }
251        }
252    }
253
254    pub fn interleaved_mask(&self) -> usize {
255        self.pad_rows + 3 * self.aux_rows
256    }
257
258    pub fn linear_mask(&self) -> usize {
259        self.interleaved_mask() + 1
260    }
261
262    pub fn linear_mask_hi(&self) -> usize {
263        self.interleaved_mask() + 2
264    }
265
266    pub fn quadratic_mask(&self) -> usize {
267        self.interleaved_mask() + 3
268    }
269
270    pub fn quadratic_mask_hi(&self) -> usize {
271        self.interleaved_mask() + 4
272    }
273
274    pub fn total_rows(&self) -> usize {
275        self.interleaved_mask() + OUTER_MASK_ROWS
276    }
277
278    pub fn hadamard_triples(&self) -> Vec<[usize; 3]> {
279        (0..self.aux_rows)
280            .map(|i| {
281                [
282                    self.pad_rows + i,
283                    self.pad_rows + self.aux_rows + i,
284                    self.pad_rows + 2 * self.aux_rows + i,
285                ]
286            })
287            .collect()
288    }
289}
290
291pub struct EvalRecord<F> {
292    pub mask_form: Vec<(Unknown, Flat<F>)>,
293    pub claim_masked: Flat<F>,
294    pub fin: Flat<F>,
295}
296
297/// One table's share of the statement, in transcript order.
298pub struct TableRecord<F> {
299    pub claims_masked: Vec<Flat<F>>,
300    pub predicate: PredicateRows<F>,
301    pub consistency: ConsistencyInputs<F>,
302
303    /// `(pad, claim, public value)` per fixed column.
304    pub fixed: Vec<(u32, u32, Flat<F>)>,
305
306    pub trace_eval: EvalRecord<F>,
307    pub gadget: Option<RingGadget<F>>,
308
309    /// `(bus id, pad, masked claimed sum)` per bus.
310    pub claimed_sums: Vec<(String, u32, Flat<F>)>,
311
312    pub mul_wires: u32,
313}
314
315/// The assembled statement: affine rows, and the ring
316/// gadgets whose tie coefficients [`linear_weights`]
317/// streams onto their chain rows.
318pub struct OuterRows<F> {
319    pub affine: Vec<AffineRow<F>>,
320    pub gadgets: Vec<RingGadget<F>>,
321}
322
323/// `H_0` of an eval sumcheck over the pad.
324pub struct InitialForm<F> {
325    pub form: Vec<(Unknown, Flat<F>)>,
326    pub gadget: Option<RingGadget<F>>,
327}
328
329/// First pad index of each masked group of a table,
330/// in the prover's masking order.
331#[derive(Clone, Copy, Debug)]
332pub struct TablePads {
333    pub claimed_sums: u32,
334    pub zerocheck: u32,
335    pub h_evals: u32,
336    pub claims: u32,
337}
338
339/// Public data of one eval sumcheck on masked claims.
340pub struct EvalInputs<'a, F> {
341    pub plan: &'a RingSwitchPlan,
342    pub eta: F,
343    pub r_mix: &'a [Block128],
344    pub shifted_claims: bool,
345    pub challenges: &'a [Flat<F>],
346    pub claim_masked: Flat<F>,
347    pub fin: Flat<F>,
348    pub claim_pads: &'a [u32],
349    pub round_first: u32,
350    pub first_wire: u32,
351}
352
353pub struct EvalOutput<F> {
354    pub record: EvalRecord<F>,
355    pub gadget: Option<RingGadget<F>>,
356}
357
358/// Everything a table's record is built from; both
359/// sides hold all of it after the table's eval phase.
360pub struct TableInputs<'a, F: TowerField, A> {
361    pub air: &'a A,
362    pub ast: &'a ConstraintAst<F>,
363    pub instance: &'a ProgramInstance<F>,
364    pub num_vars: usize,
365    pub trace_width: usize,
366    pub blinding_columns: usize,
367
368    pub alpha: Flat<F>,
369    pub gamma: Flat<F>,
370    pub beta: Flat<F>,
371    pub r_zerocheck: &'a [Flat<F>],
372    pub r_final: &'a [Flat<F>],
373    pub lookup_bus_points: &'a BTreeMap<String, Vec<Flat<F>>>,
374
375    pub claimed_sums_masked: &'a [(String, F)],
376    pub h_evals_masked: &'a [(String, F)],
377    pub val_final_masked: Flat<F>,
378    pub claims_masked: Vec<Flat<F>>,
379
380    pub pads: TablePads,
381    pub trace_eval: EvalRecord<F>,
382    pub gadget: Option<RingGadget<F>>,
383}
384
385pub fn eval_initial_form<F: HardwareField + Into<Block128> + From<u128>>(
386    plan: &RingSwitchPlan,
387    eta: F,
388    r_mix: &[Block128],
389    shifted_claims: bool,
390    claim_pads: &[u32],
391    first_wire: u32,
392) -> errors::Result<InitialForm<F>> {
393    let weights = claim_weights(plan, eta, shifted_claims);
394    if weights.len() != claim_pads.len() {
395        return Err(errors::Error::Protocol {
396            protocol: "outer",
397            message: "one pad entry per eval claim is required",
398        });
399    }
400
401    let mut form = Vec::new();
402    let mut ring = Vec::new();
403
404    for ((is_ring, weight), &pad) in weights.iter().zip(claim_pads) {
405        let a: Flat<F> = F::from(weight.0).to_hardware();
406        match is_ring {
407            true => ring.push((pad, a)),
408            false => form.push((Unknown::Pad(pad), a)),
409        }
410    }
411
412    if ring.is_empty() {
413        return Ok(InitialForm { form, gadget: None });
414    }
415
416    let mu: Vec<Flat<F>> = linearized_coeffs(&eq_tensor_b(r_mix))
417        .iter()
418        .map(|m| F::from(m.to_tower().0).to_hardware())
419        .collect();
420
421    let gadget = RingGadget::new(ring, mu, first_wire);
422    form.extend(gadget.delta_form());
423
424    Ok(InitialForm {
425        form,
426        gadget: Some(gadget),
427    })
428}
429
430/// The eval row's mask form propagated
431/// through the rounds, plus the ring gadget.
432pub fn eval_record<F>(inputs: &EvalInputs<'_, F>) -> errors::Result<EvalOutput<F>>
433where
434    F: HardwareField + Into<Block128> + From<u128>,
435{
436    let initial = eval_initial_form(
437        inputs.plan,
438        inputs.eta,
439        inputs.r_mix,
440        inputs.shifted_claims,
441        inputs.claim_pads,
442        inputs.first_wire,
443    )?;
444
445    let round_pads: Vec<[u32; 2]> = (0..inputs.challenges.len() as u32)
446        .map(|i| [inputs.round_first + 2 * i, inputs.round_first + 2 * i + 1])
447        .collect();
448    let round_refs: Vec<&[u32]> = round_pads.iter().map(|p| p.as_slice()).collect();
449
450    Ok(EvalOutput {
451        record: EvalRecord {
452            mask_form: sumcheck_mask_form(&initial.form, &round_refs, inputs.challenges, 2),
453            claim_masked: inputs.claim_masked,
454            fin: inputs.fin,
455        },
456        gadget: initial.gadget,
457    })
458}
459
460pub fn table_record<F, A>(inputs: TableInputs<'_, F, A>) -> errors::Result<TableRecord<F>>
461where
462    F: TowerField + HardwareField,
463    A: Air<F>,
464{
465    let air = inputs.air;
466    let ast = inputs.ast;
467
468    let boundary_constraints = air.boundary_constraints();
469    let bus_specs = air.permutation_checks();
470    let shape = TableShape::from_air(air, inputs.num_vars, ast)?;
471
472    let num_vars = inputs.num_vars;
473    let trace_width = inputs.trace_width;
474    let r_final = inputs.r_final;
475    let pad_first = inputs.pads.claims;
476    let half = shape.eval_claims(inputs.blinding_columns) as u32;
477
478    if inputs.claims_masked.len() != 2 * half as usize {
479        return Err(errors::Error::Protocol {
480            protocol: "outer",
481            message: "eval claim count does not match the table shape",
482        });
483    }
484
485    let eq_row_checker = TensorProduct::new(r_final.to_vec());
486
487    let mut fixed = Vec::new();
488    for fc in &air.fixed_columns() {
489        if fc.col_idx >= trace_width {
490            return Err(errors::Error::Protocol {
491                protocol: "outer",
492                message: "fixed column col_idx out of trace_width",
493            });
494        }
495
496        fixed.push((
497            pad_first + fc.col_idx as u32,
498            fc.col_idx as u32,
499            fc.shape.evaluate(r_final),
500        ));
501    }
502
503    let mut boundary = Vec::with_capacity(boundary_constraints.len());
504    for bc in &boundary_constraints {
505        if bc.col_idx >= trace_width || bc.row_idx >= 1 << num_vars {
506            return Err(errors::Error::Protocol {
507                protocol: "outer",
508                message: "boundary constraint out of the table",
509            });
510        }
511
512        boundary.push((
513            bc.col_idx as u32,
514            bc.resolve_target(inputs.instance)?.to_hardware(),
515            eq_row_checker.evaluate_at_index(bc.row_idx),
516        ));
517    }
518
519    let mut buses = Vec::with_capacity(bus_specs.len());
520    for (spec_idx, (bus_id, spec)) in bus_specs.iter().enumerate() {
521        let eq_lookup = match spec.kind {
522            BusKind::Permutation => Flat::from_raw(F::ONE),
523            BusKind::Lookup => {
524                let r_bus =
525                    inputs
526                        .lookup_bus_points
527                        .get(bus_id)
528                        .ok_or(errors::Error::Protocol {
529                            protocol: "outer",
530                            message: "lookup bus spec missing r_bus point",
531                        })?;
532
533                if r_bus.len() < num_vars {
534                    return Err(errors::Error::Protocol {
535                        protocol: "outer",
536                        message: "r_bus shorter than table num_vars",
537                    });
538                }
539
540                let one = Flat::from_raw(F::ONE);
541                let mut eq_r_hi_at_0 = one;
542                for r_j in &r_bus[num_vars..] {
543                    eq_r_hi_at_0 *= one - *r_j;
544                }
545
546                eq_r_hi_at_0 * TensorProduct::evaluate_eq_slice(&r_bus[0..num_vars], r_final)
547            }
548        };
549
550        let mut sources = Vec::new();
551        for (source, _) in &spec.sources {
552            match source {
553                Source::Column(col) | Source::PhaseColumn(col) => {
554                    sources.push(BusSource::Claim(*col as u32))
555                }
556                Source::Columns(indices) => {
557                    sources.extend(indices.iter().map(|&col| BusSource::Claim(col as u32)));
558                }
559                Source::Const(v) => sources.push(BusSource::Public(F::from(*v).to_hardware())),
560                Source::RowIndexLeBytes(n) => {
561                    sources.push(BusSource::Public(eval_row_idx_le_mle::<F>(*n, r_final)));
562                }
563                Source::RowIndexByte(n) => {
564                    sources.push(BusSource::Public(eval_row_idx_byte_mle::<F>(*n, r_final)));
565                }
566            }
567        }
568
569        buses.push(BusRow {
570            h_pad: inputs.pads.h_evals + spec_idx as u32,
571            h_masked: inputs.h_evals_masked[spec_idx].1.to_hardware(),
572            h_wire: (shape.mul_nodes + spec_idx) as u32,
573            sources,
574            selector: spec.selector.map(|i| i as u32),
575            recv_selector: spec.recv_selector.map(|i| i as u32),
576            eq_lookup,
577        });
578    }
579
580    let logup_alpha_offset = ast.roots.len() + boundary_constraints.len() + inputs.blinding_columns;
581    let mut alpha_pow = Flat::from_raw(F::ONE);
582
583    for _ in 0..logup_alpha_offset {
584        alpha_pow *= inputs.alpha;
585    }
586
587    let mut initial = Vec::with_capacity(bus_specs.len());
588    for k in 0..bus_specs.len() as u32 {
589        initial.push((Unknown::Pad(inputs.pads.claimed_sums + k), alpha_pow));
590        alpha_pow *= inputs.alpha;
591    }
592
593    let degree = shape.sumcheck_degree;
594    let round_pads: Vec<Vec<u32>> = (0..num_vars)
595        .map(|i| {
596            let first = inputs.pads.zerocheck + (i * degree) as u32;
597
598            (0..degree as u32).map(|j| first + j).collect()
599        })
600        .collect();
601    let round_refs: Vec<&[u32]> = round_pads.iter().map(|p| p.as_slice()).collect();
602
603    let consistency = ConsistencyInputs {
604        pad_first,
605        eq_zc: TensorProduct::evaluate_eq_slice(inputs.r_zerocheck, r_final),
606        alpha: inputs.alpha,
607        gamma: inputs.gamma,
608        beta: inputs.beta,
609        boundary,
610        telescope: (0..inputs.blinding_columns as u32)
611            .map(|k| {
612                let b = trace_width as u32 + k;
613
614                (b, half + b)
615            })
616            .collect(),
617        buses,
618        val_final_form: sumcheck_mask_form(&initial, &round_refs, r_final, degree),
619        val_final_masked: inputs.val_final_masked,
620    };
621
622    Ok(TableRecord {
623        claims_masked: inputs.claims_masked,
624        predicate: compile(ast, ClaimLayout { pad_first, half }),
625        consistency,
626        fixed,
627        trace_eval: inputs.trace_eval,
628        gadget: inputs.gadget,
629        claimed_sums: inputs
630            .claimed_sums_masked
631            .iter()
632            .enumerate()
633            .map(|(k, (bus_id, sum))| {
634                (
635                    bus_id.clone(),
636                    inputs.pads.claimed_sums + k as u32,
637                    sum.to_hardware(),
638                )
639            })
640            .collect(),
641        mul_wires: shape.mul_wires() as u32,
642    })
643}
644
645/// Every affine row of the statement, claims folded
646/// into constants and wire indices offset per table.
647pub fn assemble<F: TowerField + HardwareField>(
648    records: Vec<TableRecord<F>>,
649    statement: &OuterStatement,
650) -> errors::Result<OuterRows<F>> {
651    let mut affine = Vec::new();
652    let mut gadgets = Vec::new();
653    let mut offset = 0u32;
654    let mut buses: BTreeMap<String, Vec<(u32, Flat<F>)>> = BTreeMap::new();
655
656    for mut record in records {
657        let first = affine.len();
658
659        affine.append(&mut record.predicate.affine);
660        affine.push(consistency_row(
661            &record.predicate.roots,
662            &record.consistency,
663        ));
664
665        for bus in &record.consistency.buses {
666            let [lhs, rhs] =
667                bus_operand_rows(bus, record.consistency.pad_first, record.consistency.beta);
668
669            affine.push(lhs);
670            affine.push(rhs);
671        }
672
673        for &(pad, claim, public) in &record.fixed {
674            affine.push(fixed_column_row(pad, claim, public));
675        }
676
677        affine.push(eval_final_row(
678            &record.trace_eval.mask_form,
679            record.trace_eval.claim_masked,
680            record.trace_eval.fin,
681        ));
682
683        let half = (record.claims_masked.len() / 2) as u32;
684        let num_buses = record.consistency.buses.len() as u32;
685        let pad_first = record.consistency.pad_first;
686
687        for (k, bus) in record.consistency.buses.iter().enumerate() {
688            let claim = half - num_buses + k as u32;
689
690            affine.push(h_claim_pin_row(
691                bus.h_pad,
692                bus.h_masked,
693                pad_first + claim,
694                claim,
695            ));
696        }
697
698        if let Some(mut gadget) = record.gadget {
699            gadget.first_row = affine.len();
700            affine.extend(gadget.wire_rows());
701            gadgets.push(gadget);
702        }
703
704        for row in &mut affine[first..] {
705            resolve_table_row(row, &record.claims_masked, offset);
706        }
707
708        offset += record.mul_wires;
709
710        for (bus_id, pad, masked) in record.claimed_sums {
711            buses.entry(bus_id).or_default().push((pad, masked));
712        }
713    }
714
715    if offset as usize != statement.mul_wires {
716        return Err(errors::Error::Protocol {
717            protocol: "outer",
718            message: "table wire counts do not sum to the statement",
719        });
720    }
721
722    for endpoints in buses.values() {
723        affine.push(bus_sum_row(endpoints));
724    }
725
726    Ok(OuterRows { affine, gadgets })
727}
728
729pub fn reachable<F: TowerField>(ast: &ConstraintAst<F>) -> Vec<bool> {
730    let n = ast.arena.len();
731    let mut live = vec![false; n];
732
733    for root in &ast.roots {
734        live[root.0 as usize] = true;
735    }
736
737    // Children sit at lower indices than their parent
738    for i in (0..n).rev() {
739        if !live[i] {
740            continue;
741        }
742
743        let mut mark = |id: &ExprId| live[id.0 as usize] = true;
744        match ast.arena.get(ExprId(i as u32)) {
745            ConstraintExpr::Cell(_) | ConstraintExpr::Const(_) => {}
746            ConstraintExpr::Add(a, b) | ConstraintExpr::Mul(a, b) => {
747                mark(a);
748                mark(b);
749            }
750            ConstraintExpr::Scale(_, a) => mark(a),
751            ConstraintExpr::Sum(children) => children.iter().for_each(mark),
752        }
753    }
754
755    live
756}
757
758pub fn mul_node_count<F: TowerField>(ast: &ConstraintAst<F>) -> usize {
759    reachable(ast)
760        .iter()
761        .enumerate()
762        .filter(|&(i, &live)| {
763            live && matches!(ast.arena.get(ExprId(i as u32)), ConstraintExpr::Mul(..))
764        })
765        .count()
766}
767
768pub fn consistency_row<F: TowerField + HardwareField>(
769    roots: &[Form<F>],
770    inputs: &ConsistencyInputs<F>,
771) -> AffineRow<F> {
772    let mut row = AffineRow {
773        unknowns: inputs.val_final_form.to_vec(),
774        claims: Vec::new(),
775        constant: inputs.val_final_masked,
776    };
777
778    let mut alpha_pow = Flat::from_raw(F::ONE);
779
780    for form in roots {
781        scale_into(&mut row, form, alpha_pow * inputs.eq_zc);
782        alpha_pow *= inputs.alpha;
783    }
784
785    for &(claim, public, eq_row) in &inputs.boundary {
786        let weight = eq_row * alpha_pow;
787
788        push_claim(&mut row, inputs.pad_first, claim, weight);
789        row.constant += public * weight;
790
791        alpha_pow *= inputs.alpha;
792    }
793
794    for &(b, b_next) in &inputs.telescope {
795        push_claim(&mut row, inputs.pad_first, b, alpha_pow);
796        push_claim(&mut row, inputs.pad_first, b_next, alpha_pow);
797
798        alpha_pow *= inputs.alpha;
799    }
800
801    for bus in &inputs.buses {
802        let consistency = alpha_pow * inputs.eq_zc;
803        let gamma_weight = consistency * inputs.gamma;
804
805        row.unknowns.push((Unknown::Pad(bus.h_pad), gamma_weight));
806
807        row.constant += gamma_weight * bus.h_masked;
808
809        row.unknowns.push((
810            Unknown::Wire {
811                mul: bus.h_wire,
812                role: WireRole::Product,
813            },
814            consistency,
815        ));
816
817        match bus.selector {
818            Some(s) => push_claim(&mut row, inputs.pad_first, s, consistency),
819            None => row.constant += consistency,
820        }
821
822        if let Some(s) = bus.recv_selector {
823            push_claim(&mut row, inputs.pad_first, s, consistency);
824        }
825
826        let lookup_weight = alpha_pow * bus.eq_lookup;
827
828        row.unknowns.push((Unknown::Pad(bus.h_pad), lookup_weight));
829
830        row.constant += lookup_weight * bus.h_masked;
831
832        alpha_pow *= inputs.alpha;
833    }
834
835    row
836}
837
838/// Pad form of `claim'_n - claim_n` for a sumcheck whose
839/// replay reconstructs `p'_0 = claim' - p'_1` each round:
840/// entries of round `i` carry `Π_{j>i} L_0(r_j)`.
841pub fn sumcheck_mask_form<F: TowerField + HardwareField>(
842    initial: &[(Unknown, Flat<F>)],
843    round_pads: &[&[u32]],
844    challenges: &[Flat<F>],
845    degree: usize,
846) -> Vec<(Unknown, Flat<F>)> {
847    let rounds: Vec<Vec<Flat<F>>> = challenges
848        .iter()
849        .map(|&r| lagrange_weights::<F>(degree, r))
850        .collect();
851
852    let mut suffix = vec![Flat::from_raw(F::ONE); rounds.len() + 1];
853    for i in (0..rounds.len()).rev() {
854        suffix[i] = suffix[i + 1] * rounds[i][0];
855    }
856
857    let mut form: Vec<(Unknown, Flat<F>)> = initial
858        .iter()
859        .map(|&(u, coeff)| (u, coeff * suffix[0]))
860        .collect();
861
862    for (i, (pads, weights)) in round_pads.iter().zip(&rounds).enumerate() {
863        let after = suffix[i + 1];
864
865        form.push((Unknown::Pad(pads[0]), weights[0] * after));
866
867        for (j, &pad) in pads.iter().enumerate() {
868            form.push((Unknown::Pad(pad), weights[j + 1] * after));
869        }
870    }
871
872    form
873}
874
875/// `claim_n = fin` for an eval sumcheck: `fin` is public,
876/// `claim_n` is the masked running claim minus its form.
877pub fn eval_final_row<F: TowerField + HardwareField>(
878    mask_form: &[(Unknown, Flat<F>)],
879    claim_masked: Flat<F>,
880    fin: Flat<F>,
881) -> AffineRow<F> {
882    AffineRow {
883        unknowns: mask_form.to_vec(),
884        claims: Vec::new(),
885        constant: claim_masked + fin,
886    }
887}
888
889/// `c[claim] = public`, i.e. `h[pad] = c'[claim] + public`.
890pub fn fixed_column_row<F: TowerField + HardwareField>(
891    pad: u32,
892    claim: u32,
893    public: Flat<F>,
894) -> AffineRow<F> {
895    let one = Flat::from_raw(F::ONE);
896
897    AffineRow {
898        unknowns: vec![(Unknown::Pad(pad), one)],
899        claims: vec![(claim, one)],
900        constant: public,
901    }
902}
903
904/// `c[claim] = h_eval`: the eval argument's base `h`
905/// claim and the reported `h_eval` unmask to one value.
906pub fn h_claim_pin_row<F: TowerField + HardwareField>(
907    h_pad: u32,
908    h_masked: Flat<F>,
909    claim_pad: u32,
910    claim: u32,
911) -> AffineRow<F> {
912    let one = Flat::from_raw(F::ONE);
913
914    AffineRow {
915        unknowns: vec![(Unknown::Pad(h_pad), one), (Unknown::Pad(claim_pad), one)],
916        claims: vec![(claim, one)],
917        constant: h_masked,
918    }
919}
920
921/// `Σ claimed_sum = 0` over one bus id:
922/// `(pad, masked sum)` per endpoint.
923pub fn bus_sum_row<F: TowerField + HardwareField>(endpoints: &[(u32, Flat<F>)]) -> AffineRow<F> {
924    let mut constant = Flat::from_raw(F::ZERO);
925    for &(_, masked) in endpoints {
926        constant += masked;
927    }
928
929    AffineRow {
930        unknowns: endpoints
931            .iter()
932            .map(|&(pad, _)| (Unknown::Pad(pad), Flat::from_raw(F::ONE)))
933            .collect(),
934        claims: Vec::new(),
935        constant,
936    }
937}
938
939/// Folds a table's public masked claims into the constant
940/// and shifts its table-local wire indices by `mul_offset`.
941pub fn resolve_table_row<F: TowerField + HardwareField>(
942    row: &mut AffineRow<F>,
943    claims: &[Flat<F>],
944    mul_offset: u32,
945) {
946    for &(idx, coeff) in &row.claims {
947        row.constant += coeff * claims[idx as usize];
948    }
949
950    row.claims.clear();
951
952    for (unknown, _) in row.unknowns.iter_mut() {
953        if let Unknown::Wire { mul, .. } = unknown {
954            *mul += mul_offset;
955        }
956    }
957}
958
959pub fn bus_key_form<F: TowerField + HardwareField>(
960    bus: &BusRow<F>,
961    pad_first: u32,
962    beta: Flat<F>,
963) -> Form<F> {
964    let mut form = Form::default();
965    let mut beta_pow = Flat::from_raw(F::ONE);
966
967    for source in &bus.sources {
968        match source {
969            BusSource::Claim(idx) => {
970                form.unknowns
971                    .push((Unknown::Pad(pad_first + idx), beta_pow));
972                form.claims.push((*idx, beta_pow));
973            }
974            BusSource::Public(v) => form.constant += *v * beta_pow,
975        }
976
977        beta_pow *= beta;
978    }
979
980    form
981}
982
983/// Pins the bus Hadamard wire `h · key`:
984/// `Lhs = h` and `Rhs = key` over the pad.
985pub fn bus_operand_rows<F: TowerField + HardwareField>(
986    bus: &BusRow<F>,
987    pad_first: u32,
988    beta: Flat<F>,
989) -> [AffineRow<F>; 2] {
990    let one = Flat::from_raw(F::ONE);
991
992    let lhs = AffineRow {
993        unknowns: vec![
994            (
995                Unknown::Wire {
996                    mul: bus.h_wire,
997                    role: WireRole::Lhs,
998                },
999                one,
1000            ),
1001            (Unknown::Pad(bus.h_pad), one),
1002        ],
1003        claims: Vec::new(),
1004        constant: bus.h_masked,
1005    };
1006
1007    let key = bus_key_form(bus, pad_first, beta);
1008
1009    let mut unknowns = key.unknowns;
1010    unknowns.push((
1011        Unknown::Wire {
1012            mul: bus.h_wire,
1013            role: WireRole::Rhs,
1014        },
1015        one,
1016    ));
1017
1018    let rhs = AffineRow {
1019        unknowns,
1020        claims: key.claims,
1021        constant: key.constant,
1022    };
1023
1024    [lhs, rhs]
1025}
1026
1027/// PAD oracle rows: the pad in [`OuterLayout::slot`] order,
1028/// `filler` (exactly `pad_rows · (code_len - message_len)`
1029/// entries) on the message tail.
1030pub fn build_pad_rows<F: TowerField + HardwareField>(
1031    layout: &OuterLayout,
1032    pad: &[Flat<F>],
1033    code_len: usize,
1034    filler: &[Flat<F>],
1035) -> errors::Result<Vec<Vec<Flat<F>>>> {
1036    let tail = code_len - layout.message_len;
1037
1038    if pad.len() > layout.pad_rows * layout.message_len || filler.len() != layout.pad_rows * tail {
1039        return Err(errors::Error::Protocol {
1040            protocol: "outer",
1041            message: "pad rows do not hold the pad and its filler",
1042        });
1043    }
1044
1045    let mut rows = vec![vec![Flat::from_raw(F::ZERO); layout.domain_len]; layout.pad_rows];
1046    for (i, value) in pad.iter().enumerate() {
1047        let (r, c) = layout.slot(Unknown::Pad(i as u32));
1048        rows[r][c] = *value;
1049    }
1050
1051    for (row, fill) in rows.iter_mut().zip(filler.chunks_exact(tail)) {
1052        row[layout.message_len..code_len].copy_from_slice(fill);
1053    }
1054
1055    Ok(rows)
1056}
1057
1058/// AUX oracle rows: the wire blocks and the five mask
1059/// rows, `filler` on every message tail, then the
1060/// zero-sum mask row's first `code_len - 1` cells and
1061/// the messages of the interleaved and both high masks.
1062pub fn build_aux_rows<F: TowerField + HardwareField>(
1063    layout: &OuterLayout,
1064    wires: &[[Flat<F>; 3]],
1065    code_len: usize,
1066    filler: &[Flat<F>],
1067) -> errors::Result<Vec<Vec<Flat<F>>>> {
1068    let count = layout.total_rows() - layout.pad_rows;
1069    let tail = code_len - layout.message_len;
1070
1071    if wires.len() > layout.aux_rows * layout.message_len
1072        || filler.len() != aux_filler_len(layout, code_len)
1073    {
1074        return Err(errors::Error::Protocol {
1075            protocol: "outer",
1076            message: "aux rows do not hold the wires and their filler",
1077        });
1078    }
1079
1080    let mut rows = vec![vec![Flat::from_raw(F::ZERO); layout.domain_len]; count];
1081
1082    for (mul, wire) in wires.iter().enumerate() {
1083        for role in [WireRole::Lhs, WireRole::Rhs, WireRole::Product] {
1084            let (r, c) = layout.slot(Unknown::Wire {
1085                mul: mul as u32,
1086                role,
1087            });
1088
1089            rows[r - layout.pad_rows][c] = wire[role as usize];
1090        }
1091    }
1092
1093    let (tails, rest) = filler.split_at(count * tail);
1094    for (row, fill) in rows.iter_mut().zip(tails.chunks_exact(tail)) {
1095        row[layout.message_len..code_len].copy_from_slice(fill);
1096    }
1097
1098    let (zero_sum_fill, rest) = rest.split_at(code_len - 1);
1099
1100    let zero_sum = layout.linear_mask() - layout.pad_rows;
1101    let mut acc = Flat::from_raw(F::ZERO);
1102
1103    for (slot, &v) in rows[zero_sum][..code_len - 1].iter_mut().zip(zero_sum_fill) {
1104        *slot = v;
1105        acc += v;
1106    }
1107
1108    rows[zero_sum][code_len - 1] = acc;
1109
1110    // The quadratic mask keeps its zero message; every
1111    // other mask row is uniform over the whole message.
1112    let uniform = [
1113        layout.interleaved_mask(),
1114        layout.linear_mask_hi(),
1115        layout.quadratic_mask_hi(),
1116    ];
1117
1118    for (row, fill) in uniform
1119        .into_iter()
1120        .zip(rest.chunks_exact(layout.message_len))
1121    {
1122        rows[row - layout.pad_rows][..layout.message_len].copy_from_slice(fill);
1123    }
1124
1125    Ok(rows)
1126}
1127
1128pub fn aux_filler_len(layout: &OuterLayout, code_len: usize) -> usize {
1129    let count = layout.total_rows() - layout.pad_rows;
1130
1131    count * (code_len - layout.message_len) + code_len - 1 + 3 * layout.message_len
1132}
1133
1134/// Challenge count for `outer_r_lin`. The tensor
1135/// batches the affine rows at `k / |F|`, not the
1136/// `1 / |F|` of one challenge per row.
1137pub fn linear_tensor_vars<F>(rows: &OuterRows<F>) -> usize {
1138    rows.affine.len().next_power_of_two().ilog2() as usize
1139}
1140
1141pub fn linear_weights<F: TowerField + HardwareField>(
1142    layout: &OuterLayout,
1143    rows: &OuterRows<F>,
1144    tensor: &[F],
1145) -> errors::Result<LinearBatch<F>> {
1146    if tensor.len() != linear_tensor_vars(rows) {
1147        return Err(errors::Error::Protocol {
1148            protocol: "outer",
1149            message: "linear batch takes ceil(log2(affine rows)) challenges",
1150        });
1151    }
1152
1153    if rows.affine.iter().any(|row| !row.claims.is_empty()) {
1154        return Err(errors::Error::Protocol {
1155            protocol: "outer",
1156            message: "affine rows must be resolved before weighting",
1157        });
1158    }
1159
1160    let scales = expand_batch_tensor(tensor, rows.affine.len());
1161
1162    let mut weights = vec![vec![Flat::from_raw(F::ZERO); layout.message_len]; layout.total_rows()];
1163    let mut target = Flat::from_raw(F::ZERO);
1164
1165    for (row, &scale) in rows.affine.iter().zip(&scales) {
1166        for &(unknown, coeff) in &row.unknowns {
1167            let (orow, ocol) = layout.slot(unknown);
1168            weights[orow][ocol] += scale * coeff;
1169        }
1170
1171        target += scale * row.constant;
1172    }
1173
1174    for gadget in &rows.gadgets {
1175        gadget.for_each_tie_row(|j, tie| {
1176            let scale = scales[gadget.row_of(j)];
1177            for (&(pad, _), &t) in gadget.ring.iter().zip(tie) {
1178                let (orow, ocol) = layout.slot(Unknown::Pad(pad));
1179                weights[orow][ocol] += scale * t;
1180            }
1181        });
1182    }
1183
1184    let used: Vec<usize> = (0..layout.total_rows())
1185        .filter(|&r| weights[r].iter().any(|v| *v != Flat::from_raw(F::ZERO)))
1186        .collect();
1187
1188    let selected = used.iter().map(|&r| mem::take(&mut weights[r])).collect();
1189
1190    Ok(LinearBatch {
1191        weights: selected,
1192        rows: used,
1193        target,
1194    })
1195}
1196
1197fn push_claim<F: TowerField>(row: &mut AffineRow<F>, pad_first: u32, idx: u32, coeff: Flat<F>) {
1198    row.unknowns.push((Unknown::Pad(pad_first + idx), coeff));
1199    row.claims.push((idx, coeff));
1200}
1201
1202fn scale_into<F: TowerField + HardwareField>(dst: &mut AffineRow<F>, src: &Form<F>, c: Flat<F>) {
1203    for &(u, k) in &src.unknowns {
1204        dst.unknowns.push((u, k * c));
1205    }
1206
1207    for &(i, k) in &src.claims {
1208        dst.claims.push((i, k * c));
1209    }
1210
1211    dst.constant += src.constant * c;
1212}
1213
1214fn expand_batch_tensor<F: TowerField + HardwareField>(tensor: &[F], len: usize) -> Vec<Flat<F>> {
1215    let mut scales = vec![Flat::from_raw(F::ZERO); len];
1216
1217    if len == 0 {
1218        return scales;
1219    }
1220
1221    scales[0] = Flat::from_raw(F::ONE);
1222
1223    for (j, &r) in tensor.iter().enumerate() {
1224        let half = 1usize << j;
1225        let scale = r.to_hardware();
1226
1227        for i in half..len.min(half << 1) {
1228            scales[i] = scales[i - half] * scale;
1229        }
1230    }
1231
1232    scales
1233}
1234
1235fn lagrange_weights<F: TowerField + HardwareField>(degree: usize, r: Flat<F>) -> Vec<Flat<F>> {
1236    (0..=degree)
1237        .map(|j| {
1238            let mut evals = vec![F::ZERO; degree + 1];
1239            evals[j] = F::ONE;
1240
1241            UnivariatePoly::new(evals).evaluate_hw(r)
1242        })
1243        .collect()
1244}
1245
1246#[cfg(test)]
1247mod tests {
1248    use super::*;
1249    use crate::ProgramCell;
1250    use crate::constraint::ConstraintArena;
1251    use hekate_math::Block128;
1252
1253    type F = Block128;
1254
1255    const BUS_PAD_FIRST: u32 = 10;
1256    const BUS_H_PAD: u32 = 20;
1257    const BUS_PUBLIC: u128 = 77;
1258
1259    /// `half - num_buses` at `bus_record`'s `half = 2` and one bus.
1260    const H_CLAIM: usize = 1;
1261
1262    fn shape(num_vars: usize, num_columns: usize, degree: usize, buses: usize) -> TableShape {
1263        TableShape {
1264            num_vars,
1265            num_columns,
1266            sumcheck_degree: degree,
1267            num_buses: buses,
1268            mul_nodes: 0,
1269            ring_units: false,
1270        }
1271    }
1272
1273    fn mix(seed: u128) -> Flat<F> {
1274        F::from(
1275            seed.wrapping_mul(0x9e37_79b9_7f4a_7c15)
1276                .wrapping_add(0x51ed_2701),
1277        )
1278        .to_hardware()
1279    }
1280
1281    /// One table:
1282    /// `c0 · c1` as the only constraint, two columns,
1283    /// one Permutation bus over `[Claim(0), Public(77)]`
1284    /// with selector `c1`.
1285    fn bus_record(claims_masked: Vec<Flat<F>>, h_masked: Flat<F>, beta: Flat<F>) -> TableRecord<F> {
1286        let mut arena = ConstraintArena::<F>::new();
1287
1288        let a = arena.cell(ProgramCell::current(0));
1289        let b = arena.cell(ProgramCell::current(1));
1290        let root = arena.mul(a, b);
1291
1292        let ast = ConstraintAst {
1293            arena,
1294            roots: vec![root],
1295            labels: vec![None],
1296        };
1297
1298        let predicate = compile(
1299            &ast,
1300            ClaimLayout {
1301                pad_first: BUS_PAD_FIRST,
1302                half: 2,
1303            },
1304        );
1305
1306        let bus = BusRow {
1307            h_pad: BUS_H_PAD,
1308            h_masked,
1309            h_wire: predicate.mul_nodes,
1310            sources: vec![BusSource::Claim(0), BusSource::Public(mix(BUS_PUBLIC))],
1311            selector: Some(1),
1312            recv_selector: None,
1313            eq_lookup: Flat::from_raw(F::ONE),
1314        };
1315
1316        let mul_wires = predicate.mul_nodes + 1;
1317
1318        TableRecord {
1319            claims_masked,
1320            predicate,
1321            consistency: ConsistencyInputs {
1322                pad_first: BUS_PAD_FIRST,
1323                eq_zc: mix(1),
1324                alpha: mix(2),
1325                gamma: mix(3),
1326                beta,
1327                boundary: Vec::new(),
1328                telescope: Vec::new(),
1329                buses: vec![bus],
1330                val_final_form: Vec::new(),
1331                val_final_masked: mix(4),
1332            },
1333            fixed: Vec::new(),
1334            trace_eval: EvalRecord {
1335                mask_form: Vec::new(),
1336                claim_masked: mix(5),
1337                fin: mix(6),
1338            },
1339            gadget: None,
1340            claimed_sums: vec![(String::from("bus"), 0, mix(7))],
1341            mul_wires,
1342        }
1343    }
1344
1345    fn mentions(rows: &[AffineRow<F>], mul: u32, role: WireRole) -> usize {
1346        rows.iter()
1347            .filter(|r| {
1348                r.unknowns
1349                    .iter()
1350                    .any(|&(u, _)| u == Unknown::Wire { mul, role })
1351            })
1352            .count()
1353    }
1354
1355    #[test]
1356    fn ring_units_add_squaring_chain() {
1357        let plain = shape(12, 100, 3, 0);
1358        let ring = TableShape {
1359            ring_units: true,
1360            ..plain
1361        };
1362
1363        assert_eq!(ring.mul_wires() - plain.mul_wires(), 127);
1364    }
1365
1366    #[test]
1367    fn ring_blind_planes_are_masked_only_under_blinding() {
1368        let plain = shape(12, 100, 3, 0);
1369        let ring = TableShape {
1370            ring_units: true,
1371            ..plain
1372        };
1373
1374        assert_eq!(ring.masked_scalars(0), plain.masked_scalars(0));
1375        assert_eq!(
1376            ring.masked_scalars(1) - plain.masked_scalars(1),
1377            2 * RING_BLIND_BITS
1378        );
1379    }
1380
1381    #[test]
1382    fn empty_arena_has_no_wires() {
1383        let ast = ConstraintAst::<F> {
1384            arena: ConstraintArena::new(),
1385            roots: vec![],
1386            labels: vec![],
1387        };
1388
1389        assert_eq!(mul_node_count(&ast), 0);
1390    }
1391
1392    #[test]
1393    fn unreachable_mul_nodes_carry_no_wire() {
1394        let mut arena = ConstraintArena::<F>::new();
1395
1396        let a = arena.cell(ProgramCell::current(0));
1397        let b = arena.cell(ProgramCell::current(1));
1398        let live = arena.mul(a, b);
1399        let orphan = arena.cell(ProgramCell::current(2));
1400
1401        arena.mul(orphan, b);
1402
1403        let ast = ConstraintAst {
1404            arena,
1405            roots: vec![live],
1406            labels: vec![None],
1407        };
1408
1409        assert_eq!(mul_node_count(&ast), 1);
1410    }
1411
1412    #[test]
1413    fn shared_mul_node_counts_once() {
1414        let mut arena = ConstraintArena::<F>::new();
1415
1416        let a = arena.cell(ProgramCell::current(0));
1417        let b = arena.cell(ProgramCell::current(1));
1418        let shared = arena.mul(a, b);
1419        let left = arena.add(shared, a);
1420        let right = arena.add(shared, b);
1421
1422        let ast = ConstraintAst {
1423            arena,
1424            roots: vec![left, right],
1425            labels: vec![None, None],
1426        };
1427
1428        assert_eq!(mul_node_count(&ast), 1);
1429    }
1430
1431    #[test]
1432    fn nested_mul_nodes_each_carry_wire() {
1433        let mut arena = ConstraintArena::<F>::new();
1434
1435        let a = arena.cell(ProgramCell::current(0));
1436        let b = arena.cell(ProgramCell::current(1));
1437        let c = arena.cell(ProgramCell::current(2));
1438        let inner = arena.mul(a, b);
1439        let outer = arena.mul(inner, c);
1440
1441        let ast = ConstraintAst {
1442            arena,
1443            roots: vec![outer],
1444            labels: vec![None],
1445        };
1446
1447        assert_eq!(mul_node_count(&ast), 2);
1448    }
1449
1450    #[test]
1451    fn mul_under_sum_and_scale_is_reachable() {
1452        let mut arena = ConstraintArena::<F>::new();
1453
1454        let a = arena.cell(ProgramCell::current(0));
1455        let b = arena.cell(ProgramCell::current(1));
1456        let product = arena.mul(a, b);
1457        let scaled = arena.scale(F::ONE, product);
1458        let summed = arena.sum(vec![scaled, a]);
1459
1460        let ast = ConstraintAst {
1461            arena,
1462            roots: vec![summed],
1463            labels: vec![None],
1464        };
1465
1466        assert_eq!(mul_node_count(&ast), 1);
1467    }
1468
1469    #[test]
1470    fn busless_table_pays_for_no_h_claims() {
1471        let with_bus = shape(12, 100, 3, 1);
1472        let without = TableShape {
1473            num_buses: 0,
1474            ..with_bus
1475        };
1476
1477        let blind_units = 1;
1478        let claimed_sum_and_h_eval = 2;
1479        let h_claim_halves = 2;
1480
1481        assert_eq!(
1482            with_bus.masked_scalars(blind_units) - without.masked_scalars(blind_units),
1483            claimed_sum_and_h_eval + h_claim_halves
1484        );
1485    }
1486
1487    #[test]
1488    fn masked_scalars_match_transcript_tally() {
1489        let s = shape(12, 1977, 5, 3);
1490
1491        let blind_units = 1;
1492        let claimed_sums_and_h_evals = 2 * 3;
1493        let zerocheck_rounds = 12 * 5;
1494        let claims = 2 * (1977 + blind_units + 3);
1495        let trace_eval_rounds = 2 * 12;
1496
1497        assert_eq!(
1498            s.masked_scalars(blind_units),
1499            claimed_sums_and_h_evals + zerocheck_rounds + claims + trace_eval_rounds
1500        );
1501    }
1502
1503    #[test]
1504    fn every_bus_adds_hadamard_row() {
1505        let s = TableShape {
1506            mul_nodes: 40,
1507            ..shape(12, 100, 3, 3)
1508        };
1509
1510        assert_eq!(s.mul_wires(), 43);
1511    }
1512
1513    #[test]
1514    fn statement_sums_over_every_table() {
1515        let main = TableShape {
1516            mul_nodes: 40,
1517            ..shape(12, 100, 3, 2)
1518        };
1519        let chiplet = TableShape {
1520            mul_nodes: 7,
1521            ..shape(10, 30, 4, 1)
1522        };
1523
1524        let statement = OuterStatement::new(&[main, chiplet], 2);
1525
1526        assert_eq!(
1527            statement.masked_scalars,
1528            main.masked_scalars(2) + chiplet.masked_scalars(2)
1529        );
1530        assert_eq!(statement.mul_wires, 50);
1531    }
1532
1533    #[test]
1534    fn every_hadamard_wire_operand_is_pinned_by_row() {
1535        let claims: Vec<Flat<F>> = (0..4).map(|i| mix(100 + i)).collect();
1536        let record = bus_record(claims, mix(8), mix(9));
1537        let statement = OuterStatement {
1538            masked_scalars: 32,
1539            mul_wires: record.mul_wires as usize,
1540        };
1541
1542        let rows = assemble(vec![record], &statement).unwrap();
1543
1544        for mul in 0..statement.mul_wires as u32 {
1545            for role in [WireRole::Lhs, WireRole::Rhs, WireRole::Product] {
1546                assert!(
1547                    mentions(&rows.affine, mul, role) >= 1,
1548                    "wire {mul} {role:?}"
1549                );
1550            }
1551        }
1552    }
1553
1554    #[test]
1555    fn bus_operand_rows_hold_on_honest_wires() {
1556        let plain: Vec<Flat<F>> = (0..4).map(|i| mix(200 + i)).collect();
1557        let pad: Vec<Flat<F>> = (0..32).map(|i| mix(300 + i)).collect();
1558        let claims_masked: Vec<Flat<F>> = plain
1559            .iter()
1560            .enumerate()
1561            .map(|(i, &c)| c + pad[BUS_PAD_FIRST as usize + i])
1562            .collect();
1563
1564        let beta = mix(9);
1565        let h = mix(400);
1566        let h_masked = h + pad[BUS_H_PAD as usize];
1567        let key = plain[0] + beta * mix(BUS_PUBLIC);
1568
1569        let record = bus_record(claims_masked.clone(), h_masked, beta);
1570        let bus = &record.consistency.buses[0];
1571        let [lhs, rhs] = bus_operand_rows(bus, BUS_PAD_FIRST, beta);
1572
1573        let wire = |mul: u32, role: WireRole| -> Flat<F> {
1574            assert_eq!(mul, bus.h_wire);
1575            match role {
1576                WireRole::Lhs => h,
1577                WireRole::Rhs => key,
1578                WireRole::Product => h * key,
1579            }
1580        };
1581
1582        for row in [&lhs, &rhs] {
1583            let mut left = Flat::from_raw(F::ZERO);
1584            for &(u, coeff) in &row.unknowns {
1585                let value = match u {
1586                    Unknown::Pad(i) => pad[i as usize],
1587                    Unknown::Wire { mul, role } => wire(mul, role),
1588                };
1589
1590                left += coeff * value;
1591            }
1592
1593            let mut right = row.constant;
1594            for &(idx, coeff) in &row.claims {
1595                right += coeff * claims_masked[idx as usize];
1596            }
1597
1598            assert_eq!(left, right);
1599        }
1600
1601        assert_eq!(mentions(&[lhs], bus.h_wire, WireRole::Lhs), 1);
1602        assert_eq!(mentions(&[rhs], bus.h_wire, WireRole::Rhs), 1);
1603    }
1604
1605    #[test]
1606    fn bus_operand_rows_reject_free_operand() {
1607        let plain: Vec<Flat<F>> = (0..4).map(|i| mix(200 + i)).collect();
1608        let pad: Vec<Flat<F>> = (0..32).map(|i| mix(300 + i)).collect();
1609        let claims_masked: Vec<Flat<F>> = plain
1610            .iter()
1611            .enumerate()
1612            .map(|(i, &c)| c + pad[BUS_PAD_FIRST as usize + i])
1613            .collect();
1614
1615        let beta = mix(9);
1616        let h = mix(400);
1617        let h_masked = h + pad[BUS_H_PAD as usize];
1618        let key = plain[0] + beta * mix(BUS_PUBLIC);
1619
1620        let record = bus_record(claims_masked.clone(), h_masked, beta);
1621        let bus = &record.consistency.buses[0];
1622        let [lhs, _] = bus_operand_rows(bus, BUS_PAD_FIRST, beta);
1623
1624        let one = Flat::from_raw(F::ONE);
1625        let forged = |role: WireRole| match role {
1626            WireRole::Lhs => one,
1627            WireRole::Rhs | WireRole::Product => h * key,
1628        };
1629
1630        assert_eq!(
1631            forged(WireRole::Lhs) * forged(WireRole::Rhs),
1632            forged(WireRole::Product)
1633        );
1634
1635        let mut left = Flat::from_raw(F::ZERO);
1636        for &(u, coeff) in &lhs.unknowns {
1637            let value = match u {
1638                Unknown::Pad(i) => pad[i as usize],
1639                Unknown::Wire { role, .. } => forged(role),
1640            };
1641
1642            left += coeff * value;
1643        }
1644
1645        let mut right = lhs.constant;
1646        for &(idx, coeff) in &lhs.claims {
1647            right += coeff * claims_masked[idx as usize];
1648        }
1649
1650        assert_ne!(left, right);
1651    }
1652
1653    #[test]
1654    fn layout_holds_whole_statement() {
1655        let cfg = hekate_core::config::Config::prod();
1656
1657        for (scalars, wires) in [(1usize, 0usize), (4_075, 2_384), (12_000, 20_000)] {
1658            let geom = cfg.outer_geom(scalars, wires, 128).unwrap();
1659            let layout = OuterLayout::new(&geom, scalars, wires).unwrap();
1660
1661            assert!(layout.pad_rows * layout.message_len >= scalars);
1662            assert!(layout.aux_rows * layout.message_len >= wires);
1663
1664            let (r, _) = layout.slot(Unknown::Pad(scalars as u32 - 1));
1665            assert!(r < layout.pad_rows);
1666
1667            if wires > 0 {
1668                let (r, _) = layout.slot(Unknown::Wire {
1669                    mul: wires as u32 - 1,
1670                    role: WireRole::Product,
1671                });
1672
1673                assert!(r < layout.interleaved_mask());
1674            }
1675
1676            assert_eq!(layout.total_rows(), layout.quadratic_mask_hi() + 1);
1677        }
1678    }
1679
1680    #[test]
1681    fn every_mask_row_interleaved_test_reveals_is_uniform() {
1682        let cfg = hekate_core::config::Config::prod();
1683        let (scalars, wires) = (4_075usize, 2_384usize);
1684
1685        let geom = cfg.outer_geom(scalars, wires, 128).unwrap();
1686        let layout = OuterLayout::new(&geom, scalars, wires).unwrap();
1687
1688        let triples: Vec<[Flat<F>; 3]> = (0..wires as u128)
1689            .map(|i| [mix(i), mix(i + (1 << 20)), mix(i + (1 << 40))])
1690            .collect();
1691
1692        let filler: Vec<Flat<F>> = (0..aux_filler_len(&layout, geom.code_len) as u128)
1693            .map(|i| mix(i + (1 << 60)))
1694            .collect();
1695
1696        let rows = build_aux_rows(&layout, &triples, geom.code_len, &filler).unwrap();
1697
1698        let zero = Flat::from_raw(F::ZERO);
1699
1700        for mask in [
1701            layout.interleaved_mask(),
1702            layout.linear_mask_hi(),
1703            layout.quadratic_mask_hi(),
1704        ] {
1705            assert!(
1706                rows[mask - layout.pad_rows][..layout.message_len]
1707                    .iter()
1708                    .all(|v| *v != zero)
1709            );
1710        }
1711
1712        let quadratic = layout.quadratic_mask() - layout.pad_rows;
1713        assert!(
1714            rows[quadratic][..layout.message_len]
1715                .iter()
1716                .all(|v| *v == zero)
1717        );
1718
1719        assert!(
1720            build_aux_rows(&layout, &triples, geom.code_len, &filler[1..]).is_err(),
1721            "filler length is exact"
1722        );
1723    }
1724
1725    #[test]
1726    fn mask_form_matches_round_by_round_recursion() {
1727        let degree = 3;
1728        let initial: Vec<(Unknown, Flat<F>)> = (0..4u32)
1729            .map(|i| (Unknown::Pad(i), mix(500 + i as u128)))
1730            .collect();
1731        let round_pads: Vec<Vec<u32>> = (0..5u32)
1732            .map(|i| {
1733                (0..degree as u32)
1734                    .map(|j| 100 + i * degree as u32 + j)
1735                    .collect()
1736            })
1737            .collect();
1738        let round_refs: Vec<&[u32]> = round_pads.iter().map(|p| p.as_slice()).collect();
1739        let challenges: Vec<Flat<F>> = (0..5).map(|i| mix(600 + i)).collect();
1740
1741        let mut expected: Vec<(Unknown, Flat<F>)> = initial.clone();
1742        for (pads, &r) in round_refs.iter().zip(&challenges) {
1743            let weights = lagrange_weights::<F>(degree, r);
1744            for (_, coeff) in expected.iter_mut() {
1745                *coeff *= weights[0];
1746            }
1747
1748            expected.push((Unknown::Pad(pads[0]), weights[0]));
1749
1750            for (j, &pad) in pads.iter().enumerate() {
1751                expected.push((Unknown::Pad(pad), weights[j + 1]));
1752            }
1753        }
1754
1755        assert_eq!(
1756            sumcheck_mask_form(&initial, &round_refs, &challenges, degree),
1757            expected
1758        );
1759    }
1760
1761    #[test]
1762    fn h_claim_pin_row_ties_h_eval_to_committed_claim() {
1763        let pad: Vec<Flat<F>> = (0..32).map(|i| mix(300 + i)).collect();
1764        let one = Flat::from_raw(F::ONE);
1765        let h = mix(400);
1766        let h_masked = h + pad[BUS_H_PAD as usize];
1767
1768        let pin_form = [
1769            (Unknown::Pad(BUS_H_PAD), one),
1770            (Unknown::Pad(BUS_PAD_FIRST + H_CLAIM as u32), one),
1771        ];
1772
1773        for (h_claim, holds) in [(h, true), (h + one, false)] {
1774            let mut plain: Vec<Flat<F>> = (0..4).map(|i| mix(200 + i)).collect();
1775            plain[H_CLAIM] = h_claim;
1776
1777            let claims_masked: Vec<Flat<F>> = plain
1778                .iter()
1779                .enumerate()
1780                .map(|(i, &c)| c + pad[BUS_PAD_FIRST as usize + i])
1781                .collect();
1782
1783            let record = bus_record(claims_masked.clone(), h_masked, mix(9));
1784            let statement = OuterStatement {
1785                masked_scalars: 32,
1786                mul_wires: record.mul_wires as usize,
1787            };
1788
1789            let rows = assemble(vec![record], &statement).unwrap();
1790            let pins: Vec<&AffineRow<F>> = rows
1791                .affine
1792                .iter()
1793                .filter(|r| r.unknowns == pin_form)
1794                .collect();
1795
1796            assert_eq!(pins.len(), 1, "one pin row per bus");
1797            assert_eq!(pins[0].constant, h_masked + claims_masked[H_CLAIM]);
1798
1799            let left = pad[BUS_H_PAD as usize] + pad[BUS_PAD_FIRST as usize + H_CLAIM];
1800
1801            assert_eq!(left == pins[0].constant, holds);
1802        }
1803    }
1804
1805    #[test]
1806    fn batch_tensor_is_product_over_set_bits() {
1807        let a = F::from(0x1234_5678_9abc_def0u128);
1808        let b = F::from(0x0fed_cba9_8765_4321u128);
1809        let c = F::from(0xdead_beef_cafe_babeu128);
1810
1811        let (fa, fb, fc) = (a.to_hardware(), b.to_hardware(), c.to_hardware());
1812        let one = Flat::from_raw(F::ONE);
1813
1814        assert_eq!(expand_batch_tensor::<F>(&[], 1), vec![one]);
1815        assert_eq!(expand_batch_tensor(&[a], 2), vec![one, fa]);
1816        assert_eq!(expand_batch_tensor(&[a, b], 4), vec![one, fa, fb, fa * fb]);
1817
1818        assert_eq!(
1819            expand_batch_tensor(&[a, b, c], 8),
1820            vec![one, fa, fb, fa * fb, fc, fa * fc, fb * fc, fa * fb * fc,],
1821        );
1822    }
1823
1824    #[test]
1825    fn batch_tensor_truncates_below_power_of_two() {
1826        let a = F::from(0x1234_5678_9abc_def0u128);
1827        let b = F::from(0x0fed_cba9_8765_4321u128);
1828
1829        let full = expand_batch_tensor(&[a, b], 4);
1830
1831        for len in 1..=4 {
1832            assert_eq!(expand_batch_tensor(&[a, b], len), full[..len]);
1833        }
1834    }
1835}