cryptix-bigint 0.1.0

use core::ops::{Add, Sub, Shl, Mul, Shr, Rem, Neg};

use crate::digit::Digit;
use crate::property::IsBigInt;
use crate::BigUInt;

use super::WrapAround;
use super::slice::slice_mul_dig;

impl<T, const LEN: usize> Add for BigUInt<T, LEN> 
where
    T: Copy + Default + Digit
{
    type Output = (Self, T);

    fn add(self, rhs: Self) -> (Self, T) {
        let mut output = Self::ZERO;
        let c = output.iter_mut().zip(
            self.into_iter().zip(rhs.into_iter())
        ).fold(false, |mut c, (o, (a, b))| {
            (*o, c) = a.carry_add(b, c);
            c
        });
        (output, T::from_bool(c))
    }
}

impl<T, const LEN: usize> Add for &BigUInt<T, LEN> 
where
    T: Copy + Default + Digit
{
    type Output = (BigUInt<T, LEN>, T);

    fn add(self, rhs: Self) -> (BigUInt<T, LEN>, T) {
        let mut output = BigUInt::<T, LEN>::ZERO;
        let c = output.iter_mut().zip(
            self.into_iter().zip(rhs.into_iter())
        ).fold(false, |mut c, (o, (a, b))| {
            (*o, c) = a.carry_add(b, c);
            c
        });
        (output, T::from_bool(c))
    }
}

impl<T, const LEN: usize> Add<BigUInt<T, LEN>> for (BigUInt<T, LEN>, T)
where
    T: Copy + Default + Digit
{
    type Output = (BigUInt<T, LEN>, T);

    fn add(self, rhs: BigUInt<T, LEN>) -> (BigUInt<T, LEN>, T) {
        let (out, c) = self.0 + rhs;
        (out, self.1 + c)
    }
}

impl<T, const LEN: usize> Add<BigUInt<T, LEN>> for &(BigUInt<T, LEN>, T)
where
    T: Copy + Default + Digit
{
    type Output = (BigUInt<T, LEN>, T);

    fn add(self, rhs: BigUInt<T, LEN>) -> (BigUInt<T, LEN>, T) {
        let (out, c) = self.0 + rhs;
        (out, self.1 + c)
    }
}

impl<T, const LEN: usize> Add<T> for BigUInt<T, LEN> 
where
    T: Copy + Default + Digit
{
    type Output = (Self, T);

    // riscv optimized bigint add
    fn add(self, rhs: T) -> (Self, T) {
        let mut output = Self::ZERO;

        let c = output.iter_mut().zip(self.into_iter()).fold(rhs, |rhs, (o, a)| {
            let c;
            (*o, c) = a.overflow_add(rhs);
            T::from_bool(c)
        });
        (output, c)
    }
}

impl<T, const LEN: usize> Sub for BigUInt<T, LEN> 
where
    T: Copy + Default + Digit
{
    type Output = (Self, T);

    /// Calculate self - rhs, and an extra digit indicating if there is a borrow.
    /// If a borrows occurs, the extra digit will become -1
    fn sub(self, rhs: Self) -> (Self, T) {
        let mut output = Self::ZERO;
        let c = output.iter_mut().zip(
            self.into_iter().zip(rhs.into_iter())
        ).fold(false, |mut c, (o, (a, b))| {
            (*o, c) = a.borrow_sub(b, c);
            c
        });
        // here if a borrow occurs, the extra digit will become -1
        (output, T::ZERO.overflow_sub(T::from_bool(c)).0)
    }
}

impl<T, const LEN: usize> Sub<BigUInt<T, LEN>> for (BigUInt<T, LEN>, T)
where
    T: Copy + Default + Digit
{
    type Output = (BigUInt<T, LEN>, T);

    fn sub(self, rhs: BigUInt<T, LEN>) -> (BigUInt<T, LEN>, T) {
        let (out, c) = self.0 - rhs;
        (out, self.1.overflow_sub(c).0)
    }
}

impl<T, const LEN: usize> Sub<T> for BigUInt<T, LEN> 
where
    T: Copy + Default + Digit
{
    type Output = (Self, T);

    fn sub(self, rhs: <Self as IsBigInt>::Dig) -> (Self, T) {
        let mut output = Self::ZERO;

        let c = output.iter_mut().zip(self.into_iter()).fold(rhs, |rhs, (o, a)| {
            let c;
            (*o, c) = a.overflow_sub(rhs);
            T::from_bool(c)
        });
        (output, T::ZERO.overflow_sub(c).0)
    }
}

