nicolas 0.1.1

//! Finite Field Computation.

use std::ops::{Add, Div, Mul, Sub};

/// Binary Polynomials.
#[derive(Debug, Clone, Copy, Eq, PartialEq, Ord, PartialOrd, BinaryArithmetic)]
pub struct Poly2(u64);
impl Poly2 {
    /// Creates a new `Poly2` instance.
    pub fn new(n: u64) -> Self {
        Self(n)
    }
    /// Converts `Poly2` to `u64` representation.
    pub fn as_u64(self) -> u64 {
        self.0
    }
}

/// Galois Field modulo 2.
#[derive(Debug, Clone, Copy, Eq, PartialEq, Ord, PartialOrd, BinaryArithmetic)]
pub struct GF2(u8);
impl GF2 {
    /// Creates a new `GF2` instance.
    pub fn new(n: u8) -> Self {
        Self(n)
    }
    /// Converts `GF2` to `u8` representation.
    pub fn as_u8(self) -> u8 {
        self.0
    }
}
impl Div for GF2 {
    type Output = Self;
    fn div(self, _rhs: Self) -> Self::Output {
        self
    }
}

/// Galois Extension Filed Generator
#[derive(Debug, Clone, Copy, Eq, PartialEq, Ord, PartialOrd)]
pub struct GenGF2E {
    /// The order of GF(2^m).
    extension: u16,
    /// The irreducible polynomial, which generates GF(2^m).
    polynomial: Poly2,
}
impl GenGF2E {
    /// Creates a new `GenGF2E` instance.
    pub fn new(extension: u16, polynomial: Poly2) -> Self {
        Self {
            extension,
            polynomial,
        }
    }
    /// Creates a new `GF2E` instance with the given polynomial.
    pub fn gen(&self, n: Poly2) -> GF2E {
        GF2E::new(n, self.extension, self.polynomial)
    }
}

/// GF(2^m)(Galois Extension Filed) over GF(2)
#[derive(Debug, Clone, Copy, Eq, PartialEq, Ord, PartialOrd)]
pub struct GF2E(u64, u16, Poly2);
impl GF2E {
    fn new(n: Poly2, ext: u16, p: Poly2) -> Self {
        Self(n.as_u64(), ext, p)
    }
    /// Converts `GF2E` to `u64` representation.
    pub fn as_u64(&self) -> u64 {
        self.0
    }
    /// Returns the multiplicative inverse of this value.
    pub fn inverse(&self) -> Option<Self> {
        let p = self.polynomial();
        let ext = self.extension();
        gf2_inverse(self.0, (self.2).0, self.1).map(|n| Self(n, ext, p))
    }
    /// Returns the order of this extended field.
    fn extension(&self) -> u16 {
        self.1
    }
    /// Returns the irreducible polynomial, which generates this extended field.
    fn polynomial(&self) -> Poly2 {
        self.2
    }
}
impl Add for GF2E {
    type Output = Self;
    fn add(self, rhs: Self) -> Self::Output {
        let n = gf2_add(self.0, rhs.0);
        GF2E(n, self.1, self.2)
    }
}
impl Sub for GF2E {
    type Output = Self;
    fn sub(self, rhs: Self) -> Self::Output {
        // NOTE: Subtracting is the same as adding, because
        // -x = x is true in GF(p^m).
        let n = gf2_sub(self.0, rhs.0);
        GF2E(n, self.1, self.2)
    }
}
impl Mul for GF2E {
    type Output = Self;
    fn mul(self, rhs: Self) -> Self::Output {
        let p = (self.2).0;
        let n = gf2_multiply(self.0, rhs.0, p, self.1);
        GF2E(n, self.1, self.2)
    }
}

#[inline]
fn gf2_add(a: u64, b: u64) -> u64 {
    a ^ b
}

#[inline]
fn gf2_sub(a: u64, b: u64) -> u64 {
    a ^ b
}

fn gf2_multiply(mut a: u64, mut b: u64, p: u64, ext: u16) -> u64 {
    let mut c = 0;
    let overflow = 2u64.pow(u32::from(ext - 1));
    while a != 0 && b != 0 {
        if (b & 1) != 0 {
            c ^= a;
        }
        if a >= overflow {
            a = (a << 1) ^ p;
        } else {
            a <<= 1;
        }
        b >>= 1;
    }
    c
}

#[allow(dead_code)]
fn gf2_inverse_primitive(a: u64, ext: u16) -> Option<u64> {
    // irreducible polynomial が 原始多項式の時にしか計算できない
    if a == 0 {
        return None;
    }
    // a^(2^p - 2)
    let mut b = 1;
    let i = 2u64.pow(u32::from(ext)) - 2;
    for _ in 0..i {
        b *= a;
    }
    Some(b)
}

