1use 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
22pub 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
39pub 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
74pub 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
122pub 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
186pub 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
226pub 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
263pub 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
289pub 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}