impl<T, const LEN: usize> Neg for BigUInt<T, LEN> 
where
    T: Digit
{
    type Output = Self;

    fn neg(self) -> Self::Output {
        (Self::ZERO - self).0
    }
}

impl<T, const LEN: usize> Mul<T> for BigUInt<T, LEN> 
where
    T: Copy + Default + Digit
{
    type Output = (BigUInt<T, LEN>, T);

    fn mul(self, rhs: T) -> (Self, T) {
        let mut output = BigUInt([T::ZERO; LEN]);
        let c = output.iter_mut().zip(self.into_iter()).fold(T::ZERO, |mut c, (o, d)| {
            (*o, c) = d.carry_mul(rhs, c);
            c
        });
        (output, c)
    }
}

impl<T, const LEN: usize> BigUInt<T, LEN> 
where
    T: Copy + Default + Digit
{
    pub fn mul_dig(self, rhs: T) -> BigUInt<T, { LEN + 1 }> {
        let mut output = BigUInt([T::ZERO; LEN + 1]);
        let c = output.iter_mut().zip(self.into_iter()).fold(T::ZERO, |mut c, (o, d)| {
            (*o, c) = d.carry_mul(rhs, c);
            c
        });
        output[LEN] = c;
        output
    }
}

impl<T, const LEN: usize> Mul for BigUInt<T, LEN> 
where 
    T: Default + Copy + Digit,
    Self: Mul<T, Output = (Self, T)>,
    [(); 2 * LEN]: ,// add this where clause to makes sure LEN * 2 can be evaluated, i.e. no overflow
    BigUInt<T, { 2 * LEN }>: Add<Output = (BigUInt<T, { 2 * LEN }>, T)>
{
    type Output = [Self; 2];

    fn mul(self, rhs: Self) -> Self::Output {
        rhs.into_iter().enumerate().fold(BigUInt::ZERO, |o, (i, dig)| {
            let mut tmp: BigUInt<T, { 2 * LEN }> = BigUInt::ZERO;
            slice_mul_dig(&mut tmp[i..i + LEN + 1], &self.0, dig);
            (o + tmp).0
        }).into()
    }
}

impl<T, const LEN: usize> Shl<usize> for BigUInt<T, LEN> 
where
    T: Default + Copy + Digit,
    Self: IsBigInt<Dig = T>
{
    type Output = Self;

    /// this is overflowing shift for big integers
    fn shl(self, rhs: usize) -> Self::Output {
        let mut output = Self::ZERO;
        let shifted_segs = rhs >> Self::DIG_BIT_LEN_SHT;
        let shifted_bits = rhs & Self::DIG_BIT_LEN_MSK;

        // overflowing shift
        if shifted_segs >= LEN { return output }

        // can shift with copying segments
        if shifted_bits == 0 {
            output[shifted_segs..].copy_from_slice(&self[..LEN - shifted_segs]);
            return output
        }
        
        output[shifted_segs] |= self[0] << shifted_bits;
        if shifted_segs + 1 >= LEN { return output }
        
        output[shifted_segs + 1..].iter_mut().zip(self.array_windows::<2>()).for_each(|(o, &[l, h])| {
            // its guaranteed that shifted_bits < Self::SEG_BIT_LEN, no panic will occur here
            *o = l >> (Self::DIG_BIT_LEN - shifted_bits);
            *o |= h << shifted_bits;
        });
        
        output
    }
}

impl<T, const LEN: usize> Shr<usize> for BigUInt<T, LEN> 
where
    T: Default + Copy + Digit,
    Self: IsBigInt<Dig = T>
{
    type Output = Self;

    /// this is overflowing shift for big integers
    fn shr(self, rhs: usize) -> Self::Output {
        let mut output = Self::ZERO;
        let shifted_segs = rhs >> Self::DIG_BIT_LEN_SHT;
        let shifted_bits = rhs & Self::DIG_BIT_LEN_MSK;

        // overflowing shift
        if shifted_segs >= LEN {
            return output;
        }

        // can shift with copying segments
        if shifted_bits == 0 {
            output[..LEN - shifted_segs].copy_from_slice(&self[shifted_segs..]);
            return output
        }

        output.iter_mut().zip(self[shifted_segs..].array_windows::<2>()).for_each(|(o, &[l, h])| {
            // its guaranteed that shifted_bits < Self::SEG_BIT_LEN, no panic will occur here
            *o = l >> shifted_bits;
            *o |= h << (Self::DIG_BIT_LEN - shifted_bits);
        });
        output[LEN - shifted_segs - 1] |= self[LEN - 1] >> shifted_bits;
        
        output
    }
}

