From 2fe6739b42954953227cd0f621e54b7dae898e14 Mon Sep 17 00:00:00 2001
From: chriseth <chris@ethereum.org>
Date: Thu, 23 Jan 2025 16:20:12 +0100
Subject: [PATCH] Also propagate range constraints if the symbolic evaluation
 is not exhausted (#2381)

---
 executor/src/witgen/jit/witgen_inference.rs | 44 ++++++++-------------
 1 file changed, 16 insertions(+), 28 deletions(-)

diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs
index 715f82ca08..14c5580947 100644
--- a/executor/src/witgen/jit/witgen_inference.rs
+++ b/executor/src/witgen/jit/witgen_inference.rs
@@ -237,42 +237,30 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> WitgenInference<'a, T, F
         offset: i32,
         rhs: &VariableOrValue<T, Variable>,
     ) -> Result<ProcessResult<T, Variable>, Error> {
-        // First we try to find a new assignment to a variable in the equality.
-
-        let evaluator = Evaluator::new(self);
-        let Some(lhs_evaluated) = evaluator.evaluate(lhs, offset) else {
-            return Ok(ProcessResult::empty());
-        };
-
-        let rhs_evaluated = match rhs {
-            VariableOrValue::Variable(v) => evaluator.evaluate_variable(v.clone()),
-            VariableOrValue::Value(v) => (*v).into(),
-        };
-
-        let result = (lhs_evaluated - rhs_evaluated).solve()?;
-        if result.complete && result.effects.is_empty() {
-            // A complete result without effects means that there were no unknowns
-            // in the constraint.
-            // We try again, but this time we treat all non-concrete variables
-            // as unknown and in that way try to find new concrete values for
-            // already known variables.
-            let result = self.process_equality_on_row_concrete(lhs, offset, rhs)?;
-            if !result.effects.is_empty() {
-                return Ok(result);
-            }
-        }
+        // Try to find a new assignment to a variable in the equality.
+        let mut result = self.process_equality_on_row_using_evaluator(lhs, offset, rhs, false)?;
+        // Try to propagate range constraints.
+        let result_concrete =
+            self.process_equality_on_row_using_evaluator(lhs, offset, rhs, true)?;
+
+        // We only use the effects of the second evaluation,
+        // its `complete` flag is ignored.
+        result.effects.extend(result_concrete.effects);
         Ok(result)
     }
 
-    /// Process an equality but only consider concrete variables as known
-    /// and thus propagate range constraints across the equality.
-    fn process_equality_on_row_concrete(
+    fn process_equality_on_row_using_evaluator(
         &self,
         lhs: &Expression<T>,
         offset: i32,
         rhs: &VariableOrValue<T, Variable>,
+        only_concrete_known: bool,
     ) -> Result<ProcessResult<T, Variable>, Error> {
-        let evaluator = Evaluator::new(self).only_concrete_known();
+        let evaluator = if only_concrete_known {
+            Evaluator::new(self).only_concrete_known()
+        } else {
+            Evaluator::new(self)
+        };
         let Some(lhs_evaluated) = evaluator.evaluate(lhs, offset) else {
             return Ok(ProcessResult::empty());
         };