Simplify already in the evaluator in witgen. (#2390)

The evaluator in witgen_inference would always try to fully evaluate
each sub-expression into an affine expression. This could fail even if
we would have multiplied by zero later on anyway. As a bonus, we can now
use exactly this evaluator for the debug formatter, that tries to
simplify things like multiplication by zero.

executor/src/witgen/jit/debug_formatter.rs

@@ -1,17 +1,16 @@
 use itertools::Itertools;
 use powdr_ast::analyzed::{
     AlgebraicBinaryOperation, AlgebraicBinaryOperator, AlgebraicExpression as Expression,
-    AlgebraicUnaryOperation, AlgebraicUnaryOperator, Identity, LookupIdentity, PermutationIdentity,
-    PhantomLookupIdentity, PhantomPermutationIdentity, PolynomialIdentity, PolynomialType,
-    SelectedExpressions,
+    AlgebraicUnaryOperation, Identity, LookupIdentity, PermutationIdentity, PhantomLookupIdentity,
+    PhantomPermutationIdentity, PolynomialIdentity, PolynomialType, SelectedExpressions,
 };
 use powdr_number::FieldElement;
 
 use crate::witgen::range_constraints::RangeConstraint;
 
 use super::{
     variable::Variable,
-    witgen_inference::{FixedEvaluator, Value, WitgenInference},
+    witgen_inference::{FixedEvaluator, WitgenInference},
 };
 
 /// Returns a human-readable summary of the identities.
@@ -293,54 +292,8 @@ impl<T: FieldElement, FixedEval: FixedEvaluator<T>> DebugFormatter<'_, T, FixedE
     }
 
     fn try_to_known(&self, e: &Expression<T>, row_offset: i32) -> Option<T> {
-        match e {
-            Expression::Reference(r) => {
-                match r.poly_id.ptype {
-                    PolynomialType::Constant | PolynomialType::Committed => {
-                        let variable = Variable::from_reference(r, row_offset);
-                        match self.witgen.value(&variable) {
-                            Value::Concrete(v) => Some(v),
-                            _ => None,
-                        }
-                    }
-                    PolynomialType::Intermediate => {
-                        // TODO
-                        None
-                    }
-                }
-            }
-            Expression::PublicReference(_) => {
-                // TODO we need to introduce a variable type for those.
-                None
-            }
-            Expression::Challenge(_) => {
-                // TODO we need to introduce a variable type for those.
-                None
-            }
-            Expression::Number(n) => Some(*n),
-            Expression::BinaryOperation(op) => {
-                let left = self.try_to_known(&op.left, row_offset);
-                let right = self.try_to_known(&op.right, row_offset);
-                match op.op {
-                    AlgebraicBinaryOperator::Add => Some(left? + right?),
-                    AlgebraicBinaryOperator::Sub => Some(left? - right?),
-                    AlgebraicBinaryOperator::Mul => match (left, right) {
-                        (Some(a), _) | (_, Some(a)) if a == 0.into() => Some(0.into()),
-                        (Some(a), b) | (b, Some(a)) if a == 1.into() => b,
-                        (Some(a), b) | (b, Some(a)) if a == (-1).into() => b.map(|b| -b),
-                        (Some(l), Some(r)) => Some(l * r),
-                        _ => None,
-                    },
-                    AlgebraicBinaryOperator::Pow => Some(left?.pow(right?.to_integer())),
-                }
-            }
-            Expression::UnaryOperation(op) => {
-                let inner = self.try_to_known(&op.expr, row_offset);
-                match op.op {
-                    AlgebraicUnaryOperator::Minus => inner.map(|i| -i),
-                }
-            }
-        }
+        let v = self.witgen.evaluate(e, row_offset)?;
+        v.try_to_known()?.try_to_number()
     }
 }
 

executor/src/witgen/jit/witgen_inference.rs

@@ -21,6 +21,7 @@ use crate::witgen::{
 use super::{
     affine_symbolic_expression::{AffineSymbolicExpression, Error, ProcessResult},
     effect::{BranchCondition, Effect},
+    symbolic_expression::SymbolicExpression,
     variable::{Cell, MachineCallVariable, Variable},
 };
 
@@ -484,7 +485,7 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> WitgenInference<'a, T, F
             .unwrap_or_default()
     }
 
-    fn evaluate(
+    pub fn evaluate(
         &self,
         expr: &Expression<T>,
         offset: i32,
@@ -566,21 +567,28 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> Evaluator<'a, T, FixedEv
         op: &AlgebraicBinaryOperation<T>,
         offset: i32,
     ) -> Option<AffineSymbolicExpression<T, Variable>> {
-        let left = self.evaluate(&op.left, offset)?;
-        let right = self.evaluate(&op.right, offset)?;
+        let left = self.evaluate(&op.left, offset);
+        let right = self.evaluate(&op.right, offset);
         match op.op {
-            AlgebraicBinaryOperator::Add => Some(&left + &right),
-            AlgebraicBinaryOperator::Sub => Some(&left - &right),
-            AlgebraicBinaryOperator::Mul => left.try_mul(&right),
+            AlgebraicBinaryOperator::Add => Some(&left? + &right?),
+            AlgebraicBinaryOperator::Sub => Some(&left? - &right?),
+            AlgebraicBinaryOperator::Mul => {
+                if is_known_zero(&left) || is_known_zero(&right) {
+                    Some(SymbolicExpression::from(T::from(0)).into())
+                } else {
+                    left?.try_mul(&right?)
+                }
+            }
             AlgebraicBinaryOperator::Pow => {
-                let result = left
+                let result = left?
                     .try_to_known()?
                     .try_to_number()?
-                    .pow(right.try_to_known()?.try_to_number()?.to_integer());
+                    .pow(right?.try_to_known()?.try_to_number()?.to_integer());
                 Some(AffineSymbolicExpression::from(result))
             }
         }
     }
+
     fn evaluate_unary_operation(
         &self,
         op: &AlgebraicUnaryOperation<T>,
@@ -593,6 +601,12 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> Evaluator<'a, T, FixedEv
     }
 }
 
+fn is_known_zero<T: FieldElement>(x: &Option<AffineSymbolicExpression<T, Variable>>) -> bool {
+    x.as_ref()
+        .and_then(|x| x.try_to_known().map(|x| x.is_known_zero()))
+        .unwrap_or(false)
+}
+
 /// An equality constraint between an algebraic expression evaluated
 /// on a certain row offset and a variable or fixed constant value.
 #[derive(Clone, Ord, PartialOrd, Eq, PartialEq, Debug)]