From 5ad9656dfce69e17d31430efa8ce1225864ca311 Mon Sep 17 00:00:00 2001
From: feltroid Prime <96737978+feltroidprime@users.noreply.github.com>
Date: Thu, 24 Oct 2024 01:04:25 +0700
Subject: [PATCH] Polynomial: differentiate & print methods.  (#929)

* add differentiate method to poly

* print_as_sage_poly

* avoid duplicate call to .degree().

* Fix wasm target compilation.

---------

Co-authored-by: Diego K <43053772+diegokingston@users.noreply.github.com>
---
 math/src/polynomial/mod.rs | 78 ++++++++++++++++++++++++++++++++++++--
 1 file changed, 75 insertions(+), 3 deletions(-)

diff --git a/math/src/polynomial/mod.rs b/math/src/polynomial/mod.rs
index ee27923d0..679a83aec 100644
--- a/math/src/polynomial/mod.rs
+++ b/math/src/polynomial/mod.rs
@@ -1,8 +1,8 @@
 use super::field::element::FieldElement;
-use crate::field::traits::{IsField, IsSubFieldOf};
-use alloc::{borrow::ToOwned, vec, vec::Vec};
+use crate::field::traits::{IsField, IsPrimeField, IsSubFieldOf};
+use alloc::string::{String, ToString};
+use alloc::{borrow::ToOwned, format, vec, vec::Vec};
 use core::{fmt::Display, ops};
-
 pub mod dense_multilinear_poly;
 mod error;
 pub mod sparse_multilinear_poly;
@@ -140,6 +140,19 @@ impl<F: IsField> Polynomial<FieldElement<F>> {
         self.coefficients().len()
     }
 
+    /// Returns the derivative of the polynomial with respect to x.
+    pub fn differentiate(&self) -> Self {
+        let degree = self.degree();
+        if degree == 0 {
+            return Polynomial::zero();
+        }
+        let mut derivative = Vec::with_capacity(degree);
+        for (i, coeff) in self.coefficients().iter().enumerate().skip(1) {
+            derivative.push(FieldElement::<F>::from(i as u64) * coeff);
+        }
+        Polynomial::new(&derivative)
+    }
+
     /// Computes quotient with `x - b` in place.
     pub fn ruffini_division_inplace(&mut self, b: &FieldElement<F>) {
         let mut c = FieldElement::zero();
@@ -302,6 +315,44 @@ impl<F: IsField> Polynomial<FieldElement<F>> {
     }
 }
 
+impl<F: IsPrimeField> Polynomial<FieldElement<F>> {
+    // Print the polynomial as a string ready to be used in SageMath, or just for pretty printing.
+    pub fn print_as_sage_poly(&self, var_name: Option<char>) -> String {
+        let var_name = var_name.unwrap_or('x');
+        if self.coefficients.is_empty()
+            || self.coefficients.len() == 1 && self.coefficients[0] == FieldElement::zero()
+        {
+            return String::new();
+        }
+
+        let mut string = String::new();
+        let zero = FieldElement::<F>::zero();
+
+        for (i, coeff) in self.coefficients.iter().rev().enumerate() {
+            if *coeff == zero {
+                continue;
+            }
+
+            let coeff_str = coeff.representative().to_string();
+
+            if i == self.coefficients.len() - 1 {
+                string.push_str(&coeff_str);
+            } else if i == self.coefficients.len() - 2 {
+                string.push_str(&format!("{}*{} + ", coeff_str, var_name));
+            } else {
+                string.push_str(&format!(
+                    "{}*{}^{} + ",
+                    coeff_str,
+                    var_name,
+                    self.coefficients.len() - 1 - i
+                ));
+            }
+        }
+
+        string
+    }
+}
+
 pub fn pad_with_zero_coefficients_to_length<F: IsField>(
     pa: &mut Polynomial<FieldElement<F>>,
     n: usize,
@@ -1177,4 +1228,25 @@ mod tests {
         assert_eq!(lhs, g);
         assert_eq!(g, p3);
     }
+
+    #[test]
+    fn test_differentiate() {
+        // 3x^2 + 2x + 42
+        let px = Polynomial::new(&[FE::new(42), FE::new(2), FE::new(3)]);
+        // 6x + 2
+        let dpdx = px.differentiate();
+        assert_eq!(dpdx, Polynomial::new(&[FE::new(2), FE::new(6)]));
+
+        // 128
+        let px = Polynomial::new(&[FE::new(128)]);
+        // 0
+        let dpdx = px.differentiate();
+        assert_eq!(dpdx, Polynomial::new(&[FE::new(0)]));
+    }
+
+    #[test]
+    fn test_print_as_sage_poly() {
+        let p = Polynomial::new(&[FE::new(1), FE::new(2), FE::new(3)]);
+        assert_eq!(p.print_as_sage_poly(None), "3*x^2 + 2*x + 1");
+    }
 }