From 958bd88478e10b2caffae4b0ccee41fa28e6dee9 Mon Sep 17 00:00:00 2001
From: Georg Wiese <georgwiese@gmail.com>
Date: Thu, 7 Sep 2023 16:16:04 +0200
Subject: [PATCH] Bug fix in `AffineExpression::assign()`

---
 executor/src/witgen/affine_expression.rs | 32 +++++++++++++++++++++---
 1 file changed, 29 insertions(+), 3 deletions(-)

diff --git a/executor/src/witgen/affine_expression.rs b/executor/src/witgen/affine_expression.rs
index 5ae9f85da..c2d1554b8 100644
--- a/executor/src/witgen/affine_expression.rs
+++ b/executor/src/witgen/affine_expression.rs
@@ -133,14 +133,12 @@ where
     /// Incorporates the case where the symbolic variable `key` is assigned
     /// the value `value`.
     pub fn assign(&mut self, key: K, value: T) {
-        let mut offset = 0.into();
         for (k, coeff) in &mut self.coefficients {
             if *k == key {
-                offset += *coeff * value;
+                self.offset += *coeff * value;
                 *coeff = 0.into();
             }
         }
-        self.offset -= offset;
         self.clean = false;
     }
 }
@@ -576,6 +574,34 @@ mod test {
         input.iter().map(|x| (*x).into()).enumerate().collect()
     }
 
+    #[test]
+    pub fn test_affine_assign() {
+        let mut a = AffineExpression::<_, GoldilocksField> {
+            coefficients: convert(vec![2, 3]),
+            offset: 3.into(),
+            clean: true,
+        };
+        a.assign(0, 3.into());
+        assert_eq!(
+            a,
+            AffineExpression {
+                coefficients: convert(vec![0, 3]),
+                offset: 9.into(),
+                clean: false
+            },
+        );
+
+        // Now, the expression is 3b + 9. It should be able to solve for b
+        // such that 3b + 9 = 0.
+        let updates = a.solve().unwrap();
+        assert_eq!(
+            updates.constraints,
+            [(1, Constraint::Assignment((-3).into()))]
+        );
+        a.assign(1, (-3).into());
+        assert_eq!(a.constant_value().unwrap(), 0.into());
+    }
+
     #[test]
     pub fn test_affine_neg() {
         let a = AffineExpression {