From d3d5f8917482ce5649602695829862a5df4ea712 Mon Sep 17 00:00:00 2001
From: Tom French <15848336+TomAFrench@users.noreply.github.com>
Date: Tue, 28 Feb 2023 15:49:30 +0000
Subject: [PATCH] feat(acir): add useful methods from `noirc_evaluator` onto
 `Expression` (#125)

fix(acir): correctly display expressions with non-unit coefficients
---
 acir/src/native_types/arithmetic.rs | 87 ++++++++++++++++++++++++++---
 1 file changed, 80 insertions(+), 7 deletions(-)

diff --git a/acir/src/native_types/arithmetic.rs b/acir/src/native_types/arithmetic.rs
index 67838ad38..1a25ddb82 100644
--- a/acir/src/native_types/arithmetic.rs
+++ b/acir/src/native_types/arithmetic.rs
@@ -41,8 +41,8 @@ impl Default for Expression {
 
 impl std::fmt::Display for Expression {
     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
-        if self.mul_terms.is_empty() && self.linear_combinations.len() == 1 && self.q_c.is_zero() {
-            write!(f, "x{}", self.linear_combinations[0].1.witness_index())
+        if let Some(witness) = self.to_witness() {
+            write!(f, "x{}", witness.witness_index())
         } else {
             write!(f, "%{:?}%", crate::circuit::opcodes::Opcode::Arithmetic(self.clone()))
         }
@@ -149,11 +149,6 @@ impl Expression {
         Ok(expr)
     }
 
-    // TODO: possibly rename, since linear can have one mul_term
-    pub fn is_linear(&self) -> bool {
-        self.mul_terms.is_empty()
-    }
-
     pub fn term_addition(&mut self, coefficient: acir_field::FieldElement, variable: Witness) {
         self.linear_combinations.push((coefficient, variable))
     }
@@ -166,10 +161,88 @@ impl Expression {
         self.mul_terms.push((coefficient, lhs, rhs))
     }
 
+    /// Returns `true` if the expression represents a constant polynomial.
+    ///
+    /// Examples:
+    /// -  f(x,y) = x + y would return false
+    /// -  f(x,y) = xy would return false, the degree here is 2
+    /// -  f(x,y) = 5 would return true, the degree is 0
     pub fn is_const(&self) -> bool {
         self.mul_terms.is_empty() && self.linear_combinations.is_empty()
     }
 
+    /// Returns `true` if highest degree term in the expression is one.
+    ///
+    /// - `mul_term` in an expression contains degree-2 terms
+    /// - `linear_combinations` contains degree-1 terms
+    /// Hence, it is sufficient to check that there are no `mul_terms`
+    ///
+    /// Examples:
+    /// -  f(x,y) = x + y would return true
+    /// -  f(x,y) = xy would return false, the degree here is 2
+    /// -  f(x,y) = 0 would return true, the degree is 0
+    pub fn is_linear(&self) -> bool {
+        self.mul_terms.is_empty()
+    }
+
+    /// Returns `true` if the expression can be seen as a degree-1 univariate polynomial
+    ///
+    /// - `mul_terms` in an expression can be univariate, however unless the coefficient
+    /// is zero, it is always degree-2.
+    /// - `linear_combinations` contains the sum of degree-1 terms, these terms do not
+    /// need to contain the same variable and so it can be multivariate. However, we
+    /// have thus far only checked if `linear_combinations` contains one term, so this
+    /// method will return false, if the `Expression` has not been simplified.
+    ///
+    /// Hence, we check in the simplest case if an expression is a degree-1 univariate,
+    /// by checking if it contains no `mul_terms` and it contains one `linear_combination` term.
+    ///
+    /// Examples:
+    /// - f(x,y) = x would return true
+    /// - f(x,y) = x + 6 would return true
+    /// - f(x,y) = 2*y + 6 would return true
+    /// - f(x,y) = x + y would return false
+    /// - f(x,y) = x + x should return true, but we return false *** (we do not simplify)
+    /// - f(x,y) = 5 would return false
+    pub fn is_degree_one_univariate(&self) -> bool {
+        self.is_linear() && self.linear_combinations.len() == 1
+    }
+
+    /// Returns a `FieldElement` if the expression represents a constant polynomial.
+    /// Otherwise returns `None`.
+    ///
+    /// Examples:
+    /// - f(x,y) = x would return `None`
+    /// - f(x,y) = x + 6 would return `None`
+    /// - f(x,y) = 2*y + 6 would return `None`
+    /// - f(x,y) = x + y would return `None`
+    /// - f(x,y) = 5 would return `FieldElement(5)`
+    pub fn to_const(&self) -> Option<FieldElement> {
+        self.is_const().then_some(self.q_c)
+    }
+
+    /// Returns a `Witness` if the `Expression` can be represented as a degree-1
+    /// univariate polynomial. Otherwise returns `None`.
+    ///
+    /// Note that `Witness` is only capable of expressing polynomials of the form
+    /// f(x) = x and not polynomials of the form f(x) = mx+c , so this method has
+    /// extra checks to ensure that m=1 and c=0
+    pub fn to_witness(&self) -> Option<Witness> {
+        if self.is_degree_one_univariate() {
+            // If we get here, we know that our expression is of the form `f(x) = mx+c`
+            // We want to now restrict ourselves to expressions of the form f(x) = x
+            // ie where the constant term is 0 and the coefficient in front of the variable is
+            // one.
+            let (coefficient, variable) = self.linear_combinations[0];
+            let constant = self.q_c;
+
+            if coefficient.is_one() && constant.is_zero() {
+                return Some(variable);
+            }
+        }
+        None
+    }
+
     fn get_max_idx(&self) -> WitnessIdx {
         WitnessIdx {
             linear: self.linear_combinations.len(),