Add a simpler ExpressionEvaluator

backend/src/mock/connection_constraint_checker.rs

@@ -5,6 +5,7 @@ use std::ops::ControlFlow;
 
 use itertools::Itertools;
 use powdr_ast::analyzed::AlgebraicExpression;
+use powdr_ast::analyzed::AlgebraicReference;
 use powdr_ast::analyzed::Analyzed;
 use powdr_ast::analyzed::{
     Identity, LookupIdentity, PermutationIdentity, PhantomLookupIdentity,
@@ -13,15 +14,12 @@ use powdr_ast::analyzed::{
 use powdr_ast::parsed::visitor::ExpressionVisitable;
 use powdr_ast::parsed::visitor::VisitOrder;
 use powdr_backend_utils::referenced_namespaces_algebraic_expression;
-use powdr_executor::witgen::ExpressionEvaluator;
+use powdr_executor::witgen::evaluators::expression_evaluator::ExpressionEvaluator;
+use powdr_executor::witgen::evaluators::expression_evaluator::TraceValues;
 use powdr_number::FieldElement;
 use rayon::iter::IntoParallelIterator;
 use rayon::iter::ParallelIterator;
 
-use crate::mock::evaluator::evaluate_to_fe;
-
-use super::evaluator::EmptyVariables;
-use super::evaluator::Variables;
 use super::machine::Machine;
 
 #[derive(PartialEq, Eq, Debug)]
@@ -256,21 +254,19 @@ impl<'a, F: FieldElement> ConnectionConstraintChecker<'a, F> {
                 Some(machine) => (0..machine.size)
                     .into_par_iter()
                     .filter_map(|row| {
-                        let variables = Variables {
-                            machine,
-                            row,
-                            challenges: self.challenges,
-                        };
-                        let mut evaluator =
-                            ExpressionEvaluator::new(&variables, &machine.intermediate_definitions);
-                        let result = evaluate_to_fe(&mut evaluator, &selected_expressions.selector);
+                        let mut evaluator = ExpressionEvaluator::new(
+                            machine.trace_values.row(row),
+                            &machine.intermediate_definitions,
+                            self.challenges,
+                        );
+                        let result = evaluator.evaluate(&selected_expressions.selector);
 
                         assert!(result.is_zero() || result.is_one(), "Non-binary selector");
                         result.is_one().then(|| {
                             let values = selected_expressions
                                 .expressions
                                 .iter()
-                                .map(|expression| evaluate_to_fe(&mut evaluator, expression))
+                                .map(|expression| evaluator.evaluate(expression))
                                 .collect::<Vec<_>>();
                             Tuple { values, row }
                         })
@@ -283,8 +279,13 @@ impl<'a, F: FieldElement> ConnectionConstraintChecker<'a, F> {
             None => {
                 let empty_variables = EmptyVariables {};
                 let empty_definitions = BTreeMap::new();
-                let mut evaluator = ExpressionEvaluator::new(empty_variables, &empty_definitions);
-                let selector_value = evaluate_to_fe(&mut evaluator, &selected_expressions.selector);
+                let empty_challenges = BTreeMap::new();
+                let mut evaluator = ExpressionEvaluator::new(
+                    empty_variables,
+                    &empty_definitions,
+                    &empty_challenges,
+                );
+                let selector_value = evaluator.evaluate(&selected_expressions.selector);
 
                 match selector_value.to_degree() {
                     // Selected expressions is of the form `0 $ [ <constants> ]`
@@ -307,7 +308,7 @@ impl<'a, F: FieldElement> ConnectionConstraintChecker<'a, F> {
                         let values = selected_expressions
                             .expressions
                             .iter()
-                            .map(|expression| evaluate_to_fe(&mut evaluator, expression))
+                            .map(|expression| evaluator.evaluate(expression))
                             .collect::<Vec<_>>();
                         vec![Tuple { values, row: 0 }]
                     }
@@ -318,6 +319,17 @@ impl<'a, F: FieldElement> ConnectionConstraintChecker<'a, F> {
     }
 }
 
+struct EmptyVariables;
+
+impl<T> TraceValues<T> for EmptyVariables
+where
+    T: FieldElement,
+{
+    fn get(&self, _reference: &AlgebraicReference) -> T {
+        panic!()
+    }
+}
+
 /// Converts a slice to a multi-set, represented as a map from elements to their count.
 fn to_multi_set<T: Ord>(a: &[T]) -> BTreeMap<&T, usize> {
     a.iter()

backend/src/mock/machine.rs

@@ -1,16 +1,19 @@
 use std::collections::BTreeMap;
 
 use itertools::Itertools;
-use powdr_ast::analyzed::{AlgebraicExpression, AlgebraicReferenceThin, Analyzed, PolyID};
+use powdr_ast::analyzed::{AlgebraicExpression, AlgebraicReferenceThin, Analyzed};
 use powdr_backend_utils::{machine_fixed_columns, machine_witness_columns};
-use powdr_executor::{constant_evaluator::VariablySizedColumn, witgen::WitgenCallback};
+use powdr_executor::{
+    constant_evaluator::VariablySizedColumn,
+    witgen::{evaluators::expression_evaluator::OwnedTraceValues, WitgenCallback},
+};
 use powdr_number::{DegreeType, FieldElement};
 
 /// A collection of columns with self-contained constraints.
 pub struct Machine<'a, F> {
     pub machine_name: String,
     pub size: usize,
-    pub columns: BTreeMap<PolyID, Vec<F>>,
+    pub trace_values: OwnedTraceValues<F>,
     pub pil: &'a Analyzed<F>,
     pub intermediate_definitions: BTreeMap<AlgebraicReferenceThin, AlgebraicExpression<F>>,
 }
@@ -46,31 +49,20 @@ impl<'a, F: FieldElement> Machine<'a, F> {
 
         let fixed = machine_fixed_columns(fixed, pil);
         let fixed = fixed.get(&(size as DegreeType)).unwrap();
+        let fixed = fixed
+            .iter()
+            // TODO: Avoid clone?
+            .map(|(name, col)| (name.clone(), col.to_vec()))
+            .collect::<Vec<_>>();
 
         let intermediate_definitions = pil.intermediate_definitions();
 
-        let mut columns_by_name = witness
-            .into_iter()
-            // TODO: Avoid clone?
-            .chain(fixed.iter().map(|(name, col)| (name.clone(), col.to_vec())))
-            .collect::<BTreeMap<_, _>>();
-
-        let columns = pil
-            .committed_polys_in_source_order()
-            .chain(pil.constant_polys_in_source_order())
-            .flat_map(|(symbol, _)| symbol.array_elements())
-            .map(|(name, poly_id)| {
-                let column = columns_by_name
-                    .remove(&name)
-                    .unwrap_or_else(|| panic!("Missing column: {name}"));
-                (poly_id, column)
-            })
-            .collect();
+        let trace_values = OwnedTraceValues::new(pil, witness, fixed);
 
         Some(Self {
             machine_name,
             size,
-            columns,
+            trace_values,
             pil,
             intermediate_definitions,
         })

backend/src/mock/mod.rs

@@ -20,7 +20,6 @@ use powdr_number::{DegreeType, FieldElement};
 use crate::{Backend, BackendFactory, BackendOptions, Error, Proof};
 
 mod connection_constraint_checker;
-mod evaluator;
 mod machine;
 mod polynomial_constraint_checker;
 

backend/src/mock/polynomial_constraint_checker.rs

@@ -1,15 +1,15 @@
 use std::{collections::BTreeMap, fmt};
 
 use powdr_ast::{
-    analyzed::{Identity, PolynomialIdentity},
+    analyzed::{AlgebraicExpression, Identity, PolynomialIdentity},
     parsed::visitor::AllChildren,
 };
-use powdr_executor::witgen::{AffineExpression, AlgebraicVariable, ExpressionEvaluator};
+use powdr_executor::witgen::{
+    evaluators::expression_evaluator::ExpressionEvaluator, AlgebraicVariable,
+};
 use powdr_number::FieldElement;
 use rayon::iter::{IntoParallelIterator, ParallelIterator};
 
-use crate::mock::evaluator::Variables;
-
 use super::machine::Machine;
 
 pub struct PolynomialConstraintChecker<'a, F> {
@@ -55,35 +55,36 @@ impl<'a, F: FieldElement> PolynomialConstraintChecker<'a, F> {
         row: usize,
         identities: &[&'a Identity<F>],
     ) -> Vec<FailingPolynomialConstraint<'a, F>> {
-        let variables = Variables {
-            machine: self.machine,
-            row,
-            challenges: self.challenges,
-        };
-        let mut evaluator =
-            ExpressionEvaluator::new(&variables, &self.machine.intermediate_definitions);
+        let mut evaluator = ExpressionEvaluator::new(
+            self.machine.trace_values.row(row),
+            &self.machine.intermediate_definitions,
+            self.challenges,
+        );
         identities
             .iter()
             .filter_map(|identity| {
                 let identity = match identity {
                     Identity::Polynomial(polynomial_identity) => polynomial_identity,
                     _ => unreachable!("Unexpected identity: {}", identity),
                 };
-                let result = evaluator.evaluate(&identity.expression).unwrap();
-                let result = match result {
-                    AffineExpression::Constant(c) => c,
-                    _ => unreachable!("Unexpected result: {:?}", result),
-                };
+                let result = evaluator.evaluate(&identity.expression);
 
                 if result != F::zero() {
-                    let used_variables = identity
-                        .all_children()
-                        .filter_map(|child| child.try_into().ok());
+                    let used_variables = identity.all_children().filter(|expr| match expr {
+                        AlgebraicExpression::Reference(_)
+                        | AlgebraicExpression::PublicReference(_)
+                        | AlgebraicExpression::Challenge(_) => true,
+                        AlgebraicExpression::Number(_)
+                        | AlgebraicExpression::BinaryOperation(_)
+                        | AlgebraicExpression::UnaryOperation(_) => false,
+                    });
                     Some(FailingPolynomialConstraint {
                         row,
                         identity,
                         assignments: used_variables
-                            .map(|variable| (variable, variables.constant_value(variable)))
+                            .map(|variable| {
+                                (variable.try_into().unwrap(), evaluator.evaluate(variable))
+                            })
                             .collect(),
                     })
                 } else {