impl<T, const LEN: usize> Rem for BigUInt<T, LEN> 
where
    T: Default + Copy + Digit,
    Self: IsBigInt<Dig = T> + Ord,
    Self: Sub<Output = (Self, T)>
{
    type Output = Self;

    // compute a % b in time O(len(a) - len(b))
    fn rem(self, rhs: Self) -> Self::Output {
        if self < rhs {
            return self
        }

        let mut a = self;
        let b = rhs;

        // we want to compute a % b
        // 1. get bit length of a and b, where len(a) >= len(b)
        let mut a_len = a.bit_len();
        let b_len = b.bit_len();

        if a_len == 0 {
            panic!("mod by zero.")
        }

        // 2.1. len(a) == len(b), then len(a - b) < len(b) thus a % b = a - b
        while a_len > b_len {
            // 2.2. now len(a) > len(b), we left shift b to make len(a) - 1 = len(b)
            let b_shift_amt = a_len - b_len - 1;
            let tmp = b << b_shift_amt;

            // 3. a = a - b
            a = (a - tmp).0;

            // 4. if a > b still holds, substract again to make sure a < b
            if a.bit(a_len - 1) { a = (a - tmp).0 }
            // FIXME: may be there are some optimizations to avoid calculate bit length again
            a_len = a.bit_len();
        }

        if a < b { a } else { (a - b).0 }
    }
}

impl<T, const LEN: usize> WrapAround<[BigUInt<T, LEN>; 2]> for BigUInt<T, LEN>
where
    BigUInt<T, { 2 * LEN }>: Rem<Output = BigUInt<T, { 2 * LEN }>>,
    T: Default + Copy + Digit,
{
    // compute rhs % self in time O(len(a) - len(b))
    fn wrap(self, rhs: [BigUInt<T, LEN>; 2]) -> Self {
        let a: BigUInt<T, { 2 * LEN }> = self.resize();
        let b: BigUInt<T, { 2 * LEN }> = rhs.into();

        (b % a).resize()
    }
}

impl<T, const LEN: usize> WrapAround<BigUInt<T, LEN>> for BigUInt<T, LEN>
where
    T: Default + Copy + Digit,
{
    // compute rhs % self in time O(len(a) - len(b))
    fn wrap(self, rhs: BigUInt<T, LEN>) -> Self {
        rhs % self
    }
}

impl<T: Copy + PartialEq, const LEN: usize> PartialEq for BigUInt<T, LEN> {
    fn eq(&self, other: &Self) -> bool {
        self.0.eq(&other.0)
    }
}

impl<T: Copy + PartialEq, const LEN: usize> PartialEq<[BigUInt<T, LEN>; 2]> for BigUInt<T, LEN> 
where
    BigUInt<T, LEN>: IsBigInt<Dig = T>
{
    fn eq(&self, other: &[BigUInt<T, LEN>; 2]) -> bool {
        other[1].is_zero() && self == &other[0]
    }
}

impl<T, const LEN: usize> Eq for BigUInt<T, LEN> where BigUInt<T, LEN>: PartialEq {}

impl<T: Copy + PartialOrd, const LEN: usize> PartialOrd for BigUInt<T, LEN> {
    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
        self.iter().rev().partial_cmp(other.iter().rev())
    }
}

impl<T: Copy + PartialOrd, const LEN: usize> PartialOrd<[BigUInt<T, LEN>; 2]> for BigUInt<T, LEN> 
where
    BigUInt<T, LEN>: IsBigInt<Dig = T>
{
    fn partial_cmp(&self, other: &[BigUInt<T, LEN>; 2]) -> Option<core::cmp::Ordering> {
        if !other[1].is_zero() {
            return Some(core::cmp::Ordering::Less)
        }
        self.iter().rev().partial_cmp(other[0].iter().rev())
    }
}

impl<T: Copy + Ord, const LEN: usize> Ord for BigUInt<T, LEN> {
    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
        self.iter().rev().cmp(other.iter().rev())
    }
}