// TODO
fn _inverse(a: u64, p: u64) -> Option<u64> {
    // https://engineering.purdue.edu/kak/compsec/NewLectures/Lecture7.pdf
    let (g, x, _) = gcde2(p, a);
    if g == 1 {
        let (_, r) = euclidean_division2(x, p);
        Some(r)
    } else {
        None
    }
}

// See https://engineering.purdue.edu/kak/compsec/NewLectures/Lecture7.pdf.
fn gf2_inverse(a: u64, p: u64, ext: u16) -> Option<u64> {
    let mut x = p;
    let mut x1 = 1;
    let mut y = 1;
    let mut y1 = 0;
    let mut n = a;
    let mut m = p;
    while m != 0 {
        let (q, r) = euclidean_division2(n, m);
        n = m;
        m = r;
        //println!("n={:b},m={:b},q={:b},r={:b}", n, m, q, r);
        let x2 = x;
        x = x1 ^ gf2_multiply(q, x, p, ext);
        x1 = x2;
        let y2 = y;
        y = y1 ^ gf2_multiply(q, y, p, ext);
        y1 = y2;
    }
    if n == 1 {
        let (_, r) = euclidean_division2(x1 ^ p, p);
        Some(r)
    } else {
        None
    }
}

// n is a polynomial bit pattern and its coefficients are an element of GF(2).
fn degree(n: u64) -> u8 {
    match sffs::clz(n) {
        64 => 0,
        n => 63 - n as u8,
    }
}

// NOTE: The binary representation of a integer is alwasy 1 or 0.
fn lc2(n: u64) -> u8 {
    if n == 0 || n == 1 {
        0
    } else {
        1
    }
}

/// # Algorithm
/// f and g are polynomials(f >= g).
/// euclidean_division(f, g):
/// q <- 0
/// r <- f
/// while (deg(r) >= deg(g)) {
///   t <- lc(r) / lc(g) * x^(deg(r) - deg(g))
///   r <- r - t * g
///   q <- q + t
/// }
fn euclidean_division2(a: u64, b: u64) -> (u64, u64) {
    if b == 1 {
        return (a, 0);
    }
    let mut q = 0;
    let mut r = a;
    let db = degree(b);
    let lcb = lc2(b);
    while degree(r) >= db {
        let lcr = lc2(r);
        let t = if lcr == 0 || lcb == 0 {
            0
        } else {
            u64::from(degree(r) - db)
        };
        //println!("t={:b}, b={:b}, t&b={:b}", t, b, t & b);
        r ^= b << t;
        q ^= 1 << t;
        //println!(
        // "r={:b}, q={:b}, t={:b}, db={}, lcr={}, lcb={}",
        //r, q, t, db, lcr, lcb
        //);
    }
    (q, r)
}

