polynomial: add some extra functionality enabled by FieldVec

Author: apoelstra

src/lib.rs

@@ -134,6 +134,8 @@
 #![cfg_attr(docsrs, feature(doc_auto_cfg))]
 // Coding conventions
 #![deny(missing_docs)]
+#![allow(clippy::suspicious_arithmetic_impl)] // this lint is literally always wrong
+#![allow(clippy::suspicious_op_assign_impl)] // ...and "always wrong" loves company
 
 #[cfg(bench)]
 extern crate test;

src/primitives/fieldvec.rs

@@ -59,7 +59,7 @@
 
 #[cfg(all(feature = "alloc", not(feature = "std")))]
 use alloc::vec::Vec;
-use core::{iter, ops, slice};
+use core::{iter, mem, ops, slice};
 
 use super::Field;
 use crate::primitives::correction::NO_ALLOC_MAX_LENGTH;
@@ -206,12 +206,45 @@ impl<F: Default> FieldVec<F> {
             self.inner_a[self.len - 1] = item;
         } else {
             if self.len == NO_ALLOC_MAX_LENGTH + 1 {
-                let inner_a = core::mem::take(&mut self.inner_a);
+                let inner_a = mem::take(&mut self.inner_a);
                 self.inner_v = inner_a.into();
             }
             self.inner_v.push(item);
         }
     }
+
+    /// Pops an item off the end of the vector.
+    ///
+    /// # Panics
+    ///
+    /// Panics if [`Self::has_data`] is false.
+    pub fn pop(&mut self) -> Option<F> {
+        self.assert_has_data();
+        if self.len == 0 {
+            return None;
+        }
+
+        self.len -= 1;
+        #[cfg(not(feature = "alloc"))]
+        {
+            Some(mem::take(&mut self.inner_a[self.len]))
+        }
+
+        #[cfg(feature = "alloc")]
+        if self.len < NO_ALLOC_MAX_LENGTH {
+            Some(mem::take(&mut self.inner_a[self.len]))
+        } else {
+            use core::convert::TryFrom;
+
+            let ret = self.inner_v.pop();
+            let inner_v = mem::take(&mut self.inner_v);
+            match <[F; NO_ALLOC_MAX_LENGTH]>::try_from(inner_v) {
+                Ok(arr) => self.inner_a = arr,
+                Err(vec) => self.inner_v = vec,
+            }
+            ret
+        }
+    }
 }
 
 impl<F: Clone + Default> iter::FromIterator<F> for FieldVec<F> {

src/primitives/polynomial.rs

@@ -2,7 +2,7 @@
 
 //! Polynomials over Finite Fields
 
-use core::{iter, ops};
+use core::{fmt, iter, ops, slice};
 
 use super::{ExtensionField, Field, FieldVec};
 
@@ -15,6 +15,13 @@ pub struct Polynomial<F> {
 }
 
 impl<F: Field> Polynomial<F> {
+    /// Determines whether the residue is representable, given the current
+    /// compilation context.
+    pub fn has_data(&self) -> bool { self.inner.has_data() }
+
+    /// Panics if [`Self::has_data`] is false, with an informative panic message.
+    pub fn assert_has_data(&self) { self.inner.assert_has_data() }
+
     /// Provides access to the underlying [`FieldVec`].
     pub fn into_inner(self) -> FieldVec<F> { self.inner }
 
@@ -39,6 +46,18 @@ impl<F: Field> Polynomial<F> {
         degree_without_leading_zeros - leading_zeros
     }
 
+    /// An iterator over the coefficients of the polynomial.
+    ///
+    /// Yields value in "little endian" order; that is, the constant term is returned first.
+    ///
+    /// # Panics
+    ///
+    /// Panics if [`Self::has_data`] is false.
+    pub fn iter(&self) -> slice::Iter<F> {
+        self.assert_has_data();
+        self.inner.iter()
+    }
+
     /// The leading term of the polynomial.
     ///
     /// For the constant 0 polynomial, will return 0.
@@ -55,7 +74,15 @@ impl<F: Field> Polynomial<F> {
     /// factor of the polynomial.
     pub fn zero_is_root(&self) -> bool { self.inner.is_empty() || self.leading_term() == F::ZERO }
 
-    /// An iterator over the roots of the residue, when interpreted as a polynomial.
+    /// Helper function to add leading 0 terms until the polynomial has a specified
+    /// length.
+    fn zero_pad_up_to(&mut self, len: usize) {
+        while self.inner.len() < len {
+            self.inner.push(F::default());
+        }
+    }
+
+    /// An iterator over the roots of the polynomial.
     ///
     /// Takes a base element `base`. The roots of the residue will be yielded as
     /// nonnegative integers between 0 and 1 less than the order of the base,
@@ -201,6 +228,19 @@ impl<F: Field> Polynomial<F> {
     }
 }
 
+impl<F: Field> fmt::Display for Polynomial<F> {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        if self.has_data() {
+            for fe in self.iter() {
+                write!(f, "{}", fe)?;
+            }
+            Ok(())
+        } else {
+            f.write_str("<residue>")
+        }
+    }
+}
+
 impl<F: Field> iter::FromIterator<F> for Polynomial<F> {
     #[inline]
     fn from_iter<I>(iter: I) -> Self
@@ -215,6 +255,59 @@ impl<F> From<FieldVec<F>> for Polynomial<F> {
     fn from(inner: FieldVec<F>) -> Self { Self { inner } }
 }
 
+impl<F: Field> ops::Add<&Polynomial<F>> for Polynomial<F> {
+    type Output = Polynomial<F>;
+
+    fn add(mut self, other: &Polynomial<F>) -> Polynomial<F> {
+        self += other;
+        self
+    }
+}
+
+impl<F: Field> ops::Add<Polynomial<F>> for Polynomial<F> {
+    type Output = Polynomial<F>;
+    fn add(self, other: Polynomial<F>) -> Polynomial<F> { self + &other }
+}
+
+impl<F: Field> ops::Sub<&Polynomial<F>> for Polynomial<F> {
+    type Output = Polynomial<F>;
+    fn sub(mut self, other: &Polynomial<F>) -> Polynomial<F> {
+        self -= other;
+        self
+    }
+}
+
+impl<F: Field> ops::Sub<Polynomial<F>> for Polynomial<F> {
+    type Output = Polynomial<F>;
+    fn sub(self, other: Polynomial<F>) -> Polynomial<F> { self - &other }
+}
+
+impl<F: Field> ops::AddAssign<&Polynomial<F>> for Polynomial<F> {
+    fn add_assign(&mut self, other: &Self) {
+        self.zero_pad_up_to(other.inner.len());
+        for i in 0..other.inner.len() {
+            self.inner[i] += &other.inner[i];
+        }
+    }
+}
+
+impl<F: Field> ops::AddAssign for Polynomial<F> {
+    fn add_assign(&mut self, other: Polynomial<F>) { *self += &other; }
+}
+
+impl<F: Field> ops::SubAssign<&Polynomial<F>> for Polynomial<F> {
+    fn sub_assign(&mut self, other: &Polynomial<F>) {
+        self.zero_pad_up_to(other.inner.len());
+        for i in 0..other.inner.len() {
+            self.inner[i] -= &other.inner[i];
+        }
+    }
+}
+
+impl<F: Field> ops::SubAssign for Polynomial<F> {
+    fn sub_assign(&mut self, other: Polynomial<F>) { *self -= &other; }
+}
+
 /// An iterator over the roots of a polynomial.
 ///
 /// This iterator is constructed by the [`Polynomial::find_nonzero_distinct_roots`]