Skip to main content

miden_ace_codegen/
testing.rs

1//! Test helpers for ACE codegen, available under the `testing` feature or `#[cfg(test)]`.
2//!
3//! These provide reference evaluators for symbolic expressions, periodic columns,
4//! constraint folding, quotient recomposition, and DAG evaluation, suitable for
5//! validating the ACE pipeline from both within ace-codegen tests and from
6//! downstream integration tests (e.g. in miden-air).
7
8use miden_core::{Felt, field::QuadFelt};
9use miden_crypto::{
10    field::{ExtensionField, Field, TwoAdicField},
11    stark::{
12        air::symbolic::{
13            BaseEntry, BaseLeaf, ConstraintLayout, ExtEntry, ExtLeaf, SymbolicExpression,
14            SymbolicExpressionExt,
15        },
16        dft::{Radix2DitParallel, TwoAdicSubgroupDft},
17    },
18};
19
20use crate::{AceDag, AceError, InputKey, InputLayout};
21
22/// Deterministic input filler for layout-sized buffers.
23///
24/// Generates pseudo-random `QuadFelt` values using a simple LCG, suitable for
25/// testing against hand-computed reference values.
26pub fn fill_inputs(layout: &InputLayout) -> Vec<QuadFelt> {
27    let mut values = Vec::with_capacity(layout.total_inputs);
28    let mut state = 0x9e37_79b9_7f4a_7c15u64;
29    for _ in 0..layout.total_inputs {
30        state = state.wrapping_mul(6364136223846793005).wrapping_add(1);
31        let lo = Felt::new(state);
32        state = state.wrapping_mul(6364136223846793005).wrapping_add(1);
33        let hi = Felt::new(state);
34        values.push(QuadFelt::new([lo, hi]));
35    }
36    values
37}
38
39/// Evaluate periodic columns at a point by computing the polynomial in
40/// coefficient form via inverse DFT, then evaluating with Horner's method.
41pub fn eval_periodic_values<F, EF>(periodic_columns: &[Vec<F>], z_k: EF) -> Vec<EF>
42where
43    F: TwoAdicField + Ord,
44    EF: ExtensionField<F>,
45{
46    if periodic_columns.is_empty() {
47        return Vec::new();
48    }
49    let max_len = periodic_columns.iter().map(|col| col.len()).max().unwrap_or(0);
50    let dft = Radix2DitParallel::<F>::default();
51
52    periodic_columns
53        .iter()
54        .map(|col| {
55            if col.is_empty() {
56                return EF::ZERO;
57            }
58            let coeffs = dft.idft(col.clone());
59            let ratio = max_len / col.len();
60            let log_pow = ratio.ilog2() as usize;
61            let mut z_col = z_k;
62            for _ in 0..log_pow {
63                z_col *= z_col;
64            }
65            let mut acc = EF::ZERO;
66            for coeff in coeffs.iter().rev() {
67                acc = acc * z_col + EF::from(*coeff);
68            }
69            acc
70        })
71        .collect()
72}
73
74/// Evaluate a base-field symbolic expression at concrete inputs.
75pub fn eval_base_expr<F, EF>(
76    expr: &SymbolicExpression<F>,
77    inputs: &[EF],
78    layout: &InputLayout,
79    periodic_values: &[EF],
80) -> EF
81where
82    F: Field,
83    EF: ExtensionField<F>,
84{
85    match expr {
86        SymbolicExpression::Leaf(leaf) => match leaf {
87            BaseLeaf::Variable(v) => match v.entry {
88                BaseEntry::Main { offset } => {
89                    let key = InputKey::Main { offset, index: v.index };
90                    inputs[layout.index(key).unwrap()]
91                },
92                BaseEntry::Public => {
93                    let key = InputKey::Public(v.index);
94                    inputs[layout.index(key).unwrap()]
95                },
96                BaseEntry::Periodic => periodic_values[v.index],
97                BaseEntry::Preprocessed { .. } => panic!("preprocessed not supported in test"),
98            },
99            BaseLeaf::IsFirstRow => inputs[layout.index(InputKey::IsFirst).unwrap()],
100            BaseLeaf::IsLastRow => inputs[layout.index(InputKey::IsLast).unwrap()],
101            BaseLeaf::IsTransition => inputs[layout.index(InputKey::IsTransition).unwrap()],
102            BaseLeaf::Constant(c) => EF::from(*c),
103        },
104        SymbolicExpression::Add { x, y, .. } => {
105            eval_base_expr::<F, EF>(x, inputs, layout, periodic_values)
106                + eval_base_expr::<F, EF>(y, inputs, layout, periodic_values)
107        },
108        SymbolicExpression::Sub { x, y, .. } => {
109            eval_base_expr::<F, EF>(x, inputs, layout, periodic_values)
110                - eval_base_expr::<F, EF>(y, inputs, layout, periodic_values)
111        },
112        SymbolicExpression::Mul { x, y, .. } => {
113            eval_base_expr::<F, EF>(x, inputs, layout, periodic_values)
114                * eval_base_expr::<F, EF>(y, inputs, layout, periodic_values)
115        },
116        SymbolicExpression::Neg { x, .. } => {
117            -eval_base_expr::<F, EF>(x, inputs, layout, periodic_values)
118        },
119    }
120}
121
122/// Evaluate an extension-field symbolic expression at concrete inputs.
123pub fn eval_ext_expr<F, EF>(
124    expr: &SymbolicExpressionExt<F, EF>,
125    inputs: &[EF],
126    layout: &InputLayout,
127    periodic_values: &[EF],
128) -> EF
129where
130    F: Field,
131    EF: ExtensionField<F>,
132{
133    match expr {
134        SymbolicExpressionExt::Leaf(leaf) => match leaf {
135            ExtLeaf::Base(base_expr) => {
136                eval_base_expr::<F, EF>(base_expr, inputs, layout, periodic_values)
137            },
138            ExtLeaf::ExtVariable(v) => match v.entry {
139                ExtEntry::Permutation { offset } => {
140                    let mut acc = EF::ZERO;
141                    for coord in 0..EF::DIMENSION {
142                        let basis = EF::ith_basis_element(coord).unwrap();
143                        let key = InputKey::AuxCoord { offset, index: v.index, coord };
144                        let value = inputs[layout.index(key).unwrap()];
145                        acc += basis * value;
146                    }
147                    acc
148                },
149                ExtEntry::Challenge => {
150                    let alpha = inputs[layout.index(InputKey::AuxRandAlpha).unwrap()];
151                    let beta = inputs[layout.index(InputKey::AuxRandBeta).unwrap()];
152                    match v.index {
153                        0 => alpha,
154                        1 => beta,
155                        _ => panic!(
156                            "challenge index {} exceeds the 2-element randomness convention",
157                            v.index
158                        ),
159                    }
160                },
161                ExtEntry::PermutationValue => {
162                    let key = InputKey::AuxBusBoundary(v.index);
163                    inputs[layout.index(key).unwrap()]
164                },
165            },
166            ExtLeaf::ExtConstant(c) => *c,
167        },
168        SymbolicExpressionExt::Add { x, y, .. } => {
169            eval_ext_expr::<F, EF>(x, inputs, layout, periodic_values)
170                + eval_ext_expr::<F, EF>(y, inputs, layout, periodic_values)
171        },
172        SymbolicExpressionExt::Sub { x, y, .. } => {
173            eval_ext_expr::<F, EF>(x, inputs, layout, periodic_values)
174                - eval_ext_expr::<F, EF>(y, inputs, layout, periodic_values)
175        },
176        SymbolicExpressionExt::Mul { x, y, .. } => {
177            eval_ext_expr::<F, EF>(x, inputs, layout, periodic_values)
178                * eval_ext_expr::<F, EF>(y, inputs, layout, periodic_values)
179        },
180        SymbolicExpressionExt::Neg { x, .. } => {
181            -eval_ext_expr::<F, EF>(x, inputs, layout, periodic_values)
182        },
183    }
184}
185
186/// Evaluate the folded constraint accumulator from symbolic constraints.
187///
188/// Merges base and extension constraints in evaluation order (using the
189/// `ConstraintLayout`), then folds them via Horner with `alpha`.
190pub fn eval_folded_constraints<F, EF>(
191    base_constraints: &[SymbolicExpression<F>],
192    ext_constraints: &[SymbolicExpressionExt<F, EF>],
193    constraint_layout: &ConstraintLayout,
194    inputs: &[EF],
195    layout: &InputLayout,
196    periodic_values: &[EF],
197) -> EF
198where
199    F: Field,
200    EF: ExtensionField<F>,
201{
202    let alpha = inputs[layout.index(InputKey::Alpha).unwrap()];
203
204    let total = constraint_layout.base_indices.len() + constraint_layout.ext_indices.len();
205    let mut ordered: Vec<(usize, bool, usize)> = Vec::with_capacity(total);
206    for (i, &pos) in constraint_layout.base_indices.iter().enumerate() {
207        ordered.push((pos, false, i));
208    }
209    for (i, &pos) in constraint_layout.ext_indices.iter().enumerate() {
210        ordered.push((pos, true, i));
211    }
212    ordered.sort_by_key(|(pos, ..)| *pos);
213
214    let mut acc = EF::ZERO;
215    for &(_, is_ext, idx) in &ordered {
216        let val = if is_ext {
217            eval_ext_expr::<F, EF>(&ext_constraints[idx], inputs, layout, periodic_values)
218        } else {
219            eval_base_expr::<F, EF>(&base_constraints[idx], inputs, layout, periodic_values)
220        };
221        acc = acc * alpha + val;
222    }
223    acc
224}
225
226/// Evaluate the quotient recomposition at `zeta` using provided inputs.
227pub fn eval_quotient<F, EF>(layout: &InputLayout, inputs: &[EF]) -> EF
228where
229    F: Field,
230    EF: ExtensionField<F>,
231{
232    let k = layout.counts.num_quotient_chunks;
233    let z_pow_n = inputs[layout.index(InputKey::ZPowN).expect("ZPowN in layout")];
234    let s0 = inputs[layout.index(InputKey::S0).expect("S0 in layout")];
235    let f = inputs[layout.index(InputKey::F).expect("F in layout")];
236    let weight0 = inputs[layout.index(InputKey::Weight0).expect("Weight0 in layout")];
237
238    let (deltas, weights) = compute_deltas_and_weights(k, z_pow_n, s0, f, weight0);
239
240    let mut quotient = EF::ZERO;
241    for chunk in 0..k {
242        let mut chunk_value = EF::ZERO;
243        for coord in 0..EF::DIMENSION {
244            let basis = EF::ith_basis_element(coord).expect("basis index within extension degree");
245            let coord_value = inputs[layout
246                .index(InputKey::QuotientChunkCoord { offset: 0, chunk, coord })
247                .expect("quotient chunk coord in layout")];
248            chunk_value += basis * coord_value;
249        }
250
251        let mut prod = EF::ONE;
252        for (j, delta) in deltas.iter().enumerate() {
253            if j != chunk {
254                prod *= *delta;
255            }
256        }
257        quotient += weights[chunk] * prod * chunk_value;
258    }
259
260    quotient
261}
262
263/// Compute the barycentric kernel value (`zps`) for a single quotient chunk.
264pub fn zps_for_chunk<F, EF>(layout: &InputLayout, inputs: &[EF], chunk: usize) -> EF
265where
266    F: Field,
267    EF: ExtensionField<F>,
268{
269    let k = layout.counts.num_quotient_chunks;
270    assert!(chunk < k, "quotient chunk {chunk} out of range (k={k})");
271
272    let z_pow_n = inputs[layout.index(InputKey::ZPowN).expect("ZPowN in layout")];
273    let s0 = inputs[layout.index(InputKey::S0).expect("S0 in layout")];
274    let f = inputs[layout.index(InputKey::F).expect("F in layout")];
275    let weight0 = inputs[layout.index(InputKey::Weight0).expect("Weight0 in layout")];
276
277    let (deltas, weights) = compute_deltas_and_weights(k, z_pow_n, s0, f, weight0);
278
279    let mut prod = EF::ONE;
280    for (j, delta) in deltas.iter().enumerate() {
281        if j != chunk {
282            prod *= *delta;
283        }
284    }
285
286    weights[chunk] * prod
287}
288
289/// Evaluate a lowered DAG against concrete inputs.
290pub fn eval_dag<EF>(dag: &AceDag<EF>, inputs: &[EF], layout: &InputLayout) -> Result<EF, AceError>
291where
292    EF: Field,
293{
294    if inputs.len() != layout.total_inputs {
295        return Err(AceError::InvalidInputLength {
296            expected: layout.total_inputs,
297            got: inputs.len(),
298        });
299    }
300
301    let mut values: Vec<EF> = vec![EF::ZERO; dag.nodes().len()];
302    for (idx, node) in dag.nodes().iter().enumerate() {
303        let value = match node {
304            crate::dag::NodeKind::Input(key) => {
305                let input_idx = layout.index(*key).ok_or_else(|| AceError::InvalidInputLayout {
306                    message: format!("missing input key in layout: {key:?}"),
307                })?;
308                inputs[input_idx]
309            },
310            crate::dag::NodeKind::Constant(c) => *c,
311            crate::dag::NodeKind::Add(a, b) => values[a.index()] + values[b.index()],
312            crate::dag::NodeKind::Sub(a, b) => values[a.index()] - values[b.index()],
313            crate::dag::NodeKind::Mul(a, b) => values[a.index()] * values[b.index()],
314            crate::dag::NodeKind::Neg(a) => -values[a.index()],
315        };
316        values[idx] = value;
317    }
318
319    Ok(values[dag.root().index()])
320}
321
322fn compute_deltas_and_weights<EF: Field>(
323    k: usize,
324    z_pow_n: EF,
325    s0: EF,
326    f: EF,
327    weight0: EF,
328) -> (Vec<EF>, Vec<EF>) {
329    let mut deltas = Vec::with_capacity(k);
330    let mut weights = Vec::with_capacity(k);
331    let mut shift = s0;
332    let mut weight = weight0;
333    for _ in 0..k {
334        deltas.push(z_pow_n - shift);
335        weights.push(weight);
336        shift *= f;
337        weight *= f;
338    }
339    (deltas, weights)
340}