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::{BoundaryConstraint, ConstraintAst, ConstraintExpr, ExprId};
7use crate::expander::{
8    PoolLayout, RING_BLIND_BITS, RingSwitchPlan, SlotLayout, claim_weights, eq_tensor_b,
9};
10use crate::linearized::{self, RingGadget, linearized_coeffs};
11use crate::permutation::{
12    BusKind, PermutationCheckSpec, RankClock, Source, eval_row_idx_byte_mle, eval_row_idx_le_mle,
13};
14use crate::predicate::{AffineRow, ClaimLayout, Form, Unknown, WireRole, compile};
15use crate::{Air, FixedColumn, ProgramInstance, ShapeEvaluator};
16use alloc::collections::BTreeMap;
17use alloc::string::String;
18use alloc::vec;
19use alloc::vec::Vec;
20use core::mem;
21use hekate_core::config::Config;
22use hekate_core::errors;
23use hekate_core::outer::{OUTER_MASK_ROWS, OuterGeometry};
24use hekate_core::poly::univariate::{MAX_POINTS, UnivariatePoly};
25use hekate_core::tensor::TensorProduct;
26use hekate_math::{Block128, Flat, HardwareField, TowerField};
27
28pub struct LinearBatch<F> {
29    pub weights: Vec<Vec<Flat<F>>>,
30    pub rows: Vec<usize>,
31    pub target: Flat<F>,
32}
33
34pub enum BusSource<F> {
35    Claim(u32),
36    Public(Flat<F>),
37}
38
39pub struct BusRow<F> {
40    pub h_pad: u32,
41    pub h_masked: Flat<F>,
42    pub h_wire: u32,
43    pub sources: Vec<BusSource<F>>,
44    pub selector: Option<u32>,
45    pub recv_selector: Option<u32>,
46    pub eq_lookup: Flat<F>,
47}
48
49pub struct ConsistencyInputs<F> {
50    /// Pad index of claim 0, matching `ClaimLayout`.
51    pub pad_first: u32,
52
53    pub eq_zc: Flat<F>,
54    pub alpha: Flat<F>,
55    pub gamma: Flat<F>,
56    pub beta: Flat<F>,
57
58    /// `(claim index, public value, eq(r_final, row))`.
59    pub boundary: Vec<(u32, Flat<F>, Flat<F>)>,
60
61    /// `(b_k claim index, b_k_next claim index)`.
62    pub telescope: Vec<(u32, u32)>,
63
64    pub buses: Vec<BusRow<F>>,
65
66    /// Mask form of the ZeroCheck's final claim,
67    /// and the masked value the replay reaches.
68    pub val_final_form: Vec<(Unknown, Flat<F>)>,
69    pub val_final_masked: Flat<F>,
70}
71
72impl<F: HardwareField> ConsistencyInputs<F> {
73    /// `α^k · eq_zc` for `k` in `0..roots`, the
74    /// weight root `k` carries in the consistency row.
75    fn root_weights(&self, roots: usize) -> impl Iterator<Item = Flat<F>> + '_ {
76        core::iter::successors(Some(self.eq_zc), |w| Some(*w * self.alpha)).take(roots)
77    }
78}
79
80/// One table's AST, bus specs, fixed pins and boundary
81/// constraints, borrowed where the `Air` getters clone.
82#[derive(Clone, Copy)]
83pub struct TableStatics<'a, F: TowerField> {
84    pub ast: &'a ConstraintAst<F>,
85    pub specs: &'a [(String, PermutationCheckSpec)],
86    pub fixed: &'a [FixedColumn<F>],
87    pub boundary: &'a [BoundaryConstraint<F>],
88}
89
90/// The parts behind [`TableStatics`], read once from an
91/// `Air` that does not hold them, such as a `Program`.
92pub struct TableParts<F: TowerField> {
93    pub ast: ConstraintAst<F>,
94    pub specs: Vec<(String, PermutationCheckSpec)>,
95    pub fixed: Vec<FixedColumn<F>>,
96    pub boundary: Vec<BoundaryConstraint<F>>,
97}
98
99impl<F: TowerField> TableParts<F> {
100    pub fn of(air: &impl Air<F>) -> Self {
101        Self {
102            ast: air.constraint_ast(),
103            specs: air.permutation_checks(),
104            fixed: air.fixed_columns(),
105            boundary: air.boundary_constraints(),
106        }
107    }
108
109    pub fn statics(&self) -> TableStatics<'_, F> {
110        TableStatics {
111            ast: &self.ast,
112            specs: &self.specs,
113            fixed: &self.fixed,
114            boundary: &self.boundary,
115        }
116    }
117}
118
119#[derive(Clone, Copy, Debug, Eq, PartialEq)]
120pub struct TableShape {
121    pub num_vars: usize,
122
123    /// Virtual columns, after expansion.
124    pub num_columns: usize,
125
126    pub sumcheck_degree: usize,
127    pub num_buses: usize,
128    pub mul_nodes: usize,
129    pub ring_units: bool,
130    pub whole_blinds: usize,
131}
132
133impl TableShape {
134    pub fn from_air<F: TowerField>(
135        air: &impl Air<F>,
136        statics: &TableStatics<'_, F>,
137    ) -> errors::Result<Self> {
138        let ast = statics.ast;
139        let num_buses = statics.specs.len();
140
141        ast.validate_order()?;
142
143        let mut sumcheck_degree = ast.max_degree() + 1;
144
145        if !statics.boundary.is_empty() {
146            sumcheck_degree = sumcheck_degree.max(2);
147        }
148
149        if num_buses > 0 {
150            sumcheck_degree = sumcheck_degree.max(3);
151        }
152
153        if sumcheck_degree + 1 > MAX_POINTS {
154            return Err(errors::Error::Protocol {
155                protocol: "air",
156                message: "sumcheck degree exceeds the interpolation stack bound",
157            });
158        }
159
160        if air.num_columns() != air.virtual_column_layout().len() {
161            return Err(errors::Error::Protocol {
162                protocol: "air",
163                message: "num_columns does not match the virtual column layout",
164            });
165        }
166
167        let entries = air.virtual_expander().map(|e| e.expansion_entries());
168        let plan = RingSwitchPlan::new(air.column_layout(), entries.as_deref(), 0, 0)?;
169
170        Ok(Self {
171            num_vars: 0,
172            num_columns: air.num_columns(),
173            sumcheck_degree,
174            num_buses,
175            mul_nodes: mul_node_count(ast),
176            ring_units: plan.has_ring(),
177            whole_blinds: 0,
178        })
179    }
180
181    pub fn at(self, layout: &SlotLayout) -> Self {
182        Self {
183            num_vars: layout.num_vars,
184            whole_blinds: layout.plan.whole_blinds(),
185            ..self
186        }
187    }
188
189    /// Eval claims per half: virtual columns, whole blind
190    /// units, the ring blind unit's planes, one `h` per bus.
191    pub fn eval_claims(&self) -> usize {
192        let ring_blind = match self.ring_units && self.whole_blinds > 0 {
193            true => RING_BLIND_BITS,
194            false => 0,
195        };
196
197        self.num_columns + self.whole_blinds + ring_blind + self.num_buses
198    }
199
200    /// Final sumcheck values are the masked running
201    /// claims themselves and consume no pad entry.
202    pub fn masked_scalars(&self) -> usize {
203        let zerocheck = self.num_vars * self.sumcheck_degree;
204        let claims = 2 * self.eval_claims();
205
206        2 * self.num_buses + zerocheck + claims
207    }
208
209    /// Each bus contributes one `h · key` product.
210    pub fn mul_wires(&self) -> usize {
211        self.mul_nodes + self.num_buses
212    }
213}
214
215#[derive(Clone, Copy, Debug, Eq, PartialEq)]
216pub struct OuterStatement {
217    pub masked_scalars: usize,
218    pub mul_wires: usize,
219}
220
221impl OuterStatement {
222    pub fn new(shapes: &[TableShape], pool_vars: usize) -> Self {
223        let tables: usize = shapes.iter().map(TableShape::masked_scalars).sum();
224        let wires: usize = shapes.iter().map(TableShape::mul_wires).sum();
225
226        let gadget = match shapes.iter().any(|s| s.ring_units) {
227            true => linearized::BITS - 1,
228            false => 0,
229        };
230
231        Self {
232            masked_scalars: tables + 2 * pool_vars,
233            mul_wires: wires + gadget,
234        }
235    }
236
237    /// Statement size for a whole proof.
238    pub fn for_tables<F: TowerField, A: Air<F>>(
239        main: &A,
240        main_num_vars: usize,
241        chiplets: &[ChipletDef<F>],
242        chiplet_num_vars: &[usize],
243        config: &Config,
244    ) -> errors::Result<Self> {
245        if chiplets.len() != chiplet_num_vars.len() {
246            return Err(errors::Error::Protocol {
247                protocol: "outer",
248                message: "one height per chiplet is required",
249            });
250        }
251
252        let main_parts = TableParts::of(main);
253
254        let mut shapes = Vec::with_capacity(chiplets.len() + 1);
255        let mut plans = Vec::with_capacity(chiplets.len() + 1);
256
257        for def in chiplets {
258            shapes.push(TableShape::from_air(def, &def.statics())?);
259            plans.push(table_plan(def, def.permutation_checks.len(), config)?);
260        }
261
262        shapes.push(TableShape::from_air(main, &main_parts.statics())?);
263        plans.push(table_plan(main, main_parts.specs.len(), config)?);
264
265        let heights: Vec<(&RingSwitchPlan, usize)> = plans
266            .iter()
267            .zip(chiplet_num_vars.iter().chain([&main_num_vars]))
268            .map(|(plan, &num_vars)| (plan, num_vars))
269            .collect();
270
271        let pool = PoolLayout::new(&heights, F::BITS, config);
272
273        let shapes: Vec<TableShape> = shapes
274            .into_iter()
275            .zip(&pool.tables)
276            .map(|(shape, layout)| shape.at(layout))
277            .collect();
278
279        Ok(Self::new(&shapes, pool.num_vars()))
280    }
281}
282
283/// Oracle row order: pad, then the `Lhs`, `Rhs` and
284/// `Product` wire blocks, then the five mask rows.
285#[derive(Clone, Copy, Debug, Eq, PartialEq)]
286pub struct OuterLayout {
287    pub message_len: usize,
288    pub pad_rows: usize,
289    pub aux_rows: usize,
290    pub domain_len: usize,
291}
292
293impl OuterLayout {
294    pub fn new(
295        geom: &OuterGeometry,
296        masked_scalars: usize,
297        mul_wires: usize,
298    ) -> errors::Result<Self> {
299        if geom.message_len == 0 {
300            return Err(errors::Error::Protocol {
301                protocol: "outer",
302                message: "message length must be positive",
303            });
304        }
305
306        Ok(Self {
307            message_len: geom.message_len,
308            pad_rows: masked_scalars.div_ceil(geom.message_len),
309            aux_rows: mul_wires.div_ceil(geom.message_len).max(1),
310            domain_len: geom.domain_len,
311        })
312    }
313
314    pub fn slot(&self, unknown: Unknown) -> (usize, usize) {
315        match unknown {
316            Unknown::Pad(i) => {
317                let i = i as usize;
318
319                (i / self.message_len, i % self.message_len)
320            }
321            Unknown::Wire { mul, role } => {
322                let mul = mul as usize;
323                let block = role as usize;
324
325                (
326                    self.pad_rows + block * self.aux_rows + mul / self.message_len,
327                    mul % self.message_len,
328                )
329            }
330        }
331    }
332
333    pub fn interleaved_mask(&self) -> usize {
334        self.pad_rows + 3 * self.aux_rows
335    }
336
337    pub fn linear_mask(&self) -> usize {
338        self.interleaved_mask() + 1
339    }
340
341    pub fn linear_mask_hi(&self) -> usize {
342        self.interleaved_mask() + 2
343    }
344
345    pub fn quadratic_mask(&self) -> usize {
346        self.interleaved_mask() + 3
347    }
348
349    pub fn quadratic_mask_hi(&self) -> usize {
350        self.interleaved_mask() + 4
351    }
352
353    pub fn total_rows(&self) -> usize {
354        self.interleaved_mask() + OUTER_MASK_ROWS
355    }
356
357    pub fn hadamard_triples(&self) -> Vec<[usize; 3]> {
358        (0..self.aux_rows)
359            .map(|i| {
360                [
361                    self.pad_rows + i,
362                    self.pad_rows + self.aux_rows + i,
363                    self.pad_rows + 2 * self.aux_rows + i,
364                ]
365            })
366            .collect()
367    }
368}
369
370pub struct EvalRecord<F> {
371    pub mask_form: Vec<(Unknown, Flat<F>)>,
372    pub claim_masked: Flat<F>,
373    pub fin: Flat<F>,
374    pub gadget: Option<RingGadget<F>>,
375}
376
377/// One table's share of the statement, in transcript order.
378pub struct TableRecord<'s, F> {
379    pub claims_masked: Vec<Flat<F>>,
380    pub ast: &'s ConstraintAst<F>,
381    pub consistency: ConsistencyInputs<F>,
382
383    /// `(pad, claim, public value)` per fixed column.
384    pub fixed: Vec<(u32, u32, Flat<F>)>,
385
386    /// `(bus id, pad, masked claimed sum)` per bus.
387    pub claimed_sums: Vec<(String, u32, Flat<F>)>,
388
389    pub shape: TableShape,
390}
391
392impl<F: TowerField + HardwareField> TableRecord<'_, F> {
393    fn affine_rows(&self) -> usize {
394        let buses = self.consistency.buses.len();
395
396        2 * self.shape.mul_nodes + 1 + 3 * buses + self.fixed.len()
397    }
398
399    /// `f` on each row after the predicate rows, in statement order;
400    /// the consistency row comes without its root terms.
401    fn for_each_tail_row(&self, mut f: impl FnMut(AffineRow<F>)) {
402        let consistency = &self.consistency;
403        let pad_first = consistency.pad_first;
404
405        f(consistency_row(self.ast.roots.len(), consistency));
406
407        for bus in &consistency.buses {
408            let [lhs, rhs] = bus_operand_rows(bus, pad_first, consistency.beta);
409
410            f(lhs);
411            f(rhs);
412        }
413
414        for &(pad, claim, public) in &self.fixed {
415            f(fixed_column_row(pad, claim, public));
416        }
417
418        let half = (self.claims_masked.len() / 2) as u32;
419        let num_buses = consistency.buses.len() as u32;
420
421        for (k, bus) in consistency.buses.iter().enumerate() {
422            let claim = half - num_buses + k as u32;
423
424            f(h_claim_pin_row(
425                bus.h_pad,
426                bus.h_masked,
427                pad_first + claim,
428                claim,
429            ));
430        }
431    }
432}
433
434/// The assembled statement: affine rows, and each ring
435/// gadget with the statement row its chain starts at.
436pub struct OuterRows<'r, F> {
437    pub affine: Vec<AffineRow<F>>,
438    pub gadgets: Vec<(usize, &'r RingGadget<F>)>,
439}
440
441/// `H_0` of an eval sumcheck over the pad.
442pub struct InitialForm<F> {
443    pub form: Vec<(Unknown, Flat<F>)>,
444    pub gadget: Option<RingGadget<F>>,
445}
446
447/// First pad index of each masked group of a table,
448/// in the prover's masking order.
449#[derive(Clone, Copy, Debug)]
450pub struct TablePads {
451    pub claimed_sums: u32,
452    pub zerocheck: u32,
453    pub h_evals: u32,
454    pub claims: u32,
455}
456
457/// Public data of the pool's eval sumcheck on masked claims.
458pub struct EvalInputs<'a, F> {
459    pub pool: &'a PoolLayout,
460    pub eta: F,
461    pub rho: &'a [F],
462    pub r_mix: &'a [Block128],
463    pub shifted_claims: bool,
464    pub challenges: &'a [Flat<F>],
465    pub claim_masked: Flat<F>,
466    pub fin: Flat<F>,
467    pub claim_pads: &'a [u32],
468    pub round_first: u32,
469    pub first_wire: u32,
470}
471
472/// Everything a table's record is built from; both
473/// sides hold all of it after the table's eval phase.
474pub struct TableInputs<'a, 's, F: TowerField> {
475    pub statics: TableStatics<'s, F>,
476    pub shape: TableShape,
477    pub instance: &'a ProgramInstance<F>,
478    pub blinding_columns: usize,
479
480    pub alpha: Flat<F>,
481    pub gamma: Flat<F>,
482    pub beta: Flat<F>,
483    pub r_zerocheck: &'a [Flat<F>],
484    pub r_final: &'a [Flat<F>],
485    pub lookup_bus_points: &'a BTreeMap<String, Vec<Flat<F>>>,
486    pub clocks: &'a [Option<RankClock>],
487
488    pub claimed_sums_masked: &'a [(String, F)],
489    pub h_evals_masked: &'a [(String, F)],
490    pub val_final_masked: Flat<F>,
491    pub claims_masked: Vec<Flat<F>>,
492
493    pub pads: TablePads,
494}
495
496/// Weight rows and target of the linear
497/// test, summed over scaled affine rows.
498struct Accumulator<'l, F> {
499    layout: &'l OuterLayout,
500    weights: Vec<Vec<Flat<F>>>,
501    target: Flat<F>,
502}
503
504impl<'l, F: TowerField + HardwareField> Accumulator<'l, F> {
505    fn new(layout: &'l OuterLayout) -> Self {
506        Self {
507            layout,
508            weights: vec![vec![Flat::from_raw(F::ZERO); layout.message_len]; layout.total_rows()],
509            target: Flat::from_raw(F::ZERO),
510        }
511    }
512
513    fn add(&mut self, unknown: Unknown, value: Flat<F>) {
514        let (row, col) = self.layout.slot(unknown);
515        self.weights[row][col] += value;
516    }
517
518    /// Adds `scale` times one table's `row`: claims read
519    /// from `claims`, wires shifted by `mul_offset`.
520    fn row(&mut self, row: &AffineRow<F>, scale: Flat<F>, claims: &[Flat<F>], mul_offset: u32) {
521        for &(unknown, coeff) in &row.unknowns {
522            self.add(unknown.shifted(mul_offset), scale * coeff);
523        }
524
525        let mut constant = row.constant;
526        for &(idx, coeff) in &row.claims {
527            constant += coeff * claims[idx as usize];
528        }
529
530        self.target += scale * constant;
531    }
532
533    /// The weight rows that hold a nonzero
534    /// entry, with their oracle row indices.
535    fn batch(mut self) -> LinearBatch<F> {
536        let zero = Flat::from_raw(F::ZERO);
537        let rows: Vec<usize> = (0..self.weights.len())
538            .filter(|&r| self.weights[r].iter().any(|v| *v != zero))
539            .collect();
540
541        let weights = rows
542            .iter()
543            .map(|&r| mem::take(&mut self.weights[r]))
544            .collect();
545
546        LinearBatch {
547            weights,
548            rows,
549            target: self.target,
550        }
551    }
552}
553
554struct Sweep<F> {
555    live: Vec<bool>,
556    adjoint: Vec<Flat<F>>,
557}
558
559impl<F: TowerField + HardwareField> Sweep<F> {
560    fn new() -> Self {
561        Self {
562            live: Vec::new(),
563            adjoint: Vec::new(),
564        }
565    }
566
567    /// Adds the table's predicate rows at `scales` and its
568    /// root terms at `consistency` in one pass from the last
569    /// node to the first, each node handing down its adjoint.
570    fn run(
571        &mut self,
572        acc: &mut Accumulator<'_, F>,
573        record: &TableRecord<'_, F>,
574        scales: &[Flat<F>],
575        consistency: Flat<F>,
576        mul_offset: u32,
577    ) -> errors::Result<()> {
578        let ast = record.ast;
579        let claims = &record.claims_masked;
580
581        let layout = ClaimLayout {
582            pad_first: record.consistency.pad_first,
583            half: (claims.len() / 2) as u32,
584        };
585
586        let n = ast.arena.len();
587
588        self.live.clear();
589        self.live.resize(n, false);
590        self.adjoint.clear();
591        self.adjoint.resize(n, Flat::from_raw(F::ZERO));
592
593        let roots = record.consistency.root_weights(ast.roots.len());
594        for (root, weight) in ast.roots.iter().zip(roots) {
595            self.pass(*root, consistency * weight);
596        }
597
598        let mut muls = scales.len() / 2;
599
600        // Children sit at lower indices than their parent
601        for i in (0..n).rev() {
602            if !self.live[i] {
603                continue;
604            }
605
606            let adjoint = self.adjoint[i];
607
608            match ast.arena.get(ExprId(i as u32)) {
609                ConstraintExpr::Cell(cell) => {
610                    let idx = layout.index(*cell);
611
612                    acc.add(Unknown::Pad(layout.pad_first + idx), adjoint);
613
614                    acc.target += adjoint * claims[idx as usize];
615                }
616                ConstraintExpr::Const(v) => acc.target += adjoint * v.to_hardware(),
617                ConstraintExpr::Add(a, b) => {
618                    self.pass(*a, adjoint);
619                    self.pass(*b, adjoint);
620                }
621                ConstraintExpr::Scale(coeff, a) => self.pass(*a, adjoint * coeff.to_hardware()),
622                ConstraintExpr::Sum(children) => {
623                    for c in children {
624                        self.pass(*c, adjoint);
625                    }
626                }
627                ConstraintExpr::Mul(a, b) => {
628                    muls = muls.checked_sub(1).ok_or(errors::Error::Protocol {
629                        protocol: "outer",
630                        message: "table shape undercounts the AST's Mul nodes",
631                    })?;
632
633                    let lhs = scales[2 * muls];
634                    let rhs = scales[2 * muls + 1];
635                    let mul = mul_offset + muls as u32;
636
637                    acc.add(
638                        Unknown::Wire {
639                            mul,
640                            role: WireRole::Lhs,
641                        },
642                        lhs,
643                    );
644                    acc.add(
645                        Unknown::Wire {
646                            mul,
647                            role: WireRole::Rhs,
648                        },
649                        rhs,
650                    );
651                    acc.add(
652                        Unknown::Wire {
653                            mul,
654                            role: WireRole::Product,
655                        },
656                        adjoint,
657                    );
658
659                    self.pass(*a, lhs);
660                    self.pass(*b, rhs);
661                }
662            }
663        }
664
665        match muls {
666            0 => Ok(()),
667            _ => Err(errors::Error::Protocol {
668                protocol: "outer",
669                message: "table shape overcounts the AST's Mul nodes",
670            }),
671        }
672    }
673
674    fn pass(&mut self, child: ExprId, value: Flat<F>) {
675        let c = child.0 as usize;
676
677        self.live[c] = true;
678        self.adjoint[c] += value;
679    }
680}
681
682pub fn eval_initial_form<F: HardwareField + Into<Block128> + From<u128>>(
683    pool: &PoolLayout,
684    eta: F,
685    rho: &[F],
686    r_mix: &[Block128],
687    shifted_claims: bool,
688    claim_pads: &[u32],
689    first_wire: u32,
690) -> errors::Result<InitialForm<F>> {
691    let mut weights = Vec::with_capacity(claim_pads.len());
692    for layout in &pool.tables {
693        let rho = rho
694            .get(..layout.pack_vars())
695            .ok_or(errors::Error::Protocol {
696                protocol: "outer",
697                message: "fewer block challenges than a table packs",
698            })?;
699
700        weights.extend(claim_weights(layout, eta, rho, shifted_claims)?);
701    }
702
703    if weights.len() != claim_pads.len() {
704        return Err(errors::Error::Protocol {
705            protocol: "outer",
706            message: "one pad entry per eval claim is required",
707        });
708    }
709
710    let mut form = Vec::new();
711    let mut ring = Vec::new();
712
713    for ((is_ring, weight), &pad) in weights.iter().zip(claim_pads) {
714        let a: Flat<F> = F::from(weight.0).to_hardware();
715        match is_ring {
716            true => ring.push((pad, a)),
717            false => form.push((Unknown::Pad(pad), a)),
718        }
719    }
720
721    if ring.is_empty() {
722        return Ok(InitialForm { form, gadget: None });
723    }
724
725    let mu: Vec<Flat<F>> = linearized_coeffs(&eq_tensor_b(r_mix))
726        .iter()
727        .map(|m| F::from(m.to_tower().0).to_hardware())
728        .collect();
729
730    let gadget = RingGadget::new(ring, mu, first_wire);
731
732    form.extend(gadget.delta_form());
733
734    Ok(InitialForm {
735        form,
736        gadget: Some(gadget),
737    })
738}
739
740/// The eval row's mask form propagated
741/// through the rounds, plus the ring gadget.
742pub fn eval_record<F>(inputs: &EvalInputs<'_, F>) -> errors::Result<EvalRecord<F>>
743where
744    F: HardwareField + Into<Block128> + From<u128>,
745{
746    let initial = eval_initial_form(
747        inputs.pool,
748        inputs.eta,
749        inputs.rho,
750        inputs.r_mix,
751        inputs.shifted_claims,
752        inputs.claim_pads,
753        inputs.first_wire,
754    )?;
755
756    let round_pads: Vec<[u32; 2]> = (0..inputs.challenges.len() as u32)
757        .map(|i| [inputs.round_first + 2 * i, inputs.round_first + 2 * i + 1])
758        .collect();
759    let round_refs: Vec<&[u32]> = round_pads.iter().map(|p| p.as_slice()).collect();
760
761    Ok(EvalRecord {
762        mask_form: sumcheck_mask_form(&initial.form, &round_refs, inputs.challenges, 2),
763        claim_masked: inputs.claim_masked,
764        fin: inputs.fin,
765        gadget: initial.gadget,
766    })
767}
768
769pub fn table_record<'s, F>(inputs: TableInputs<'_, 's, F>) -> errors::Result<TableRecord<'s, F>>
770where
771    F: TowerField + HardwareField,
772{
773    let TableStatics {
774        ast,
775        specs: bus_specs,
776        fixed: fixed_columns,
777        boundary: boundary_constraints,
778    } = inputs.statics;
779
780    let shape = inputs.shape;
781    let num_vars = shape.num_vars;
782    let trace_width = shape.num_columns;
783    let r_final = inputs.r_final;
784    let pad_first = inputs.pads.claims;
785    let half = shape.eval_claims() as u32;
786
787    if inputs.claims_masked.len() != 2 * half as usize {
788        return Err(errors::Error::Protocol {
789            protocol: "outer",
790            message: "eval claim count does not match the table shape",
791        });
792    }
793
794    let eq_row_checker = TensorProduct::new(r_final.to_vec());
795
796    let mut shapes = ShapeEvaluator::new(r_final);
797    let mut fixed = Vec::new();
798
799    for fc in fixed_columns {
800        if fc.col_idx >= trace_width {
801            return Err(errors::Error::Protocol {
802                protocol: "outer",
803                message: "fixed column col_idx out of trace_width",
804            });
805        }
806
807        fixed.push((
808            pad_first + fc.col_idx as u32,
809            fc.col_idx as u32,
810            shapes.evaluate(&fc.shape),
811        ));
812    }
813
814    let mut boundary = Vec::with_capacity(boundary_constraints.len());
815    for bc in boundary_constraints {
816        if bc.col_idx >= trace_width || bc.row_idx >= 1 << num_vars {
817            return Err(errors::Error::Protocol {
818                protocol: "outer",
819                message: "boundary constraint out of the table",
820            });
821        }
822
823        boundary.push((
824            bc.col_idx as u32,
825            bc.resolve_target(inputs.instance)?.to_hardware(),
826            eq_row_checker.evaluate_at_index(bc.row_idx),
827        ));
828    }
829
830    let mut buses = Vec::with_capacity(bus_specs.len());
831    for (spec_idx, (bus_id, spec)) in bus_specs.iter().enumerate() {
832        let eq_lookup = match spec.kind {
833            BusKind::Permutation => Flat::from_raw(F::ONE),
834            BusKind::Lookup => {
835                let r_bus =
836                    inputs
837                        .lookup_bus_points
838                        .get(bus_id)
839                        .ok_or(errors::Error::Protocol {
840                            protocol: "outer",
841                            message: "lookup bus spec missing r_bus point",
842                        })?;
843
844                if r_bus.len() < num_vars {
845                    return Err(errors::Error::Protocol {
846                        protocol: "outer",
847                        message: "r_bus shorter than table num_vars",
848                    });
849                }
850
851                let one = Flat::from_raw(F::ONE);
852                let mut eq_r_hi_at_0 = one;
853                for r_j in &r_bus[num_vars..] {
854                    eq_r_hi_at_0 *= one - *r_j;
855                }
856
857                eq_r_hi_at_0 * TensorProduct::evaluate_eq_slice(&r_bus[0..num_vars], r_final)
858            }
859        };
860
861        let mut sources = Vec::new();
862        for (source, _) in &spec.sources {
863            match source {
864                Source::Column(col) | Source::PhaseColumn(col) => {
865                    sources.push(BusSource::Claim(*col as u32))
866                }
867                Source::Columns(indices) => {
868                    sources.extend(indices.iter().map(|&col| BusSource::Claim(col as u32)));
869                }
870                Source::Const(v) => sources.push(BusSource::Public(F::from(*v).to_hardware())),
871                Source::RowIndexLeBytes(n) => {
872                    sources.push(BusSource::Public(eval_row_idx_le_mle::<F>(*n, r_final)));
873                }
874                Source::RowIndexByte(n) => {
875                    sources.push(BusSource::Public(eval_row_idx_byte_mle::<F>(*n, r_final)));
876                }
877                Source::EmitRank(_) => {
878                    let clock = RankClock::for_spec(inputs.clocks, spec_idx)?;
879                    sources.push(BusSource::Public(clock.evaluate(r_final)?));
880                }
881            }
882        }
883
884        buses.push(BusRow {
885            h_pad: inputs.pads.h_evals + spec_idx as u32,
886            h_masked: inputs.h_evals_masked[spec_idx].1.to_hardware(),
887            h_wire: (shape.mul_nodes + spec_idx) as u32,
888            sources,
889            selector: spec.selector.map(|i| i as u32),
890            recv_selector: spec.recv_selector.map(|i| i as u32),
891            eq_lookup,
892        });
893    }
894
895    let logup_alpha_offset = ast.roots.len() + boundary_constraints.len() + inputs.blinding_columns;
896
897    let mut alpha_pow = Flat::from_raw(F::ONE);
898    for _ in 0..logup_alpha_offset {
899        alpha_pow *= inputs.alpha;
900    }
901
902    let mut initial = Vec::with_capacity(bus_specs.len());
903    for k in 0..bus_specs.len() as u32 {
904        initial.push((Unknown::Pad(inputs.pads.claimed_sums + k), alpha_pow));
905
906        alpha_pow *= inputs.alpha;
907    }
908
909    let degree = shape.sumcheck_degree;
910    let round_pads: Vec<Vec<u32>> = (0..num_vars)
911        .map(|i| {
912            let first = inputs.pads.zerocheck + (i * degree) as u32;
913
914            (0..degree as u32).map(|j| first + j).collect()
915        })
916        .collect();
917    let round_refs: Vec<&[u32]> = round_pads.iter().map(|p| p.as_slice()).collect();
918
919    let consistency = ConsistencyInputs {
920        pad_first,
921        eq_zc: TensorProduct::evaluate_eq_slice(inputs.r_zerocheck, r_final),
922        alpha: inputs.alpha,
923        gamma: inputs.gamma,
924        beta: inputs.beta,
925        boundary,
926        telescope: (0..inputs.blinding_columns as u32)
927            .map(|k| {
928                let b = trace_width as u32 + k;
929
930                (b, half + b)
931            })
932            .collect(),
933        buses,
934        val_final_form: sumcheck_mask_form(&initial, &round_refs, r_final, degree),
935        val_final_masked: inputs.val_final_masked,
936    };
937
938    Ok(TableRecord {
939        claims_masked: inputs.claims_masked,
940        ast,
941        consistency,
942        fixed,
943        claimed_sums: inputs
944            .claimed_sums_masked
945            .iter()
946            .enumerate()
947            .map(|(k, (bus_id, sum))| {
948                (
949                    bus_id.clone(),
950                    inputs.pads.claimed_sums + k as u32,
951                    sum.to_hardware(),
952                )
953            })
954            .collect(),
955        shape,
956    })
957}
958
959/// Every affine row of the statement, claims folded
960/// into constants and wire indices offset per table.
961pub fn assemble<'r, F: TowerField + HardwareField>(
962    records: &'r [TableRecord<'_, F>],
963    eval: &'r EvalRecord<F>,
964    statement: &OuterStatement,
965) -> errors::Result<OuterRows<'r, F>> {
966    let mut affine = Vec::new();
967    let mut gadgets = Vec::new();
968    let mut offset = 0u32;
969
970    for record in records {
971        let first = affine.len();
972
973        let layout = ClaimLayout {
974            pad_first: record.consistency.pad_first,
975            half: (record.claims_masked.len() / 2) as u32,
976        };
977
978        let predicate = compile(record.ast, layout);
979        let consistency = first + predicate.affine.len();
980
981        affine.extend(predicate.affine);
982
983        record.for_each_tail_row(|row| affine.push(row));
984
985        let weights = record.consistency.root_weights(predicate.roots.len());
986        for (form, weight) in predicate.roots.iter().zip(weights) {
987            scale_into(&mut affine[consistency], form, weight);
988        }
989
990        for row in &mut affine[first..] {
991            resolve_table_row(row, &record.claims_masked, offset);
992        }
993
994        offset += record.shape.mul_wires() as u32;
995    }
996
997    affine.push(eval_final_row(&eval.mask_form, eval.claim_masked, eval.fin));
998
999    if let Some(gadget) = &eval.gadget {
1000        gadgets.push((affine.len(), gadget));
1001        affine.extend(gadget.wire_rows());
1002    }
1003
1004    if offset as usize + eval_wires(eval) != statement.mul_wires {
1005        return Err(errors::Error::Protocol {
1006            protocol: "outer",
1007            message: "table wire counts do not sum to the statement",
1008        });
1009    }
1010
1011    for endpoints in bus_endpoints(records).values() {
1012        affine.push(bus_sum_row(endpoints));
1013    }
1014
1015    Ok(OuterRows { affine, gadgets })
1016}
1017
1018/// Affine rows of the statement [`assemble`] builds: every
1019/// table's, the pool's eval rows, one bus-sum row per bus id.
1020pub fn statement_rows<F: TowerField + HardwareField>(
1021    records: &[TableRecord<'_, F>],
1022    eval: &EvalRecord<F>,
1023    statement: &OuterStatement,
1024) -> errors::Result<usize> {
1025    let tables_wires: usize = records.iter().map(|r| r.shape.mul_wires()).sum();
1026
1027    if tables_wires + eval_wires(eval) != statement.mul_wires {
1028        return Err(errors::Error::Protocol {
1029            protocol: "outer",
1030            message: "table wire counts do not sum to the statement",
1031        });
1032    }
1033
1034    let tables: usize = records.iter().map(TableRecord::affine_rows).sum();
1035
1036    Ok(tables + eval_rows(eval) + bus_endpoints(records).len())
1037}
1038
1039pub fn reachable<F: TowerField>(ast: &ConstraintAst<F>) -> Vec<bool> {
1040    let n = ast.arena.len();
1041    let mut live = vec![false; n];
1042
1043    for root in &ast.roots {
1044        live[root.0 as usize] = true;
1045    }
1046
1047    // Children sit at lower indices than their parent
1048    for i in (0..n).rev() {
1049        if !live[i] {
1050            continue;
1051        }
1052
1053        let mut mark = |id: &ExprId| live[id.0 as usize] = true;
1054        match ast.arena.get(ExprId(i as u32)) {
1055            ConstraintExpr::Cell(_) | ConstraintExpr::Const(_) => {}
1056            ConstraintExpr::Add(a, b) | ConstraintExpr::Mul(a, b) => {
1057                mark(a);
1058                mark(b);
1059            }
1060            ConstraintExpr::Scale(_, a) => mark(a),
1061            ConstraintExpr::Sum(children) => children.iter().for_each(mark),
1062        }
1063    }
1064
1065    live
1066}
1067
1068pub fn mul_node_count<F: TowerField>(ast: &ConstraintAst<F>) -> usize {
1069    reachable(ast)
1070        .iter()
1071        .enumerate()
1072        .filter(|&(i, &live)| {
1073            live && matches!(ast.arena.get(ExprId(i as u32)), ConstraintExpr::Mul(..))
1074        })
1075        .count()
1076}
1077
1078/// The consistency row without its root terms:
1079/// root `k` enters at `α^k · eq_zc`.
1080pub fn consistency_row<F: TowerField + HardwareField>(
1081    roots: usize,
1082    inputs: &ConsistencyInputs<F>,
1083) -> AffineRow<F> {
1084    let mut row = AffineRow {
1085        unknowns: inputs.val_final_form.to_vec(),
1086        claims: Vec::new(),
1087        constant: inputs.val_final_masked,
1088    };
1089
1090    let mut alpha_pow = Flat::from_raw(F::ONE);
1091
1092    for _ in 0..roots {
1093        alpha_pow *= inputs.alpha;
1094    }
1095
1096    for &(claim, public, eq_row) in &inputs.boundary {
1097        let weight = eq_row * alpha_pow;
1098
1099        push_claim(&mut row, inputs.pad_first, claim, weight);
1100        row.constant += public * weight;
1101
1102        alpha_pow *= inputs.alpha;
1103    }
1104
1105    for &(b, b_next) in &inputs.telescope {
1106        push_claim(&mut row, inputs.pad_first, b, alpha_pow);
1107        push_claim(&mut row, inputs.pad_first, b_next, alpha_pow);
1108
1109        alpha_pow *= inputs.alpha;
1110    }
1111
1112    for bus in &inputs.buses {
1113        let consistency = alpha_pow * inputs.eq_zc;
1114        let gamma_weight = consistency * inputs.gamma;
1115
1116        row.unknowns.push((Unknown::Pad(bus.h_pad), gamma_weight));
1117
1118        row.constant += gamma_weight * bus.h_masked;
1119
1120        row.unknowns.push((
1121            Unknown::Wire {
1122                mul: bus.h_wire,
1123                role: WireRole::Product,
1124            },
1125            consistency,
1126        ));
1127
1128        match bus.selector {
1129            Some(s) => push_claim(&mut row, inputs.pad_first, s, consistency),
1130            None => row.constant += consistency,
1131        }
1132
1133        if let Some(s) = bus.recv_selector {
1134            push_claim(&mut row, inputs.pad_first, s, consistency);
1135        }
1136
1137        let lookup_weight = alpha_pow * bus.eq_lookup;
1138
1139        row.unknowns.push((Unknown::Pad(bus.h_pad), lookup_weight));
1140
1141        row.constant += lookup_weight * bus.h_masked;
1142
1143        alpha_pow *= inputs.alpha;
1144    }
1145
1146    row
1147}
1148
1149/// Pad form of `claim'_n - claim_n` for a sumcheck whose
1150/// replay reconstructs `p'_0 = claim' - p'_1` each round:
1151/// entries of round `i` carry `Π_{j>i} L_0(r_j)`.
1152pub fn sumcheck_mask_form<F: TowerField + HardwareField>(
1153    initial: &[(Unknown, Flat<F>)],
1154    round_pads: &[&[u32]],
1155    challenges: &[Flat<F>],
1156    degree: usize,
1157) -> Vec<(Unknown, Flat<F>)> {
1158    let rounds: Vec<Vec<Flat<F>>> = challenges
1159        .iter()
1160        .map(|&r| lagrange_weights::<F>(degree, r))
1161        .collect();
1162
1163    let mut suffix = vec![Flat::from_raw(F::ONE); rounds.len() + 1];
1164    for i in (0..rounds.len()).rev() {
1165        suffix[i] = suffix[i + 1] * rounds[i][0];
1166    }
1167
1168    let mut form: Vec<(Unknown, Flat<F>)> = initial
1169        .iter()
1170        .map(|&(u, coeff)| (u, coeff * suffix[0]))
1171        .collect();
1172
1173    for (i, (pads, weights)) in round_pads.iter().zip(&rounds).enumerate() {
1174        let after = suffix[i + 1];
1175
1176        form.push((Unknown::Pad(pads[0]), weights[0] * after));
1177
1178        for (j, &pad) in pads.iter().enumerate() {
1179            form.push((Unknown::Pad(pad), weights[j + 1] * after));
1180        }
1181    }
1182
1183    form
1184}
1185
1186/// `claim_n = fin` for an eval sumcheck: `fin` is public,
1187/// `claim_n` is the masked running claim minus its form.
1188pub fn eval_final_row<F: TowerField + HardwareField>(
1189    mask_form: &[(Unknown, Flat<F>)],
1190    claim_masked: Flat<F>,
1191    fin: Flat<F>,
1192) -> AffineRow<F> {
1193    AffineRow {
1194        unknowns: mask_form.to_vec(),
1195        claims: Vec::new(),
1196        constant: claim_masked + fin,
1197    }
1198}
1199
1200/// `c[claim] = public`, i.e. `h[pad] = c'[claim] + public`.
1201pub fn fixed_column_row<F: TowerField + HardwareField>(
1202    pad: u32,
1203    claim: u32,
1204    public: Flat<F>,
1205) -> AffineRow<F> {
1206    let one = Flat::from_raw(F::ONE);
1207
1208    AffineRow {
1209        unknowns: vec![(Unknown::Pad(pad), one)],
1210        claims: vec![(claim, one)],
1211        constant: public,
1212    }
1213}
1214
1215/// `c[claim] = h_eval`: the eval argument's base `h`
1216/// claim and the reported `h_eval` unmask to one value.
1217pub fn h_claim_pin_row<F: TowerField + HardwareField>(
1218    h_pad: u32,
1219    h_masked: Flat<F>,
1220    claim_pad: u32,
1221    claim: u32,
1222) -> AffineRow<F> {
1223    let one = Flat::from_raw(F::ONE);
1224
1225    AffineRow {
1226        unknowns: vec![(Unknown::Pad(h_pad), one), (Unknown::Pad(claim_pad), one)],
1227        claims: vec![(claim, one)],
1228        constant: h_masked,
1229    }
1230}
1231
1232/// `Σ claimed_sum = 0` over one bus id:
1233/// `(pad, masked sum)` per endpoint.
1234pub fn bus_sum_row<F: TowerField + HardwareField>(endpoints: &[(u32, Flat<F>)]) -> AffineRow<F> {
1235    let mut constant = Flat::from_raw(F::ZERO);
1236    for &(_, masked) in endpoints {
1237        constant += masked;
1238    }
1239
1240    AffineRow {
1241        unknowns: endpoints
1242            .iter()
1243            .map(|&(pad, _)| (Unknown::Pad(pad), Flat::from_raw(F::ONE)))
1244            .collect(),
1245        claims: Vec::new(),
1246        constant,
1247    }
1248}
1249
1250/// Folds a table's public masked claims into the constant
1251/// and shifts its table-local wire indices by `mul_offset`.
1252pub fn resolve_table_row<F: TowerField + HardwareField>(
1253    row: &mut AffineRow<F>,
1254    claims: &[Flat<F>],
1255    mul_offset: u32,
1256) {
1257    for &(idx, coeff) in &row.claims {
1258        row.constant += coeff * claims[idx as usize];
1259    }
1260
1261    row.claims.clear();
1262
1263    for (unknown, _) in row.unknowns.iter_mut() {
1264        *unknown = unknown.shifted(mul_offset);
1265    }
1266}
1267
1268pub fn bus_key_form<F: TowerField + HardwareField>(
1269    bus: &BusRow<F>,
1270    pad_first: u32,
1271    beta: Flat<F>,
1272) -> Form<F> {
1273    let mut form = Form::default();
1274    let mut beta_pow = Flat::from_raw(F::ONE);
1275
1276    for source in &bus.sources {
1277        match source {
1278            BusSource::Claim(idx) => {
1279                form.unknowns
1280                    .push((Unknown::Pad(pad_first + idx), beta_pow));
1281                form.claims.push((*idx, beta_pow));
1282            }
1283            BusSource::Public(v) => form.constant += *v * beta_pow,
1284        }
1285
1286        beta_pow *= beta;
1287    }
1288
1289    form
1290}
1291
1292/// Pins the bus Hadamard wire `h · key`:
1293/// `Lhs = h` and `Rhs = key` over the pad.
1294pub fn bus_operand_rows<F: TowerField + HardwareField>(
1295    bus: &BusRow<F>,
1296    pad_first: u32,
1297    beta: Flat<F>,
1298) -> [AffineRow<F>; 2] {
1299    let one = Flat::from_raw(F::ONE);
1300
1301    let lhs = AffineRow {
1302        unknowns: vec![
1303            (
1304                Unknown::Wire {
1305                    mul: bus.h_wire,
1306                    role: WireRole::Lhs,
1307                },
1308                one,
1309            ),
1310            (Unknown::Pad(bus.h_pad), one),
1311        ],
1312        claims: Vec::new(),
1313        constant: bus.h_masked,
1314    };
1315
1316    let key = bus_key_form(bus, pad_first, beta);
1317
1318    let mut unknowns = key.unknowns;
1319    unknowns.push((
1320        Unknown::Wire {
1321            mul: bus.h_wire,
1322            role: WireRole::Rhs,
1323        },
1324        one,
1325    ));
1326
1327    let rhs = AffineRow {
1328        unknowns,
1329        claims: key.claims,
1330        constant: key.constant,
1331    };
1332
1333    [lhs, rhs]
1334}
1335
1336/// PAD oracle rows: the pad in [`OuterLayout::slot`] order,
1337/// `filler` (exactly `pad_rows · (code_len - message_len)`
1338/// entries) on the message tail.
1339pub fn build_pad_rows<F: TowerField + HardwareField>(
1340    layout: &OuterLayout,
1341    pad: &[Flat<F>],
1342    code_len: usize,
1343    filler: &[Flat<F>],
1344) -> errors::Result<Vec<Vec<Flat<F>>>> {
1345    let tail = code_len - layout.message_len;
1346
1347    if pad.len() > layout.pad_rows * layout.message_len || filler.len() != layout.pad_rows * tail {
1348        return Err(errors::Error::Protocol {
1349            protocol: "outer",
1350            message: "pad rows do not hold the pad and its filler",
1351        });
1352    }
1353
1354    let mut rows = vec![vec![Flat::from_raw(F::ZERO); layout.domain_len]; layout.pad_rows];
1355    for (i, value) in pad.iter().enumerate() {
1356        let (r, c) = layout.slot(Unknown::Pad(i as u32));
1357        rows[r][c] = *value;
1358    }
1359
1360    for (row, fill) in rows.iter_mut().zip(filler.chunks_exact(tail)) {
1361        row[layout.message_len..code_len].copy_from_slice(fill);
1362    }
1363
1364    Ok(rows)
1365}
1366
1367/// AUX oracle rows: the wire blocks and the five mask
1368/// rows, `filler` on every message tail, then the
1369/// zero-sum mask row's first `code_len - 1` cells and
1370/// the messages of the interleaved and both high masks.
1371pub fn build_aux_rows<F: TowerField + HardwareField>(
1372    layout: &OuterLayout,
1373    wires: &[[Flat<F>; 3]],
1374    code_len: usize,
1375    filler: &[Flat<F>],
1376) -> errors::Result<Vec<Vec<Flat<F>>>> {
1377    let count = layout.total_rows() - layout.pad_rows;
1378    let tail = code_len - layout.message_len;
1379
1380    if wires.len() > layout.aux_rows * layout.message_len
1381        || filler.len() != aux_filler_len(layout, code_len)
1382    {
1383        return Err(errors::Error::Protocol {
1384            protocol: "outer",
1385            message: "aux rows do not hold the wires and their filler",
1386        });
1387    }
1388
1389    let mut rows = vec![vec![Flat::from_raw(F::ZERO); layout.domain_len]; count];
1390
1391    for (mul, wire) in wires.iter().enumerate() {
1392        for role in [WireRole::Lhs, WireRole::Rhs, WireRole::Product] {
1393            let (r, c) = layout.slot(Unknown::Wire {
1394                mul: mul as u32,
1395                role,
1396            });
1397
1398            rows[r - layout.pad_rows][c] = wire[role as usize];
1399        }
1400    }
1401
1402    let (tails, rest) = filler.split_at(count * tail);
1403    for (row, fill) in rows.iter_mut().zip(tails.chunks_exact(tail)) {
1404        row[layout.message_len..code_len].copy_from_slice(fill);
1405    }
1406
1407    let (zero_sum_fill, rest) = rest.split_at(code_len - 1);
1408
1409    let zero_sum = layout.linear_mask() - layout.pad_rows;
1410
1411    let mut acc = Flat::from_raw(F::ZERO);
1412    for (slot, &v) in rows[zero_sum][..code_len - 1].iter_mut().zip(zero_sum_fill) {
1413        *slot = v;
1414        acc += v;
1415    }
1416
1417    rows[zero_sum][code_len - 1] = acc;
1418
1419    // The quadratic mask keeps its zero message; every
1420    // other mask row is uniform over the whole message.
1421    let uniform = [
1422        layout.interleaved_mask(),
1423        layout.linear_mask_hi(),
1424        layout.quadratic_mask_hi(),
1425    ];
1426
1427    for (row, fill) in uniform
1428        .into_iter()
1429        .zip(rest.chunks_exact(layout.message_len))
1430    {
1431        rows[row - layout.pad_rows][..layout.message_len].copy_from_slice(fill);
1432    }
1433
1434    Ok(rows)
1435}
1436
1437pub fn aux_filler_len(layout: &OuterLayout, code_len: usize) -> usize {
1438    let count = layout.total_rows() - layout.pad_rows;
1439
1440    count * (code_len - layout.message_len) + code_len - 1 + 3 * layout.message_len
1441}
1442
1443/// Challenge count for `outer_r_lin`. The tensor
1444/// batches the affine rows at `k / |F|`, not the
1445/// `1 / |F|` of one challenge per row.
1446pub fn linear_tensor_vars(rows: usize) -> usize {
1447    rows.next_power_of_two().ilog2() as usize
1448}
1449
1450/// Weights and target of the linear test on
1451/// the rows [`assemble`] builds, computed
1452/// from the records without building them.
1453pub fn linear_weights<F: TowerField + HardwareField + Into<Block128>>(
1454    layout: &OuterLayout,
1455    statement: &OuterStatement,
1456    records: &[TableRecord<'_, F>],
1457    eval: &EvalRecord<F>,
1458    tensor: &[F],
1459) -> errors::Result<LinearBatch<F>> {
1460    let rows = statement_rows(records, eval, statement)?;
1461
1462    if tensor.len() != linear_tensor_vars(rows) {
1463        return Err(errors::Error::Protocol {
1464            protocol: "outer",
1465            message: "linear batch takes ceil(log2(affine rows)) challenges",
1466        });
1467    }
1468
1469    weigh_records(layout, records, eval, &expand_batch_tensor(tensor, rows))
1470}
1471
1472/// Weights and target of the linear test on built rows:
1473/// the reference [`linear_weights`] is tested against.
1474pub fn linear_weights_of_rows<F: TowerField + HardwareField + Into<Block128>>(
1475    layout: &OuterLayout,
1476    rows: &OuterRows<'_, F>,
1477    tensor: &[F],
1478) -> errors::Result<LinearBatch<F>> {
1479    if tensor.len() != linear_tensor_vars(rows.affine.len()) {
1480        return Err(errors::Error::Protocol {
1481            protocol: "outer",
1482            message: "linear batch takes ceil(log2(affine rows)) challenges",
1483        });
1484    }
1485
1486    if rows.affine.iter().any(|row| !row.claims.is_empty()) {
1487        return Err(errors::Error::Protocol {
1488            protocol: "outer",
1489            message: "affine rows must be resolved before weighting",
1490        });
1491    }
1492
1493    Ok(weigh_rows(
1494        layout,
1495        rows,
1496        &expand_batch_tensor(tensor, rows.affine.len()),
1497    ))
1498}
1499
1500pub fn table_plan<F: TowerField>(
1501    air: &impl Air<F>,
1502    num_buses: usize,
1503    config: &Config,
1504) -> errors::Result<RingSwitchPlan> {
1505    let entries = air.virtual_expander().map(|e| e.expansion_entries());
1506
1507    RingSwitchPlan::new(
1508        air.column_layout(),
1509        entries.as_deref(),
1510        config.blind_units(),
1511        num_buses,
1512    )
1513}
1514
1515fn push_claim<F: TowerField>(row: &mut AffineRow<F>, pad_first: u32, idx: u32, coeff: Flat<F>) {
1516    row.unknowns.push((Unknown::Pad(pad_first + idx), coeff));
1517    row.claims.push((idx, coeff));
1518}
1519
1520fn scale_into<F: TowerField + HardwareField>(dst: &mut AffineRow<F>, src: &Form<F>, c: Flat<F>) {
1521    for &(u, k) in &src.unknowns {
1522        dst.unknowns.push((u, k * c));
1523    }
1524
1525    for &(i, k) in &src.claims {
1526        dst.claims.push((i, k * c));
1527    }
1528
1529    dst.constant += src.constant * c;
1530}
1531
1532fn expand_batch_tensor<F: TowerField + HardwareField>(tensor: &[F], len: usize) -> Vec<Flat<F>> {
1533    let mut scales = vec![Flat::from_raw(F::ZERO); len];
1534
1535    if len == 0 {
1536        return scales;
1537    }
1538
1539    scales[0] = Flat::from_raw(F::ONE);
1540
1541    for (j, &r) in tensor.iter().enumerate() {
1542        let half = 1usize << j;
1543        let scale = r.to_hardware();
1544
1545        for i in half..len.min(half << 1) {
1546            scales[i] = scales[i - half] * scale;
1547        }
1548    }
1549
1550    scales
1551}
1552
1553fn lagrange_weights<F: TowerField + HardwareField>(degree: usize, r: Flat<F>) -> Vec<Flat<F>> {
1554    (0..=degree)
1555        .map(|j| {
1556            let mut evals = vec![F::ZERO; degree + 1];
1557            evals[j] = F::ONE;
1558
1559            UnivariatePoly::new(evals).evaluate_hw(r)
1560        })
1561        .collect()
1562}
1563
1564fn weigh_records<F: TowerField + HardwareField + Into<Block128>>(
1565    layout: &OuterLayout,
1566    records: &[TableRecord<'_, F>],
1567    eval: &EvalRecord<F>,
1568    scales: &[Flat<F>],
1569) -> errors::Result<LinearBatch<F>> {
1570    let buses = bus_endpoints(records);
1571    let tables: usize = records.iter().map(TableRecord::affine_rows).sum();
1572
1573    if scales.len() != tables + eval_rows(eval) + buses.len() {
1574        return Err(errors::Error::Protocol {
1575            protocol: "outer",
1576            message: "scale count does not match the statement rows",
1577        });
1578    }
1579
1580    let mut acc = Accumulator::new(layout);
1581    let mut sweep = Sweep::new();
1582
1583    let mut first = 0;
1584    let mut mul_offset = 0u32;
1585
1586    for record in records {
1587        let end = first + record.affine_rows();
1588
1589        weigh_table(
1590            &mut acc,
1591            &mut sweep,
1592            record,
1593            &scales[first..end],
1594            mul_offset,
1595        )?;
1596
1597        first = end;
1598        mul_offset += record.shape.mul_wires() as u32;
1599    }
1600
1601    let final_row = eval_final_row(&eval.mask_form, eval.claim_masked, eval.fin);
1602    acc.row(&final_row, scales[first], &[], 0);
1603
1604    first += 1;
1605
1606    if let Some(gadget) = &eval.gadget {
1607        let chain = &scales[first..first + linearized::CHAIN_ROWS];
1608        for (row, &scale) in gadget.wire_rows().iter().zip(chain) {
1609            acc.row(row, scale, &[], 0);
1610        }
1611
1612        gadget.for_each_tie_weight(chain, |pad, weight| acc.add(Unknown::Pad(pad), weight));
1613
1614        first += linearized::CHAIN_ROWS;
1615    }
1616
1617    for (endpoints, &scale) in buses.values().zip(&scales[first..]) {
1618        acc.row(&bus_sum_row(endpoints), scale, &[], 0);
1619    }
1620
1621    Ok(acc.batch())
1622}
1623
1624fn weigh_table<F: TowerField + HardwareField + Into<Block128>>(
1625    acc: &mut Accumulator<'_, F>,
1626    sweep: &mut Sweep<F>,
1627    record: &TableRecord<'_, F>,
1628    scales: &[Flat<F>],
1629    mul_offset: u32,
1630) -> errors::Result<()> {
1631    let (predicate, tail) = scales.split_at(2 * record.shape.mul_nodes);
1632
1633    sweep.run(acc, record, predicate, tail[0], mul_offset)?;
1634
1635    let mut next = 0;
1636    record.for_each_tail_row(|row| {
1637        acc.row(&row, tail[next], &record.claims_masked, mul_offset);
1638
1639        next += 1;
1640    });
1641
1642    Ok(())
1643}
1644
1645fn weigh_rows<F: TowerField + HardwareField + Into<Block128>>(
1646    layout: &OuterLayout,
1647    rows: &OuterRows<'_, F>,
1648    scales: &[Flat<F>],
1649) -> LinearBatch<F> {
1650    let mut acc = Accumulator::new(layout);
1651    for (row, &scale) in rows.affine.iter().zip(scales) {
1652        acc.row(row, scale, &[], 0);
1653    }
1654
1655    for &(first, gadget) in &rows.gadgets {
1656        let chain = &scales[first..first + linearized::CHAIN_ROWS];
1657        gadget.for_each_tie_weight(chain, |pad, weight| acc.add(Unknown::Pad(pad), weight));
1658    }
1659
1660    acc.batch()
1661}
1662
1663fn bus_endpoints<'r, F: TowerField>(
1664    records: &'r [TableRecord<'_, F>],
1665) -> BTreeMap<&'r str, Vec<(u32, Flat<F>)>> {
1666    let mut buses: BTreeMap<&str, Vec<(u32, Flat<F>)>> = BTreeMap::new();
1667    for record in records {
1668        for (bus_id, pad, masked) in &record.claimed_sums {
1669            buses
1670                .entry(bus_id.as_str())
1671                .or_default()
1672                .push((*pad, *masked));
1673        }
1674    }
1675
1676    buses
1677}
1678
1679fn eval_rows<F>(eval: &EvalRecord<F>) -> usize {
1680    match eval.gadget {
1681        Some(_) => 1 + linearized::CHAIN_ROWS,
1682        None => 1,
1683    }
1684}
1685
1686fn eval_wires<F>(eval: &EvalRecord<F>) -> usize {
1687    match eval.gadget {
1688        Some(_) => linearized::BITS - 1,
1689        None => 0,
1690    }
1691}
1692
1693#[cfg(test)]
1694mod tests {
1695    use super::*;
1696    use crate::ProgramCell;
1697    use crate::constraint::ConstraintArena;
1698    use core::ops::Range;
1699    use hekate_math::Block128;
1700
1701    type F = Block128;
1702
1703    const BUS_PAD_FIRST: u32 = 10;
1704    const BUS_H_PAD: u32 = 20;
1705    const BUS_PUBLIC: u128 = 77;
1706
1707    /// `half - num_buses` at `bus_record`'s `half = 2` and one bus.
1708    const H_CLAIM: usize = 1;
1709
1710    fn shape(num_vars: usize, num_columns: usize, degree: usize, buses: usize) -> TableShape {
1711        TableShape {
1712            num_vars,
1713            num_columns,
1714            sumcheck_degree: degree,
1715            num_buses: buses,
1716            mul_nodes: 0,
1717            ring_units: false,
1718            whole_blinds: 0,
1719        }
1720    }
1721
1722    fn mix(seed: u128) -> Flat<F> {
1723        F::from(
1724            seed.wrapping_mul(0x9e37_79b9_7f4a_7c15)
1725                .wrapping_add(0x51ed_2701),
1726        )
1727        .to_hardware()
1728    }
1729
1730    fn bus_ast() -> ConstraintAst<F> {
1731        let mut arena = ConstraintArena::<F>::new();
1732
1733        let a = arena.cell(ProgramCell::current(0));
1734        let b = arena.cell(ProgramCell::current(1));
1735        let root = arena.mul(a, b);
1736
1737        ConstraintAst {
1738            arena,
1739            roots: vec![root],
1740            labels: vec![None],
1741        }
1742    }
1743
1744    /// One table over [`bus_ast`]: `c0 · c1` as the only
1745    /// constraint, two columns, one Permutation bus over
1746    /// `[Claim(0), Public(77)]` with selector `c1`.
1747    fn bus_record(
1748        ast: &ConstraintAst<F>,
1749        claims_masked: Vec<Flat<F>>,
1750        h_masked: Flat<F>,
1751        beta: Flat<F>,
1752    ) -> TableRecord<'_, F> {
1753        let mul_nodes = mul_node_count(ast);
1754
1755        let bus = BusRow {
1756            h_pad: BUS_H_PAD,
1757            h_masked,
1758            h_wire: mul_nodes as u32,
1759            sources: vec![BusSource::Claim(0), BusSource::Public(mix(BUS_PUBLIC))],
1760            selector: Some(1),
1761            recv_selector: None,
1762            eq_lookup: Flat::from_raw(F::ONE),
1763        };
1764
1765        TableRecord {
1766            claims_masked,
1767            ast,
1768            consistency: ConsistencyInputs {
1769                pad_first: BUS_PAD_FIRST,
1770                eq_zc: mix(1),
1771                alpha: mix(2),
1772                gamma: mix(3),
1773                beta,
1774                boundary: Vec::new(),
1775                telescope: Vec::new(),
1776                buses: vec![bus],
1777                val_final_form: Vec::new(),
1778                val_final_masked: mix(4),
1779            },
1780            fixed: Vec::new(),
1781            claimed_sums: vec![(String::from("bus"), 0, mix(7))],
1782            shape: TableShape {
1783                mul_nodes,
1784                ..shape(1, 2, 3, 1)
1785            },
1786        }
1787    }
1788
1789    fn eval_without_gadget() -> EvalRecord<F> {
1790        EvalRecord {
1791            mask_form: Vec::new(),
1792            claim_masked: mix(5),
1793            fin: mix(6),
1794            gadget: None,
1795        }
1796    }
1797
1798    fn mentions(rows: &[AffineRow<F>], mul: u32, role: WireRole) -> usize {
1799        rows.iter()
1800            .filter(|r| {
1801                r.unknowns
1802                    .iter()
1803                    .any(|&(u, _)| u == Unknown::Wire { mul, role })
1804            })
1805            .count()
1806    }
1807
1808    fn rich_ast() -> ConstraintAst<F> {
1809        let mut arena = ConstraintArena::<F>::new();
1810
1811        let c0 = arena.cell(ProgramCell::current(0));
1812        let c1 = arena.cell(ProgramCell::current(1));
1813        let c2 = arena.cell(ProgramCell::current(2));
1814        let n0 = arena.cell(ProgramCell::next(0));
1815
1816        let k = arena.constant(F::from(0x55u128));
1817
1818        let prod = arena.mul(c0, c1);
1819        let square = arena.mul(c2, c2);
1820
1821        arena.mul(c1, n0);
1822
1823        let scaled = arena.scale(F::from(0x1234u128), prod);
1824        let sum = arena.sum(vec![scaled, n0, k, n0, square]);
1825        let nested = arena.mul(sum, prod);
1826        let root = arena.add(nested, c2);
1827
1828        ConstraintAst {
1829            arena,
1830            roots: vec![root, prod, root, square],
1831            labels: vec![None; 4],
1832        }
1833    }
1834
1835    fn linear_ast() -> ConstraintAst<F> {
1836        let mut arena = ConstraintArena::<F>::new();
1837
1838        let c0 = arena.cell(ProgramCell::current(0));
1839        let n1 = arena.cell(ProgramCell::next(1));
1840        let k = arena.constant(F::from(7u128));
1841        let scaled = arena.scale(F::from(3u128), n1);
1842        let root = arena.sum(vec![c0, scaled, k]);
1843
1844        ConstraintAst {
1845            arena,
1846            roots: vec![root],
1847            labels: vec![None],
1848        }
1849    }
1850
1851    fn empty_ast() -> ConstraintAst<F> {
1852        ConstraintAst {
1853            arena: ConstraintArena::new(),
1854            roots: Vec::new(),
1855            labels: Vec::new(),
1856        }
1857    }
1858
1859    fn weighted_record(
1860        ast: &ConstraintAst<F>,
1861        base: u32,
1862        columns: u32,
1863        buses: u32,
1864        fixed: bool,
1865    ) -> TableRecord<'_, F> {
1866        let seed = base as u128;
1867        let one = Flat::from_raw(F::ONE);
1868        let half = columns + buses;
1869        let mul_nodes = mul_node_count(ast);
1870
1871        let free = base + 2 * half;
1872        let sums_pad = free + buses;
1873        let mask_pad = sums_pad + buses;
1874
1875        let bus_rows = (0..buses)
1876            .map(|k| BusRow {
1877                h_pad: free + k,
1878                h_masked: mix(seed + 10 + k as u128),
1879                h_wire: mul_nodes as u32 + k,
1880                sources: vec![
1881                    BusSource::Claim(0),
1882                    BusSource::Public(mix(seed + 20 + k as u128)),
1883                    BusSource::Claim(columns - 1),
1884                ],
1885                selector: (k % 2 == 0).then_some(1),
1886                recv_selector: (k == 1).then_some(0),
1887                eq_lookup: match k % 2 {
1888                    0 => one,
1889                    _ => mix(seed + 30 + k as u128),
1890                },
1891            })
1892            .collect();
1893
1894        TableRecord {
1895            claims_masked: (0..2 * half as u128).map(|i| mix(seed + 300 + i)).collect(),
1896            ast,
1897            consistency: ConsistencyInputs {
1898                pad_first: base,
1899                eq_zc: mix(seed + 1),
1900                alpha: mix(seed + 2),
1901                gamma: mix(seed + 3),
1902                beta: mix(seed + 4),
1903                boundary: vec![
1904                    (1, mix(seed + 5), mix(seed + 6)),
1905                    (0, mix(seed + 7), mix(seed + 8)),
1906                ],
1907                telescope: vec![(columns - 1, half + columns - 1)],
1908                buses: bus_rows,
1909                val_final_form: (0..3)
1910                    .map(|j| (Unknown::Pad(mask_pad + j), mix(seed + 60 + j as u128)))
1911                    .collect(),
1912                val_final_masked: mix(seed + 9),
1913            },
1914            fixed: match fixed {
1915                true => vec![(base, 0, mix(seed + 70)), (base + 1, 1, mix(seed + 71))],
1916                false => Vec::new(),
1917            },
1918            claimed_sums: (0..buses)
1919                .map(|k| {
1920                    let id = match k {
1921                        0 => String::from("shared"),
1922                        _ => alloc::format!("table {base} bus {k}"),
1923                    };
1924
1925                    (id, sums_pad + k, mix(seed + 90 + k as u128))
1926                })
1927                .collect(),
1928            shape: TableShape {
1929                num_vars: 3,
1930                num_columns: columns as usize,
1931                sumcheck_degree: 3,
1932                num_buses: buses as usize,
1933                mul_nodes,
1934                ring_units: false,
1935                whole_blinds: 0,
1936            },
1937        }
1938    }
1939
1940    fn weighted_eval(base: u32, first_wire: u32) -> EvalRecord<F> {
1941        let seed = base as u128;
1942
1943        let mut mask_form: Vec<(Unknown, Flat<F>)> = (0..3)
1944            .map(|j| (Unknown::Pad(base + j), mix(seed + 40 + j as u128)))
1945            .collect();
1946
1947        let entries = (0..4)
1948            .map(|c| (base + 6 + c, mix(seed + 50 + c as u128)))
1949            .collect();
1950
1951        let mu = (0..linearized::BITS as u128)
1952            .map(|j| mix(seed + 100 + j))
1953            .collect();
1954
1955        let gadget = RingGadget::new(entries, mu, first_wire);
1956        mask_form.extend(gadget.delta_form());
1957
1958        EvalRecord {
1959            mask_form,
1960            claim_masked: mix(seed + 80),
1961            fin: mix(seed + 81),
1962            gadget: Some(gadget),
1963        }
1964    }
1965
1966    /// Three tables over `asts`: the first carries every
1967    /// row class, the bus id `shared` spans two tables.
1968    fn weighted_statement(
1969        asts: &[ConstraintAst<F>; 3],
1970    ) -> (
1971        Vec<TableRecord<'_, F>>,
1972        EvalRecord<F>,
1973        OuterStatement,
1974        OuterLayout,
1975    ) {
1976        let records = vec![
1977            weighted_record(&asts[0], 0, 4, 2, true),
1978            weighted_record(&asts[1], 600, 3, 1, false),
1979            weighted_record(&asts[2], 1200, 2, 0, true),
1980        ];
1981
1982        let table_wires: usize = records.iter().map(|r| r.shape.mul_wires()).sum();
1983        let eval = weighted_eval(1800, table_wires as u32);
1984
1985        let statement = OuterStatement {
1986            masked_scalars: 2048,
1987            mul_wires: table_wires + linearized::BITS - 1,
1988        };
1989
1990        let geom = hekate_core::config::Config::prod()
1991            .outer_geom(statement.masked_scalars, statement.mul_wires, 128)
1992            .unwrap();
1993
1994        let layout =
1995            OuterLayout::new(&geom, statement.masked_scalars, statement.mul_wires).unwrap();
1996
1997        (records, eval, statement, layout)
1998    }
1999
2000    /// Statement rows of each nonempty row class, in [`assemble`] order.
2001    fn row_classes(
2002        records: &[TableRecord<'_, F>],
2003        eval: &EvalRecord<F>,
2004    ) -> Vec<(&'static str, Range<usize>)> {
2005        let mut classes = Vec::new();
2006        let mut at = 0;
2007
2008        for record in records {
2009            let buses = record.consistency.buses.len();
2010
2011            for (class, len) in [
2012                ("predicate", 2 * record.shape.mul_nodes),
2013                ("consistency", 1),
2014                ("bus operands", 2 * buses),
2015                ("fixed", record.fixed.len()),
2016                ("h pins", buses),
2017            ] {
2018                if len > 0 {
2019                    classes.push((class, at..at + len));
2020                }
2021
2022                at += len;
2023            }
2024        }
2025
2026        let chain = eval_rows(eval) - 1;
2027
2028        classes.push(("eval final", at..at + 1));
2029
2030        at += 1;
2031
2032        if chain > 0 {
2033            classes.push(("chain", at..at + chain));
2034
2035            at += chain;
2036        }
2037
2038        classes.push(("bus sums", at..at + bus_endpoints(records).len()));
2039
2040        classes
2041    }
2042
2043    #[test]
2044    fn ring_units_add_one_squaring_chain() {
2045        let plain = shape(12, 100, 3, 0);
2046        let ring = TableShape {
2047            ring_units: true,
2048            ..plain
2049        };
2050
2051        let without = OuterStatement::new(&[plain, plain], 12);
2052        let with = OuterStatement::new(&[ring, ring], 12);
2053
2054        assert_eq!(ring.mul_wires(), plain.mul_wires());
2055        assert_eq!(with.mul_wires - without.mul_wires, 127);
2056    }
2057
2058    #[test]
2059    fn ring_blind_planes_are_masked_only_under_blinding() {
2060        let plain = shape(12, 100, 3, 0);
2061        let ring = TableShape {
2062            ring_units: true,
2063            ..plain
2064        };
2065
2066        assert_eq!(ring.masked_scalars(), plain.masked_scalars());
2067
2068        for whole_blinds in [1, 4] {
2069            let blinded = |s: TableShape| TableShape { whole_blinds, ..s };
2070
2071            assert_eq!(
2072                blinded(ring).masked_scalars() - blinded(plain).masked_scalars(),
2073                2 * RING_BLIND_BITS
2074            );
2075        }
2076    }
2077
2078    #[test]
2079    fn empty_arena_has_no_wires() {
2080        let ast = ConstraintAst::<F> {
2081            arena: ConstraintArena::new(),
2082            roots: vec![],
2083            labels: vec![],
2084        };
2085
2086        assert_eq!(mul_node_count(&ast), 0);
2087    }
2088
2089    #[test]
2090    fn unreachable_mul_nodes_carry_no_wire() {
2091        let mut arena = ConstraintArena::<F>::new();
2092
2093        let a = arena.cell(ProgramCell::current(0));
2094        let b = arena.cell(ProgramCell::current(1));
2095        let live = arena.mul(a, b);
2096        let orphan = arena.cell(ProgramCell::current(2));
2097
2098        arena.mul(orphan, b);
2099
2100        let ast = ConstraintAst {
2101            arena,
2102            roots: vec![live],
2103            labels: vec![None],
2104        };
2105
2106        assert_eq!(mul_node_count(&ast), 1);
2107    }
2108
2109    #[test]
2110    fn shared_mul_node_counts_once() {
2111        let mut arena = ConstraintArena::<F>::new();
2112
2113        let a = arena.cell(ProgramCell::current(0));
2114        let b = arena.cell(ProgramCell::current(1));
2115        let shared = arena.mul(a, b);
2116        let left = arena.add(shared, a);
2117        let right = arena.add(shared, b);
2118
2119        let ast = ConstraintAst {
2120            arena,
2121            roots: vec![left, right],
2122            labels: vec![None, None],
2123        };
2124
2125        assert_eq!(mul_node_count(&ast), 1);
2126    }
2127
2128    #[test]
2129    fn nested_mul_nodes_each_carry_wire() {
2130        let mut arena = ConstraintArena::<F>::new();
2131
2132        let a = arena.cell(ProgramCell::current(0));
2133        let b = arena.cell(ProgramCell::current(1));
2134        let c = arena.cell(ProgramCell::current(2));
2135        let inner = arena.mul(a, b);
2136        let outer = arena.mul(inner, c);
2137
2138        let ast = ConstraintAst {
2139            arena,
2140            roots: vec![outer],
2141            labels: vec![None],
2142        };
2143
2144        assert_eq!(mul_node_count(&ast), 2);
2145    }
2146
2147    #[test]
2148    fn mul_under_sum_and_scale_is_reachable() {
2149        let mut arena = ConstraintArena::<F>::new();
2150
2151        let a = arena.cell(ProgramCell::current(0));
2152        let b = arena.cell(ProgramCell::current(1));
2153        let product = arena.mul(a, b);
2154        let scaled = arena.scale(F::ONE, product);
2155        let summed = arena.sum(vec![scaled, a]);
2156
2157        let ast = ConstraintAst {
2158            arena,
2159            roots: vec![summed],
2160            labels: vec![None],
2161        };
2162
2163        assert_eq!(mul_node_count(&ast), 1);
2164    }
2165
2166    #[test]
2167    fn busless_table_pays_for_no_h_claims() {
2168        let with_bus = TableShape {
2169            whole_blinds: 1,
2170            ..shape(12, 100, 3, 1)
2171        };
2172
2173        let without = TableShape {
2174            num_buses: 0,
2175            ..with_bus
2176        };
2177
2178        let claimed_sum_and_h_eval = 2;
2179        let h_claim_halves = 2;
2180
2181        assert_eq!(
2182            with_bus.masked_scalars() - without.masked_scalars(),
2183            claimed_sum_and_h_eval + h_claim_halves
2184        );
2185    }
2186
2187    #[test]
2188    fn masked_scalars_match_transcript_tally() {
2189        for whole_blinds in [1, 8] {
2190            let s = TableShape {
2191                whole_blinds,
2192                ..shape(12, 1977, 5, 3)
2193            };
2194
2195            let claimed_sums_and_h_evals = 2 * 3;
2196            let zerocheck_rounds = 12 * 5;
2197            let claims = 2 * (1977 + whole_blinds + 3);
2198
2199            assert_eq!(
2200                s.masked_scalars(),
2201                claimed_sums_and_h_evals + zerocheck_rounds + claims
2202            );
2203        }
2204    }
2205
2206    #[test]
2207    fn every_bus_adds_hadamard_row() {
2208        let s = TableShape {
2209            mul_nodes: 40,
2210            ..shape(12, 100, 3, 3)
2211        };
2212
2213        assert_eq!(s.mul_wires(), 43);
2214    }
2215
2216    #[test]
2217    fn statement_sums_over_every_table() {
2218        let main = TableShape {
2219            mul_nodes: 40,
2220            whole_blinds: 2,
2221            ..shape(12, 100, 3, 2)
2222        };
2223
2224        let chiplet = TableShape {
2225            mul_nodes: 7,
2226            whole_blinds: 8,
2227            ..shape(10, 30, 4, 1)
2228        };
2229
2230        let statement = OuterStatement::new(&[main, chiplet], 12);
2231
2232        assert_eq!(
2233            statement.masked_scalars,
2234            main.masked_scalars() + chiplet.masked_scalars() + 2 * 12
2235        );
2236        assert_eq!(statement.mul_wires, 50);
2237    }
2238
2239    #[test]
2240    fn every_hadamard_wire_operand_is_pinned_by_row() {
2241        let claims: Vec<Flat<F>> = (0..4).map(|i| mix(100 + i)).collect();
2242        let ast = bus_ast();
2243        let record = bus_record(&ast, claims, mix(8), mix(9));
2244        let statement = OuterStatement {
2245            masked_scalars: 32,
2246            mul_wires: record.shape.mul_wires(),
2247        };
2248
2249        let records = [record];
2250        let eval = eval_without_gadget();
2251        let rows = assemble(&records, &eval, &statement).unwrap();
2252
2253        for mul in 0..statement.mul_wires as u32 {
2254            for role in [WireRole::Lhs, WireRole::Rhs, WireRole::Product] {
2255                assert!(
2256                    mentions(&rows.affine, mul, role) >= 1,
2257                    "wire {mul} {role:?}"
2258                );
2259            }
2260        }
2261    }
2262
2263    #[test]
2264    fn bus_operand_rows_hold_on_honest_wires() {
2265        let plain: Vec<Flat<F>> = (0..4).map(|i| mix(200 + i)).collect();
2266        let pad: Vec<Flat<F>> = (0..32).map(|i| mix(300 + i)).collect();
2267        let claims_masked: Vec<Flat<F>> = plain
2268            .iter()
2269            .enumerate()
2270            .map(|(i, &c)| c + pad[BUS_PAD_FIRST as usize + i])
2271            .collect();
2272
2273        let beta = mix(9);
2274        let h = mix(400);
2275        let h_masked = h + pad[BUS_H_PAD as usize];
2276        let key = plain[0] + beta * mix(BUS_PUBLIC);
2277
2278        let ast = bus_ast();
2279        let record = bus_record(&ast, claims_masked.clone(), h_masked, beta);
2280        let bus = &record.consistency.buses[0];
2281        let [lhs, rhs] = bus_operand_rows(bus, BUS_PAD_FIRST, beta);
2282
2283        let wire = |mul: u32, role: WireRole| -> Flat<F> {
2284            assert_eq!(mul, bus.h_wire);
2285            match role {
2286                WireRole::Lhs => h,
2287                WireRole::Rhs => key,
2288                WireRole::Product => h * key,
2289            }
2290        };
2291
2292        for row in [&lhs, &rhs] {
2293            let mut left = Flat::from_raw(F::ZERO);
2294            for &(u, coeff) in &row.unknowns {
2295                let value = match u {
2296                    Unknown::Pad(i) => pad[i as usize],
2297                    Unknown::Wire { mul, role } => wire(mul, role),
2298                };
2299
2300                left += coeff * value;
2301            }
2302
2303            let mut right = row.constant;
2304            for &(idx, coeff) in &row.claims {
2305                right += coeff * claims_masked[idx as usize];
2306            }
2307
2308            assert_eq!(left, right);
2309        }
2310
2311        assert_eq!(mentions(&[lhs], bus.h_wire, WireRole::Lhs), 1);
2312        assert_eq!(mentions(&[rhs], bus.h_wire, WireRole::Rhs), 1);
2313    }
2314
2315    #[test]
2316    fn bus_operand_rows_reject_free_operand() {
2317        let plain: Vec<Flat<F>> = (0..4).map(|i| mix(200 + i)).collect();
2318        let pad: Vec<Flat<F>> = (0..32).map(|i| mix(300 + i)).collect();
2319        let claims_masked: Vec<Flat<F>> = plain
2320            .iter()
2321            .enumerate()
2322            .map(|(i, &c)| c + pad[BUS_PAD_FIRST as usize + i])
2323            .collect();
2324
2325        let beta = mix(9);
2326        let h = mix(400);
2327        let h_masked = h + pad[BUS_H_PAD as usize];
2328        let key = plain[0] + beta * mix(BUS_PUBLIC);
2329
2330        let ast = bus_ast();
2331        let record = bus_record(&ast, claims_masked.clone(), h_masked, beta);
2332        let bus = &record.consistency.buses[0];
2333        let [lhs, _] = bus_operand_rows(bus, BUS_PAD_FIRST, beta);
2334
2335        let one = Flat::from_raw(F::ONE);
2336        let forged = |role: WireRole| match role {
2337            WireRole::Lhs => one,
2338            WireRole::Rhs | WireRole::Product => h * key,
2339        };
2340
2341        assert_eq!(
2342            forged(WireRole::Lhs) * forged(WireRole::Rhs),
2343            forged(WireRole::Product)
2344        );
2345
2346        let mut left = Flat::from_raw(F::ZERO);
2347        for &(u, coeff) in &lhs.unknowns {
2348            let value = match u {
2349                Unknown::Pad(i) => pad[i as usize],
2350                Unknown::Wire { role, .. } => forged(role),
2351            };
2352
2353            left += coeff * value;
2354        }
2355
2356        let mut right = lhs.constant;
2357        for &(idx, coeff) in &lhs.claims {
2358            right += coeff * claims_masked[idx as usize];
2359        }
2360
2361        assert_ne!(left, right);
2362    }
2363
2364    #[test]
2365    fn layout_holds_whole_statement() {
2366        let cfg = hekate_core::config::Config::prod();
2367
2368        for (scalars, wires) in [(1usize, 0usize), (4_075, 2_384), (12_000, 20_000)] {
2369            let geom = cfg.outer_geom(scalars, wires, 128).unwrap();
2370            let layout = OuterLayout::new(&geom, scalars, wires).unwrap();
2371
2372            assert!(layout.pad_rows * layout.message_len >= scalars);
2373            assert!(layout.aux_rows * layout.message_len >= wires);
2374
2375            let (r, _) = layout.slot(Unknown::Pad(scalars as u32 - 1));
2376            assert!(r < layout.pad_rows);
2377
2378            if wires > 0 {
2379                let (r, _) = layout.slot(Unknown::Wire {
2380                    mul: wires as u32 - 1,
2381                    role: WireRole::Product,
2382                });
2383
2384                assert!(r < layout.interleaved_mask());
2385            }
2386
2387            assert_eq!(layout.total_rows(), layout.quadratic_mask_hi() + 1);
2388        }
2389    }
2390
2391    #[test]
2392    fn every_mask_row_interleaved_test_reveals_is_uniform() {
2393        let cfg = hekate_core::config::Config::prod();
2394        let (scalars, wires) = (4_075usize, 2_384usize);
2395
2396        let geom = cfg.outer_geom(scalars, wires, 128).unwrap();
2397        let layout = OuterLayout::new(&geom, scalars, wires).unwrap();
2398
2399        let triples: Vec<[Flat<F>; 3]> = (0..wires as u128)
2400            .map(|i| [mix(i), mix(i + (1 << 20)), mix(i + (1 << 40))])
2401            .collect();
2402
2403        let filler: Vec<Flat<F>> = (0..aux_filler_len(&layout, geom.code_len) as u128)
2404            .map(|i| mix(i + (1 << 60)))
2405            .collect();
2406
2407        let rows = build_aux_rows(&layout, &triples, geom.code_len, &filler).unwrap();
2408
2409        let zero = Flat::from_raw(F::ZERO);
2410
2411        for mask in [
2412            layout.interleaved_mask(),
2413            layout.linear_mask_hi(),
2414            layout.quadratic_mask_hi(),
2415        ] {
2416            assert!(
2417                rows[mask - layout.pad_rows][..layout.message_len]
2418                    .iter()
2419                    .all(|v| *v != zero)
2420            );
2421        }
2422
2423        let quadratic = layout.quadratic_mask() - layout.pad_rows;
2424        assert!(
2425            rows[quadratic][..layout.message_len]
2426                .iter()
2427                .all(|v| *v == zero)
2428        );
2429
2430        assert!(
2431            build_aux_rows(&layout, &triples, geom.code_len, &filler[1..]).is_err(),
2432            "filler length is exact"
2433        );
2434    }
2435
2436    #[test]
2437    fn mask_form_matches_round_by_round_recursion() {
2438        let degree = 3;
2439        let initial: Vec<(Unknown, Flat<F>)> = (0..4u32)
2440            .map(|i| (Unknown::Pad(i), mix(500 + i as u128)))
2441            .collect();
2442        let round_pads: Vec<Vec<u32>> = (0..5u32)
2443            .map(|i| {
2444                (0..degree as u32)
2445                    .map(|j| 100 + i * degree as u32 + j)
2446                    .collect()
2447            })
2448            .collect();
2449        let round_refs: Vec<&[u32]> = round_pads.iter().map(|p| p.as_slice()).collect();
2450        let challenges: Vec<Flat<F>> = (0..5).map(|i| mix(600 + i)).collect();
2451
2452        let mut expected: Vec<(Unknown, Flat<F>)> = initial.clone();
2453        for (pads, &r) in round_refs.iter().zip(&challenges) {
2454            let weights = lagrange_weights::<F>(degree, r);
2455            for (_, coeff) in expected.iter_mut() {
2456                *coeff *= weights[0];
2457            }
2458
2459            expected.push((Unknown::Pad(pads[0]), weights[0]));
2460
2461            for (j, &pad) in pads.iter().enumerate() {
2462                expected.push((Unknown::Pad(pad), weights[j + 1]));
2463            }
2464        }
2465
2466        assert_eq!(
2467            sumcheck_mask_form(&initial, &round_refs, &challenges, degree),
2468            expected
2469        );
2470    }
2471
2472    #[test]
2473    fn h_claim_pin_row_ties_h_eval_to_committed_claim() {
2474        let pad: Vec<Flat<F>> = (0..32).map(|i| mix(300 + i)).collect();
2475        let one = Flat::from_raw(F::ONE);
2476        let h = mix(400);
2477        let h_masked = h + pad[BUS_H_PAD as usize];
2478
2479        let pin_form = [
2480            (Unknown::Pad(BUS_H_PAD), one),
2481            (Unknown::Pad(BUS_PAD_FIRST + H_CLAIM as u32), one),
2482        ];
2483
2484        for (h_claim, holds) in [(h, true), (h + one, false)] {
2485            let mut plain: Vec<Flat<F>> = (0..4).map(|i| mix(200 + i)).collect();
2486            plain[H_CLAIM] = h_claim;
2487
2488            let claims_masked: Vec<Flat<F>> = plain
2489                .iter()
2490                .enumerate()
2491                .map(|(i, &c)| c + pad[BUS_PAD_FIRST as usize + i])
2492                .collect();
2493
2494            let ast = bus_ast();
2495            let record = bus_record(&ast, claims_masked.clone(), h_masked, mix(9));
2496            let statement = OuterStatement {
2497                masked_scalars: 32,
2498                mul_wires: record.shape.mul_wires(),
2499            };
2500
2501            let records = [record];
2502            let eval = eval_without_gadget();
2503            let rows = assemble(&records, &eval, &statement).unwrap();
2504            let pins: Vec<&AffineRow<F>> = rows
2505                .affine
2506                .iter()
2507                .filter(|r| r.unknowns == pin_form)
2508                .collect();
2509
2510            assert_eq!(pins.len(), 1, "one pin row per bus");
2511            assert_eq!(pins[0].constant, h_masked + claims_masked[H_CLAIM]);
2512
2513            let left = pad[BUS_H_PAD as usize] + pad[BUS_PAD_FIRST as usize + H_CLAIM];
2514
2515            assert_eq!(left == pins[0].constant, holds);
2516        }
2517    }
2518
2519    #[test]
2520    fn batch_tensor_is_product_over_set_bits() {
2521        let a = F::from(0x1234_5678_9abc_def0u128);
2522        let b = F::from(0x0fed_cba9_8765_4321u128);
2523        let c = F::from(0xdead_beef_cafe_babeu128);
2524
2525        let (fa, fb, fc) = (a.to_hardware(), b.to_hardware(), c.to_hardware());
2526        let one = Flat::from_raw(F::ONE);
2527
2528        assert_eq!(expand_batch_tensor::<F>(&[], 1), vec![one]);
2529        assert_eq!(expand_batch_tensor(&[a], 2), vec![one, fa]);
2530        assert_eq!(expand_batch_tensor(&[a, b], 4), vec![one, fa, fb, fa * fb]);
2531
2532        assert_eq!(
2533            expand_batch_tensor(&[a, b, c], 8),
2534            vec![one, fa, fb, fa * fb, fc, fa * fc, fb * fc, fa * fb * fc,],
2535        );
2536    }
2537
2538    #[test]
2539    fn batch_tensor_truncates_below_power_of_two() {
2540        let a = F::from(0x1234_5678_9abc_def0u128);
2541        let b = F::from(0x0fed_cba9_8765_4321u128);
2542
2543        let full = expand_batch_tensor(&[a, b], 4);
2544
2545        for len in 1..=4 {
2546            assert_eq!(expand_batch_tensor(&[a, b], len), full[..len]);
2547        }
2548    }
2549
2550    #[test]
2551    fn reverse_pass_reproduces_assembled_weights() {
2552        let asts = [rich_ast(), linear_ast(), empty_ast()];
2553        let (records, eval, statement, layout) = weighted_statement(&asts);
2554
2555        let rows = statement_rows(&records, &eval, &statement).unwrap();
2556        let assembled = assemble(&records, &eval, &statement).unwrap();
2557
2558        assert_eq!(assembled.affine.len(), rows);
2559
2560        let tensor: Vec<F> = (0..linear_tensor_vars(rows) as u128)
2561            .map(|i| mix(9000 + i).to_tower())
2562            .collect();
2563
2564        let got = linear_weights(&layout, &statement, &records, &eval, &tensor).unwrap();
2565        let want = linear_weights_of_rows(&layout, &assembled, &tensor).unwrap();
2566
2567        assert_eq!(got.rows, want.rows);
2568        assert_eq!(got.weights, want.weights);
2569        assert_eq!(got.target, want.target);
2570    }
2571
2572    #[test]
2573    fn dropping_row_class_breaks_equality() {
2574        let asts = [rich_ast(), linear_ast(), empty_ast()];
2575        let (records, eval, statement, layout) = weighted_statement(&asts);
2576
2577        let assembled = assemble(&records, &eval, &statement).unwrap();
2578        let scales: Vec<Flat<F>> = (0..assembled.affine.len() as u128)
2579            .map(|i| mix(7000 + i))
2580            .collect();
2581
2582        let want = weigh_rows(&layout, &assembled, &scales);
2583        let full = weigh_records(&layout, &records, &eval, &scales).unwrap();
2584
2585        assert_eq!(full.weights, want.weights);
2586        assert_eq!(full.target, want.target);
2587
2588        for (class, rows) in row_classes(&records, &eval) {
2589            let mut cut = scales.clone();
2590            cut[rows].fill(Flat::from_raw(F::ZERO));
2591
2592            let got = weigh_records(&layout, &records, &eval, &cut).unwrap();
2593
2594            assert!(
2595                got.weights != want.weights || got.target != want.target,
2596                "{class}"
2597            );
2598        }
2599    }
2600
2601    #[test]
2602    fn miscounted_mul_nodes_are_rejected() {
2603        let ast = rich_ast();
2604
2605        for (delta, message) in [
2606            (-1, "table shape undercounts the AST's Mul nodes"),
2607            (1, "table shape overcounts the AST's Mul nodes"),
2608        ] {
2609            let mut record = weighted_record(&ast, 0, 4, 2, true);
2610            record.shape.mul_nodes = record.shape.mul_nodes.saturating_add_signed(delta);
2611
2612            let statement = OuterStatement {
2613                masked_scalars: 2048,
2614                mul_wires: record.shape.mul_wires(),
2615            };
2616
2617            let geom = Config::prod()
2618                .outer_geom(statement.masked_scalars, statement.mul_wires, 128)
2619                .unwrap();
2620
2621            let layout =
2622                OuterLayout::new(&geom, statement.masked_scalars, statement.mul_wires).unwrap();
2623
2624            let records = [record];
2625            let eval = eval_without_gadget();
2626            let rows = statement_rows(&records, &eval, &statement).unwrap();
2627            let tensor: Vec<F> = (0..linear_tensor_vars(rows) as u128)
2628                .map(|i| mix(8000 + i).to_tower())
2629                .collect();
2630
2631            match linear_weights(&layout, &statement, &records, &eval, &tensor) {
2632                Err(errors::Error::Protocol { message: got, .. }) => assert_eq!(got, message),
2633                _ => panic!("{message}"),
2634            }
2635        }
2636    }
2637
2638    #[test]
2639    fn mismatched_scale_count_is_rejected() {
2640        let asts = [rich_ast(), linear_ast(), empty_ast()];
2641        let (records, eval, statement, layout) = weighted_statement(&asts);
2642
2643        let rows = statement_rows(&records, &eval, &statement).unwrap();
2644        let scales: Vec<Flat<F>> = (0..=rows as u128).map(|i| mix(6000 + i)).collect();
2645
2646        assert!(weigh_records(&layout, &records, &eval, &scales[..rows]).is_ok());
2647
2648        for len in [0, rows - 1, rows + 1] {
2649            match weigh_records(&layout, &records, &eval, &scales[..len]) {
2650                Err(errors::Error::Protocol { message, .. }) => {
2651                    assert_eq!(
2652                        message, "scale count does not match the statement rows",
2653                        "{len}"
2654                    )
2655                }
2656                _ => panic!("{len}"),
2657            }
2658        }
2659    }
2660}