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