use alloc::vec::Vec;
use air::{Air, AuxRandElements, ConstraintCompositionCoefficients, EvaluationFrame};
use math::{polynom, FieldElement};
pub fn evaluate_constraints<A: Air, E: FieldElement<BaseField = A::BaseField>>(
air: &A,
composition_coefficients: ConstraintCompositionCoefficients<E>,
main_trace_frame: &EvaluationFrame<E>,
aux_trace_frame: &Option<EvaluationFrame<E>>,
aux_rand_elements: Option<&AuxRandElements<E>>,
x: E,
) -> E {
let t_constraints = air.get_transition_constraints(&composition_coefficients.transition);
let periodic_values = air
.get_periodic_column_polys()
.iter()
.map(|poly| {
let num_cycles = air.trace_length() / poly.len();
let x = x.exp_vartime((num_cycles as u32).into());
polynom::eval(poly, x)
})
.collect::<Vec<_>>();
let mut t_evaluations1 = vec![E::ZERO; t_constraints.num_main_constraints()];
air.evaluate_transition(main_trace_frame, &periodic_values, &mut t_evaluations1);
let mut t_evaluations2 = vec![E::ZERO; t_constraints.num_aux_constraints()];
if let Some(aux_trace_frame) = aux_trace_frame {
let aux_rand_elements =
aux_rand_elements.expect("expected aux rand elements to be present");
air.evaluate_aux_transition(
main_trace_frame,
aux_trace_frame,
&periodic_values,
aux_rand_elements,
&mut t_evaluations2,
);
}
let mut result = t_constraints.combine_evaluations::<E>(&t_evaluations1, &t_evaluations2, x);
let b_constraints =
air.get_boundary_constraints(aux_rand_elements, &composition_coefficients.boundary);
for group in b_constraints.main_constraints().iter() {
result += group.evaluate_at(main_trace_frame.current(), x);
}
if let Some(aux_trace_frame) = aux_trace_frame {
for group in b_constraints.aux_constraints().iter() {
result += group.evaluate_at(aux_trace_frame.current(), x);
}
}
result
}