fix: avoid unnecessarily splitting expressions with multiplication terms with a shared term

acvm-repo/acvm/src/compiler/transformers/csat.rs

@@ -87,7 +87,7 @@
     // This optimization will search for combinations of terms which can be represented in a single assert-zero opcode
     // Case 1 : qM * wL * wR + qL * wL + qR * wR + qO * wO + qC
     // This polynomial does not require any further optimizations, it can be safely represented in one opcode
     // ie a polynomial with 1 mul(bi-variate) term and 3 (univariate) terms where 2 of those terms match the bivariate term
     // wL and wR, we can represent it in one opcode
     // GENERALIZED for WIDTH: instead of the number 3, we use `WIDTH`
     //
@@ -427,15 +427,15 @@
 
     // A polynomial with width-2 fan-in terms and a single non-zero mul term can fit into one opcode
     // Example: Axy + Dz . Notice, that the mul term places a constraint on the first two terms, but not the last term
     // XXX: This would change if our arithmetic polynomial equation was changed to Axyz for example, but for now it is not.
     if expr.linear_combinations.len() <= (width - 2) {
         return true;
     }
 
-    // We now know that we have a single mul term. We also know that the mul term must match up with two other terms
+    // We now know that we have a single mul term. We also know that the mul term must match up with at least one of the other terms
     // A polynomial whose mul terms are non zero which do not match up with two terms in the fan-in cannot fit into one opcode
     // An example of this is: Axy + Bx + Cy + ...
     // Notice how the bivariate monomial xy has two univariate monomials with their respective coefficients
     // XXX: note that if x or y is zero, then we could apply a further optimization, but this would be done in another algorithm.
     // It would be the same as when we have zero coefficients - Can only work if wire is constrained to be zero publicly
     let mul_term = &expr.mul_terms[0];
@@ -461,7 +461,25 @@
         }
     }
 
-    found_x & found_y
+    // If the multiplication is a squaring then we must assign the two witnesses to separate wires and so we
+    // can never get a zero contribution to the width.
+    let multiplication_is_squaring = mul_term.1 == mul_term.2;
+
+    let mul_term_width_contribution = if !multiplication_is_squaring && (found_x & found_y) {
+        // Both witnesses involved in the multiplication exist elsewhere in the expression.
+        // They both do not contribute to the width of the expression as this would be double-counting
+        // due to their appearance in the linear terms.
+        0
+    } else if found_x || found_y {
+        // One of the witnesses involved in the multiplication exists elsewhere in the expression.
+        // The multiplication then only contributes 1 new witness to the width.
+        1
+    } else {
+        // Worst case scenario, the multiplication is using completely unique witnesses so has a contribution of 2.
+        2
+    };
+
+    mul_term_width_contribution + expr.linear_combinations.len() <= width
 }
 
 #[cfg(test)]
@@ -573,4 +591,20 @@
         let contains_b = got_optimized_opcode_a.linear_combinations.iter().any(|(_, w)| *w == b);
         assert!(contains_b);
     }
+
+    #[test]
+    fn recognize_expr_with_single_shared_witness_which_fits_in_single_identity() {
+        // Regression test for an expression which Zac found which should have been preserved but
+        // was being split into two expressions.
+        let expr = Expression {
+            mul_terms: vec![(-FieldElement::from(555u128), Witness(8), Witness(10))],
+            linear_combinations: vec![
+                (FieldElement::one(), Witness(10)),
+                (FieldElement::one(), Witness(11)),
+                (-FieldElement::one(), Witness(13)),
+            ],
+            q_c: FieldElement::zero(),
+        };
+        assert!(fits_in_one_identity(&expr, 4));
+    }
 }