/// Computes extended GCD on binary polynomials.
///
/// Suppose that there are variables x and y such that a*x + b*y = GCD(a, b),
/// a return value is (GCD(a, b), x, y).
#[allow(dead_code)]
fn gcde2(mut a: u64, mut b: u64) -> (u64, u64, u64) {
    let mut x0 = 1;
    let mut y0 = 0;
    let mut x1 = 0;
    let mut y1 = 1;
    while b != 0 {
        let (q, r) = euclidean_division2(a, b);
        a = b;
        b = r;
        let x2 = x0;
        x0 = x1;
        x1 = x2 ^ q & x1;
        let y2 = y0;
        y0 = y1;
        y1 = y2 ^ q & y1;
    }
    (a, x0, y0)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::u64;

    #[test]
    fn degree_works() {
        assert_eq!(0, degree(0));
        assert_eq!(0, degree(0b01));
        assert_eq!(1, degree(0b011));
        assert_eq!(2, degree(0b101));
        assert_eq!(63, degree(u64::MAX));
    }

    #[test]
    fn lc2_works() {
        assert_eq!(0, lc2(0));
        assert_eq!(0, lc2(0b01));
        assert_eq!(1, lc2(0b1000));
        assert_eq!(1, lc2(u64::MAX));
    }

    #[test]
    fn euclidean_division2_works() {
        assert_eq!((0b11, 0b1), euclidean_division2(0b1011, 0b110));
        assert_eq!((0b110, 0), euclidean_division2(0b110, 0b1));
        assert_eq!((0b111, 0b1), euclidean_division2(0b11010, 0b101));
        assert_eq!((0b101, 0b100), euclidean_division2(0b100011011, 0b1010011));
    }

    #[test]
    fn gcde2_works() {
        assert_eq!((0b1, 0b1, 0b1), gcde2(0b1011, 0b110));
        assert_eq!((0b110, 0, 0b1), gcde2(0b1010, 0b110));
    }

    #[test]
    fn poly2_add_works() {
        // lhs: x^8 + x^4 + x^3 + x + 1 = 0
        let p1 = Poly2(0b100011011);
        // rhs: x^8 + x^7 + x^3 + 1 = 0
        let p2 = Poly2(0b110001001);
        assert_eq!(Poly2(0b010010010), p1 + p2);
    }

    #[test]
    fn poly2_multiply_works() {
        // lhs: x^8 + x^4 + x^3 + x + 1 = 0
        let p1 = Poly2(0b100011011);
        // rhs: x^8 + x^7 + x^3 + 1 = 0
        let p2 = Poly2(0b110001001);
        assert_eq!(Poly2(0b100001001), p1 * p2);
    }

    #[test]
    fn gf2e_add_works() {
        // x^4 + x + 1 = 0
        let p = Poly2::new(0b10011);
        let g = GenGF2E::new(4, p);
        // a^9
        let a = g.gen(Poly2(0b1010));
        // a^5
        let b = g.gen(Poly2(0b0110));
        assert_eq!(g.gen(Poly2(0b1100)), a + b);
        assert_eq!(g.gen(Poly2(0b1100)), a - b);
        assert_eq!(g.gen(Poly2(0)), a + a);
        assert_eq!(g.gen(Poly2(0)), b + b);
    }

    #[test]
    fn gf2e_multiply_works() {
        // x^4 + x + 1 = 0
        let p = Poly2::new(0b10011);
        let g = GenGF2E::new(4, p);
        // a^9
        let a = g.gen(Poly2(0b1010));
        // a^5
        let b = g.gen(Poly2(0b0110));
        assert_eq!(g.gen(Poly2(0b1001)), a * b);
    }

    fn _inverse_works() {
        // x^3 + x + 1
        let p = Poly2::new(0b1011);
        let g = GenGF2E::new(3, p);
        // x
        let a = g.gen(Poly2(0b10));
        assert_eq!(Some(0b101), _inverse(a.0, p.0));
    }

    #[test]
    fn gf2_inverse_works() {
        // x^8 + x^4 + x^3 + x + 1
        let p = Poly2::new(0b100011011);
        let g = GenGF2E::new(8, p);
        // x^7
        let a = g.gen(Poly2(0b10000000));
        // x^7 + x + 1
        assert_eq!(Some(0b10000011), gf2_inverse(a.0, p.0, 8));
        // x^8 + x^4 + x^3 + x + 1

        let p = Poly2::new(0b100011011);
        let g = GenGF2E::new(8, p);
        // x^7 + x^4 + x^2 + 1
        let a = g.gen(Poly2(0b10010101));
        // x^7 + x + 1
        assert_eq!(Some(0b10001010), gf2_inverse(a.0, p.0, 8));
        // x^3 + x + 1
        let p = Poly2::new(0b1011);
        let g = GenGF2E::new(3, p);
        // x
        let a = g.gen(Poly2(0b10));
        assert_eq!(Some(0b101), gf2_inverse(a.0, p.0, 3));
    }

    #[test]
    fn gf2e_inverse_works() {
        // x^3 + x + 1
        let p = Poly2::new(0b1011);
        let g = GenGF2E::new(3, p);
        // x
        let a = g.gen(Poly2(0b10));
        assert_eq!(Some(g.gen(Poly2(0b101))), a.inverse());
        // x^4 + x + 1
        let p = Poly2::new(0b10011);
        let g = GenGF2E::new(4, p);
        // x
        let a = g.gen(Poly2(0b10));
        assert_eq!(Some(g.gen(Poly2(0b1001))), a.inverse());
    }

    #[test]
    fn gf2_inverse_primitive_works() {
        // x^3 + x + 1
        let p = Poly2::new(0b1011);
        let g = GenGF2E::new(3, p);
        // x
        let a = g.gen(Poly2(0b10));
        assert_eq!(Some(0b1000000), gf2_inverse_primitive(a.0, 3));
        // x^4 + x + 1
        let p = Poly2::new(0b10011);
        let g = GenGF2E::new(4, p);
        // x
        let a = g.gen(Poly2(0b10));
        assert_eq!(Some(1 << 14), gf2_inverse_primitive(a.0, 4));
    }
}