chore: cleanup logic in mul_with_witness (#1980)

* cleanup logic

* code review: use CoW

* clippy

* fmt

* add squaring optimization

* chore: minor documentation nits

* chore: move `mul_with_witness` into `generated_acir.rs`

* chore: better placement of `mul_with_witness`

---------

Co-authored-by: TomAFrench <[email protected]>

crates/noirc_evaluator/src/ssa_refactor/acir_gen/acir_ir/generated_acir.rs

@@ -338,6 +338,51 @@ impl GeneratedAcir {
         lhs.add_mul(FieldElement::from(2_i128), &inter.unwrap())
     }
 
+    /// Returns an expression which represents `lhs * rhs`
+    ///
+    /// If one has multiplicative term and the other is of degree one or more,
+    /// the function creates [intermediate variables][`Witness`] accordingly.
+    /// There are two cases where we can optimize the multiplication between two expressions:
+    /// 1. If both expressions have at most a total degree of 1 in each term, then we can just multiply them
+    /// as each term in the result will be degree-2.
+    /// 2. If one expression is a constant, then we can just multiply the constant with the other expression
+    ///
+    /// (1) is because an [`Expression`] can hold at most a degree-2 univariate polynomial
+    /// which is what you get when you multiply two degree-1 univariate polynomials.
+    pub(crate) fn mul_with_witness(&mut self, lhs: &Expression, rhs: &Expression) -> Expression {
+        use std::borrow::Cow;
+        let lhs_is_linear = lhs.is_linear();
+        let rhs_is_linear = rhs.is_linear();
+
+        // Case 1: Both expressions have at most a total degree of 1 in each term
+        if lhs_is_linear && rhs_is_linear {
+            return (lhs * rhs)
+                .expect("one of the expressions is a constant and so this should not fail");
+        }
+
+        // Case 2: One or both of the sides needs to be reduced to a degree-1 univariate polynomial
+        let lhs_reduced = if lhs_is_linear {
+            Cow::Borrowed(lhs)
+        } else {
+            Cow::Owned(self.get_or_create_witness(lhs).into())
+        };
+
+        // If the lhs and rhs are the same, then we do not need to reduce
+        // rhs, we only need to square the lhs.
+        if lhs == rhs {
+            return (&*lhs_reduced * &*lhs_reduced)
+                .expect("Both expressions are reduced to be degree<=1");
+        };
+
+        let rhs_reduced = if rhs_is_linear {
+            Cow::Borrowed(rhs)
+        } else {
+            Cow::Owned(self.get_or_create_witness(rhs).into())
+        };
+
+        (&*lhs_reduced * &*rhs_reduced).expect("Both expressions are reduced to be degree<=1")
+    }
+
     /// Signed division lhs /  rhs
     /// We derive the signed division from the unsigned euclidian division.
     /// note that this is not euclidian division!

crates/noirc_evaluator/src/ssa_refactor/acir_gen/acir_ir/sort.rs

@@ -1,6 +1,5 @@
-use acvm::acir::native_types::{Expression, Witness};
-
 use super::generated_acir::GeneratedAcir;
+use acvm::acir::native_types::{Expression, Witness};
 
 impl GeneratedAcir {
     // Generates gates for a sorting network
@@ -71,31 +70,4 @@ impl GeneratedAcir {
         conf.extend(w2);
         (conf, out_expr)
     }
-
-    /// Returns an expression which represents a*b
-    /// If one has multiplicative term and the other is of degree one or more,
-    /// the function creates intermediate variables accordindly
-    pub(super) fn mul_with_witness(&mut self, a: &Expression, b: &Expression) -> Expression {
-        let a_arith;
-        let a_arith = if !a.mul_terms.is_empty() && !b.is_const() {
-            let a_witness = self.get_or_create_witness(a);
-            a_arith = Expression::from(a_witness);
-            &a_arith
-        } else {
-            a
-        };
-        let b_arith;
-        let b_arith = if !b.mul_terms.is_empty() && !a.is_const() {
-            if a == b {
-                a_arith
-            } else {
-                let b_witness = self.get_or_create_witness(a);
-                b_arith = Expression::from(b_witness);
-                &b_arith
-            }
-        } else {
-            b
-        };
-        (a_arith * b_arith).expect("Both expressions are reduced to be degree<=1")
-    }
 }