Implement extended GCD and modular inverse

Author: vks

src/lib.rs

@@ -23,8 +23,9 @@ extern crate num_traits as traits;
 
 use core::mem;
 use core::ops::Add;
+use core::cmp::Ordering;
 
-use traits::{Num, Signed, Zero};
+use traits::{Num, NumRef, RefNum, Signed, Zero};
 
 mod roots;
 pub use roots::Roots;
@@ -1013,6 +1014,57 @@ impl_integer_for_usize!(usize, test_integer_usize);
 #[cfg(has_i128)]
 impl_integer_for_usize!(u128, test_integer_u128);
 
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub struct GcdResult<T> {
+    /// Greatest common divisor.
+    pub gcd: T,
+    /// Coefficients such that: gcd(a, b) = c1*a + c2*b
+    pub c1: T, pub c2: T,
+}
+
+/// Calculate greatest common divisor and the corresponding coefficients.
+pub fn extended_gcd<T: Integer + NumRef>(a: T, b: T) -> GcdResult<T>
+    where for<'a> &'a T: RefNum<T>
+{
+    // Euclid's extended algorithm
+    let (mut s, mut old_s) = (T::zero(), T::one());
+    let (mut t, mut old_t) = (T::one(), T::zero());
+    let (mut r, mut old_r) = (b, a);
+
+    while r != T::zero() {
+        let quotient = &old_r / &r;
+        old_r = old_r - &quotient * &r; std::mem::swap(&mut old_r, &mut r);
+        old_s = old_s - &quotient * &s; std::mem::swap(&mut old_s, &mut s);
+        old_t = old_t - quotient * &t; std::mem::swap(&mut old_t, &mut t);
+    }
+
+    let _quotients = (t, s); // == (a, b) / gcd
+
+    GcdResult { gcd: old_r, c1: old_s, c2: old_t }
+}
+
+/// Find the standard representation of a (mod n).
+pub fn normalize<T: Integer + NumRef>(a: T, n: &T) -> T {
+    let a = a % n;
+    match a.cmp(&T::zero()) {
+        Ordering::Less => a + n,
+        _ => a,
+    }
+}
+
+/// Calculate the inverse of a (mod n).
+pub fn inverse<T: Integer + NumRef + Clone>(a: T, n: &T) -> Option<T>
+    where for<'a> &'a T: RefNum<T>
+{
+    let GcdResult { gcd, c1: c, .. } = extended_gcd(a, n.clone());
+    if gcd == T::one() {
+        Some(normalize(c, n))
+    } else {
+        None
+    }
+}
+
+
 /// An iterator over binomial coefficients.
 pub struct IterBinomial<T> {
     a: T,
@@ -1169,6 +1221,24 @@ fn test_lcm_overflow() {
     check!(u64, 0x8000_0000_0000_0000, 0x02, 0x8000_0000_0000_0000);
 }
 
+#[test]
+fn test_extended_gcd() {
+    assert_eq!(extended_gcd(240, 46), GcdResult { gcd: 2, c1: -9, c2: 47});
+}
+
+#[test]
+fn test_normalize() {
+    assert_eq!(normalize(10, &7), 3);
+    assert_eq!(normalize(7, &7), 0);
+    assert_eq!(normalize(5, &7), 5);
+    assert_eq!(normalize(-3, &7), 4);
+}
+
+#[test]
+fn test_inverse() {
+    assert_eq!(inverse(5, &7).unwrap(), 3);
+}
+
 #[test]
 fn test_iter_binomial() {
     macro_rules! check_simple {