APK/bigint.rs

#[derive(Clone, Debug)]
pub struct BigUint {
    limbs: Vec<u32>,
}

impl BigUint {
    pub fn zero() -> Self {
        BigUint { limbs: vec![] }
    }

    pub fn from_u32(x: u32) -> Self {
        if x == 0 {
            BigUint::zero()
        } else {
            BigUint { limbs: vec![x] }
        }
    }

    pub fn from_be_bytes(bytes: &[u8]) -> Self {
        let mut b = bytes.to_vec();
        while b.len() > 1 && b[0] == 0 {
            b.remove(0);
        }
        let mut limbs = Vec::new();
        let mut i = b.len();
        while i > 0 {
            let start = i.saturating_sub(4);
            let mut limb = 0u32;
            for (k, &byte) in b[start..i].iter().enumerate() {
                let shift = (b[start..i].len() - 1 - k) * 8;
                limb |= (byte as u32) << shift;
            }
            limbs.push(limb);
            i = start;
        }
        let mut r = BigUint { limbs };
        r.normalize();
        r
    }

    pub fn to_be_bytes(&self) -> Vec<u8> {
        if self.limbs.is_empty() {
            return vec![0];
        }
        let mut out = Vec::new();
        for &limb in self.limbs.iter().rev() {
            out.extend_from_slice(&limb.to_be_bytes());
        }
        while out.len() > 1 && out[0] == 0 {
            out.remove(0);
        }
        out
    }

    pub fn to_be_bytes_padded(&self, len: usize) -> Vec<u8> {
        let mut b = self.to_be_bytes();
        if b.len() > len {
            b = b[b.len() - len..].to_vec();
        }
        while b.len() < len {
            b.insert(0, 0);
        }
        b
    }

    fn normalize(&mut self) {
        while let Some(&0) = self.limbs.last() {
            self.limbs.pop();
        }
    }

    pub fn is_zero(&self) -> bool {
        self.limbs.is_empty()
    }

    fn bit_len(&self) -> usize {
        if self.limbs.is_empty() {
            return 0;
        }
        let top = *self.limbs.last().unwrap();
        (self.limbs.len() - 1) * 32 + (32 - top.leading_zeros() as usize)
    }

    fn bit(&self, i: usize) -> bool {
        let limb = i / 32;
        let off = i % 32;
        if limb >= self.limbs.len() {
            return false;
        }
        (self.limbs[limb] >> off) & 1 == 1
    }

    fn cmp(&self, other: &BigUint) -> std::cmp::Ordering {
        use std::cmp::Ordering::*;
        if self.limbs.len() != other.limbs.len() {
            return self.limbs.len().cmp(&other.limbs.len());
        }
        for i in (0..self.limbs.len()).rev() {
            match self.limbs[i].cmp(&other.limbs[i]) {
                Equal => continue,
                ord => return ord,
            }
        }
        Equal
    }

    fn sub(&self, other: &BigUint) -> BigUint {
        let mut out = Vec::new();
        let mut borrow = 0i64;
        for i in 0..self.limbs.len() {
            let a = self.limbs[i] as i64;
            let b = *other.limbs.get(i).unwrap_or(&0) as i64;
            let mut diff = a - b - borrow;
            if diff < 0 {
                diff += 1 << 32;
                borrow = 1;
            } else {
                borrow = 0;
            }
            out.push(diff as u32);
        }
        let mut r = BigUint { limbs: out };
        r.normalize();
        r
    }

    fn shl1(&self) -> BigUint {
        let mut out = Vec::with_capacity(self.limbs.len() + 1);
        let mut carry = 0u32;
        for &limb in &self.limbs {
            let new = (limb << 1) | carry;
            carry = limb >> 31;
            out.push(new);
        }
        if carry != 0 {
            out.push(carry);
        }
        let mut r = BigUint { limbs: out };
        r.normalize();
        r
    }

    fn mul(&self, other: &BigUint) -> BigUint {
        if self.is_zero() || other.is_zero() {
            return BigUint::zero();
        }
        let mut out = vec![0u64; self.limbs.len() + other.limbs.len()];
        for (i, &a) in self.limbs.iter().enumerate() {
            let mut carry = 0u64;
            for (j, &b) in other.limbs.iter().enumerate() {
                let cur = out[i + j] + (a as u64) * (b as u64) + carry;
                out[i + j] = cur & 0xFFFF_FFFF;
                carry = cur >> 32;
            }
            out[i + other.limbs.len()] += carry;
        }
        let limbs: Vec<u32> = out.iter().map(|&x| x as u32).collect();
        let mut r = BigUint { limbs };
        r.normalize();
        r
    }

    fn rem(&self, m: &BigUint) -> BigUint {
        use std::cmp::Ordering::*;
        if m.is_zero() {
            panic!("mod by zero");
        }
        if self.cmp(m) == Less {
            return self.clone();
        }
        let mut rem = BigUint::zero();
        for i in (0..self.bit_len()).rev() {
            rem = rem.shl1();
            if self.bit(i) {
                if rem.limbs.is_empty() {
                    rem.limbs.push(1);
                } else {
                    rem.limbs[0] |= 1;
                }
            }
            if rem.cmp(m) != Less {
                rem = rem.sub(m);
            }
        }
        rem.normalize();
        rem
    }

    fn mulmod(&self, other: &BigUint, m: &BigUint) -> BigUint {
        self.mul(other).rem(m)
    }

    pub fn modpow(&self, exp: &BigUint, m: &BigUint) -> BigUint {
        if m.cmp(&BigUint::from_u32(1)) == std::cmp::Ordering::Equal {
            return BigUint::zero();
        }
        let mut result = BigUint::from_u32(1);
        let mut base = self.rem(m);
        let bits = exp.bit_len();
        for i in 0..bits {
            if exp.bit(i) {
                result = result.mulmod(&base, m);
            }
            base = base.mulmod(&base, m);
        }
        result
    }
}