1 files changed, 83 insertions, 20 deletions

diff --git a/src/lib.rs b/src/lib.rs
index db3186a..8d8c3a0 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -182,7 +182,7 @@ pub fn poly_division(a: &[BigInt], b: &[BigInt], p: BigInt) -> (Vec<BigInt>, Vec
 
         let leading_coefficient = b_leading_inverse.clone() * cand_leading;
 
-        quotient.push(leading_coefficient.clone());
+        quotient.push(leading_coefficient.clone().rem_euclid(&p));
 
         if leading_coefficient != conversion(0) {
             for (i, cand_ref) in candidate.iter_mut().skip(diff).take(b.len()).enumerate() {
@@ -367,8 +367,9 @@ fn derivative(poly: &[BigInt], p: usize) -> Vec<BigInt> {
     }
 }
 
-/// Return true if and only if the polynomial has no repeated factors
-/// modulo the prime p.
+/// Return Ok() if and only if the polynomial has no repeated factors
+/// modulo the prime p, otherwise return an error containing the gcd
+/// of the polynomial with its derivative.
 pub fn is_square_free(poly: &[BigInt], p: usize) -> Result<(), Vec<BigInt>> {
     let derivative = derivative(poly, p);
 
@@ -535,29 +536,91 @@ pub fn berlekamp(poly: &[BigInt], p: usize) -> BeRe {
 /// p.
 ///
 /// If the polynomial has repeated factors, return an error.
-#[allow(clippy::result_unit_err)]
-pub fn factors(poly: &[BigInt], p: usize) -> Result<Vec<Vec<BigInt>>, Vec<BigInt>> {
-    if let Err(gcd) = is_square_free(poly, p) {
-        return Err(gcd);
-    }
+pub fn factors(poly: &[BigInt], p: usize) -> Vec<(usize, Vec<BigInt>)> {
+    let poly = trim(poly);
+
+    let mut result: Vec<(usize, Vec<BigInt>)> = Vec::new();
+    let mut candidates: Vec<(usize, Vec<BigInt>)> = vec![(1, poly.to_vec())];
+
+    'candidate_loop: while let Some((multiplicity, poly)) = candidates.pop() {
+        if let Err(gcd) = is_square_free(&poly, p) {
+            if gcd.len() == poly.len() {
+                // the derivative is zero
+                let pth_root = fake_p_th_root(&poly, p);
+
+                if pth_root.is_empty() {
+                    return Vec::new();
+                }
+
+                candidates.push((multiplicity * p, pth_root));
+
+                continue 'candidate_loop;
+            }
+
+            let (quotient, _) = poly_division(&poly, &gcd, conversion(p));
+
+            candidates.push((multiplicity, quotient));
+            candidates.push((multiplicity, gcd));
 
-    let mut result = Vec::new();
-    let mut candidates: Vec<Vec<BigInt>> = vec![poly.to_vec()];
+            continue 'candidate_loop;
+        }
 
-    while let Some(poly) = candidates.pop() {
         let bere = berlekamp(&poly, p);
 
         match bere {
             BeRe::Prod(x, y) => {
-                candidates.push(x);
-                candidates.push(y);
+                candidates.push((multiplicity, x));
+                candidates.push((multiplicity, y));
+            }
+            BeRe::Irreducible => {
+                result.push((multiplicity, poly));
             }
-            BeRe::Irreducible => result.push(poly),
-            BeRe::NonSquareFree => return Err(Vec::new()),
+            BeRe::NonSquareFree => panic!("this should be impossible"),
         }
     }
 
-    Ok(result)
+    result
+}
+
+/// Return the p-th root of the polynomial, if it is a p-th power
+/// indeed.
+///
+/// Precisely speaking, a monomial cx^{pn+k} in the polynomial
+/// contributes to the term cx^n in the new polynomial, where k is any
+/// integer such that 0 <= k < p.
+fn fake_p_th_root(poly: &[BigInt], p: usize) -> Vec<BigInt> {
+    let big_p = conversion(p);
+    let big_zero = conversion(0);
+
+    let iter = poly.iter().enumerate().filter_map(|(n, c)| {
+        let rem = c.rem_euclid(&big_p);
+
+        (&rem != &big_zero).then(|| (n.div_euclid(p), rem))
+    });
+
+    let mut degrees_map: std::collections::HashMap<usize, BigInt> = Default::default();
+
+    degrees_map.extend(iter);
+
+    if degrees_map.is_empty() {
+        return Vec::new();
+    }
+
+    let mut max = 0usize;
+
+    for (n, _) in degrees_map.iter() {
+        max = std::cmp::max(max, *n);
+    }
+
+    let mut result: Vec<BigInt> = std::iter::repeat_with(|| conversion(0))
+        .take(max + 1)
+        .collect();
+
+    for (n, c) in degrees_map {
+        *result.get_mut(n).unwrap() = c;
+    }
+
+    result
 }
 
 /// Return a sub-slice that trims the leading zeroes.
@@ -764,12 +827,12 @@ impl<'a> std::fmt::Display for Poly<'a> {
             write!(
                 f,
                 "{}{}{}",
-                if n + 1 >= poly.len() {
+                if coeff < &conversion(0) {
+                    " - "
+                } else if n + 1 >= poly.len() {
                     ""
-                } else if coeff >= &conversion(0) {
-                    " + "
                 } else {
-                    " - "
+                    " + "
                 },
                 Coefficient(coeff.clone()),
                 Monomial(n),