relp-num 0.2.1

use std::cmp::{min, Ordering};
use std::num::NonZeroUsize;
use std::ptr;

use smallvec::SmallVec;

use crate::integer::big::BITS_PER_WORD;
use crate::integer::big::ops::building_blocks::{add_assign_slice, addmul_1, borrowing_sub_mut, carrying_add_mut, is_well_formed, is_well_formed_non_zero, mul_1, sub_assign_slice, sub_from_slice, sub_n, to_twos_complement};
use crate::integer::big::properties::cmp;

#[inline]
pub fn shr_mut<const S: usize>(values: &mut SmallVec<[usize; S]>, words: usize, bits: u32) {
    debug_assert_shr_constraints(values, words, bits);

    let original_number_words = values.len();

    if bits > 0 {
        for i in 0..(original_number_words - words - 1) {
            values[i] = values[i + words] >> bits;
            values[i] |= values[i + words + 1] << (BITS_PER_WORD - bits);
        }

        let last_shifted = values.last().unwrap() >> bits;
        let words_to_keep = if last_shifted > 0 {
            values[original_number_words - words - 1] = last_shifted;
            original_number_words - words
        } else {
            original_number_words - words - 1
        };

        values.truncate(words_to_keep);
    } else {
        let remaining_words = original_number_words - words;
        unsafe {
            // Derive both pointers from the same mutable borrow: evaluating `as_mut_ptr` would
            // invalidate a pointer taken from a shared borrow of the same values.
            let pointer = values.as_mut_ptr();
            ptr::copy(pointer.add(words), pointer, remaining_words);
        }
        values.truncate(remaining_words);
    }

    debug_assert!(is_well_formed(values));
}

#[inline]
pub fn shr<const S: usize>(values: &[usize], words_to_remove: usize, bits: u32) -> SmallVec<[usize; S]> {
    debug_assert_shr_constraints(values, words_to_remove, bits);

    if bits > 0 {
        let last_shifted = values.last().unwrap() >> bits;
        let words_to_keep = if last_shifted > 0 {
            values.len() - words_to_remove
        } else {
            values.len() - words_to_remove - 1
        };

        let mut result = SmallVec::with_capacity(words_to_keep);
        for i in 0..(values.len() - words_to_remove - 1) {
            let mut result_word = values[i + words_to_remove] >> bits;
            result_word |= values[i + words_to_remove + 1] << (BITS_PER_WORD - bits);
            result.push(result_word);
        }

        if last_shifted > 0 {
            result.push(last_shifted);
        }

        result
    } else {
        SmallVec::from_slice(&values[words_to_remove..])
    }
}

#[inline]
pub fn shl_mut<const S: usize>(values: &mut SmallVec<[usize; S]>, words: usize, bits: u32) {
    debug_assert!(is_well_formed_non_zero(values));

    if bits > 0 {
        let original_length = values.len();

        let overflows = bits > values.last().unwrap().leading_zeros();
        if overflows {
            // These values will be overwritten
            values.extend(std::iter::repeat_n(0, words + 1));
            *values.last_mut().unwrap() = values[original_length - 1] >> (BITS_PER_WORD - bits);
        } else {
            // These values will be overwritten
            values.extend(std::iter::repeat_n(0, words));
        }

        for i in (1..original_length).rev() {
            values[words + i] = values[i] << bits;
            values[words + i] |= values[i - 1] >> (BITS_PER_WORD - bits);
        }
        values[words] = values[0] << bits;
    } else {
        values.reserve(words);
        let old_length = values.len();
        unsafe {
            // The words in `words..(words + old_length)` are uninitialized until the copy below
            // fills them; `SmallVec::as_mut_ptr` does not build an intermediate slice over them.
            // Both pointers are derived from the same mutable borrow, as evaluating `as_mut_ptr`
            // would invalidate a pointer taken from a shared borrow of the same values.
            values.set_len(old_length + words);
            let pointer = values.as_mut_ptr();
            ptr::copy(pointer, pointer.add(words), old_length);
        }
    }

    values[..words].fill(0);

    debug_assert!(is_well_formed_non_zero(values));
}

/// Shift left within the same number of words, returning the words that shift out of the top.
///
/// # Panics
///
/// In debug mode, if `bits` is zero or not smaller than the number of bits in a word.
#[inline]
pub fn shl_mut_overflowing<const S: usize>(values: &mut SmallVec<[usize; S]>, bits: u32) -> Option<NonZeroUsize> {
    debug_assert_ne!(bits, 0);
    debug_assert!(bits < BITS_PER_WORD);

    let original_number_words = values.len();

    let value_highest = values.last().unwrap();
    let leading_zeros = value_highest.leading_zeros();
    let carry = if bits > leading_zeros {
        Some(NonZeroUsize::new(value_highest >> (BITS_PER_WORD - bits)).unwrap())
    } else {
        None
    };

    for i in (1..original_number_words).rev() {
        values[i] <<= bits;
        // `bits` is not zero, so this shift is not by the full width of a word.
        values[i] |= values[i - 1] >> (BITS_PER_WORD - bits);
    }

    values[0] <<= bits;

    carry
}

#[inline]
pub fn shl<const S: usize>(values: &[usize], bits: u32) -> SmallVec<[usize; S]> {
    debug_assert_ne!(bits, 0);
    debug_assert!(bits < BITS_PER_WORD);
    debug_assert!(bits <= values.last().unwrap().leading_zeros());

    let mut result = SmallVec::with_capacity(values.len());
    result.push(values[0] << bits);
    for i in 1..values.len() {
        let mut result_value = values[i] << bits;
        if bits > 0 {
            result_value |= values[i - 1] >> (BITS_PER_WORD - bits);
        }
        result.push(result_value);
    }

    result
}

#[inline]
pub fn debug_assert_shr_constraints(values: &[usize], words_to_remove: usize, bits: u32) {
    debug_assert!(is_well_formed_non_zero(values));
    debug_assert!(!values.is_empty(), "Should not be called on a zero value");
    debug_assert!(words_to_remove < values.len(), "Can't shift away all words");
    debug_assert!(bits < BITS_PER_WORD, "Use the `words` argument to shift with entire words");
}

#[inline]
pub fn add_assign<const S: usize>(values: &mut SmallVec<[usize; S]>, rhs: &[usize]) {
    debug_assert!(is_well_formed(values));
    debug_assert!(is_well_formed(rhs));

    let shared = min(values.len(), rhs.len());
    let mut carry = add_assign_slice(&mut values[..shared], &rhs[..shared]);
    let mut i = shared;

    while i < rhs.len() {
        let (new_value, new_carry) = rhs[i].overflowing_add(carry as usize);
        values.push(new_value);
        carry = new_carry;
        i += 1;
    }

    while carry && i < values.len() {
        let (new_value, new_carry) = values[i].overflowing_add(carry as usize);
        values[i] = new_value;
        carry = new_carry;
        i += 1;
    }

    if carry {
        values.push(1);
    }
}

#[inline]
pub fn add_assign_single_non_zero<const S: usize>(
    values: &mut SmallVec<[usize; S]>,
    rhs: usize,
) {
    let mut carry = false;
    carrying_add_mut(&mut values[0], rhs, &mut carry);
    let mut i = 1;
    while carry && i < values.len() {
        carrying_add_mut(&mut values[i], 0, &mut carry);
        i += 1;
    }
    if carry {
        values.push(1);
    }
}

#[inline]
pub(crate) fn mul_assign_single_non_zero<const S: usize>(
    values: &mut SmallVec<[usize; S]>, rhs: usize,
) {
    debug_assert!(!values.is_empty());
    // Multiplying by zero would leave a denormalized `[0, ..]` behind.
    debug_assert_ne!(rhs, 0);

    let (low, mut previous_high) = values[0].carrying_mul(rhs, 0);
    values[0] = low;
    let mut carry = false;
    for i in 1..values.len() {
        let (low, high) = values[i].carrying_mul(rhs, 0);
        let (value_new, carry_new) = previous_high.carrying_add(low, carry);
        values[i] = value_new;
        carry = carry_new;
        previous_high = high;
    }

    // The high word of a product of two words is at most `usize::MAX - 1`, so adding the carry
    // can't overflow and there is never a second word to push.
    let value_new = previous_high + carry as usize;
    if value_new > 0 {
        values.push(value_new);
    }
}

#[must_use]
#[inline]
pub unsafe fn mul_non_zero<const S: usize>(values: &[usize], rhs: &[usize]) -> SmallVec<[usize; S]> {
    debug_assert!(is_well_formed(values));
    debug_assert!(is_well_formed(rhs));
    debug_assert!(!values.is_empty());
    debug_assert!(!rhs.is_empty());

    let mut result = SmallVec::with_capacity(values.len() + rhs.len());

    let (large, small) = if rhs.len() > values.len() {
        (rhs, values)
    } else {
        (values, rhs)
    };

    unsafe {
        // The capacity reserved above covers these words. Setting the length before the
        // multiplication is what gives the slice handed to `mul_1` provenance over the words it
        // writes; every one of them is written there before it is read, and `usize` has no drop
        // glue, so no value is dropped that was never initialized.
        result.set_len(large.len());
    }
    let carry = mul_1(&mut result, large, small[0]);
    result.push(carry);

    for i in 1..small.len() {
        // `result` is `large.len() + i` words long here, so this slice is exactly `large.len()`
        // words, which is what `addmul_1` writes.
        let carry = addmul_1(&mut result[i..], large, small[i]);
        result.push(carry);
    }

    if *result.last().unwrap() == 0 {
        result.pop();
    }

    debug_assert!(is_well_formed(&result));

    result
}

#[inline]
pub unsafe fn sub<const S: usize>(
    values: &SmallVec<[usize; S]>,
    rhs: &SmallVec<[usize; S]>,
) -> SmallVec<[usize; S]> {
    debug_assert!(is_well_formed(values));
    debug_assert!(is_well_formed(rhs));
    debug_assert_eq!(cmp(values, rhs), Ordering::Greater);

    let mut result = SmallVec::with_capacity(values.len());
    // Will be overwritten in the unsafe block, but this is safe and extends the length
    result.extend(std::iter::repeat_n(0, rhs.len()));

    let mut carry = sub_n(&mut result, values, rhs, rhs.len()) > 0;

    while carry {
        debug_assert!(values.len() > rhs.len());
        let (value, new_carry) = values[result.len()].borrowing_sub(0, carry);
        carry = new_carry;
        result.push(value);
    }

    if result.len() < values.len() {
        result.extend_from_slice(&values[result.len()..]);
    } else {
        while let Some(0) = result.last() {
            result.pop();
        }
    }

    result
}

#[inline]
pub(crate) fn subtracting_cmp<const S: usize>(left: &mut SmallVec<[usize; S]>, right: &[usize]) -> Ordering {
    debug_assert!(is_well_formed_non_zero(left));
    debug_assert!(is_well_formed_non_zero(right));

    match left.len().cmp(&right.len()) {
        Ordering::Less => {
            let mut i = left.len();
            let mut carry = sub_from_slice(left, &right[..i]);

            while carry {
                let (new_value, new_carry) = right[i].borrowing_sub(0, true);
                left.push(new_value);
                carry = new_carry;
                i += 1;
            }

            left.extend_from_slice(&right[i..]);

            while *left.last().unwrap() == 0 {
                left.pop();
            }

            Ordering::Less
        }
        Ordering::Equal => {
            let carry = sub_assign_slice(left, right);

            if carry {
                // result is negative; `to_twos_complement` also normalizes
                to_twos_complement(left);
                Ordering::Less
            } else {
                // result is zero or positive

                while let Some(0) = left.last() {
                    left.pop();
                }

                // TODO(PERFORMANCE): This test is often not necessary, as it is known in advance
                //  whether the result can be zero or not.
                if left.is_empty() {
                    Ordering::Equal
                } else {
                    Ordering::Greater
                }
            }
        }
        Ordering::Greater => {
            let mut carry = sub_assign_slice(&mut left[..right.len()], right);
            let mut i = 0;
            while carry && i + right.len() < left.len() {
                borrowing_sub_mut(&mut left[i + right.len()], 0, &mut carry);
                i += 1;
            }
            debug_assert!(!carry);
            while *left.last().unwrap() == 0 {
                left.pop();
            }

            Ordering::Greater
        }
    }
}

#[inline]
pub(crate) fn subtracting_cmp_ne_single<const S: usize>(left: &mut SmallVec<[usize; S]>, right: usize) -> Ordering {
    debug_assert!(is_well_formed_non_zero(left));
    debug_assert_ne!(right, 0);
    debug_assert!(left.len() > 1 || left[0] != right);

    if left.len() > 1 {
        // result won't be negative
        let mut carry = false;
        borrowing_sub_mut(&mut left[0], right, &mut carry);

        let mut i = 1;
        while carry {
            borrowing_sub_mut(&mut left[i], 0, &mut carry);
            i += 1;
        }

        while let Some(0) = left.last() {
            left.pop();
        }

        Ordering::Greater
    } else {
        // result might be negative
        if left[0] < right {
            left[0] = right - left[0];
            Ordering::Less
        } else {
            // left[0] > right
            left[0] -= right;
            Ordering::Greater
        }
    }
}

/// # Safety
///
/// `values` must not be empty.
#[must_use]
#[inline]
pub unsafe fn is_one_non_zero(values: &[usize]) -> bool {
    debug_assert!(is_well_formed(values));
    debug_assert!(!values.is_empty());

    // SAFETY: The caller guarantees that there is at least one word.
    unsafe { *values.get_unchecked(0) == 1 && values.len() == 1 }
}

/// # Safety
///
/// Neither `left` nor `right` may be empty.
#[inline]
pub unsafe fn both_not_one_non_zero(left: &[usize], right: &[usize]) -> bool {
    debug_assert!(!left.is_empty());
    debug_assert!(!right.is_empty());

    // SAFETY: The caller guarantees that both have at least one word.
    unsafe {
        (*left.get_unchecked(0) != 1 || left.len() > 1) && (*right.get_unchecked(0) != 1 || right.len() > 1)
    }
}

#[cfg(test)]
mod test {
    use std::mem;
    use std::num::NonZeroUsize;

    use smallvec::{smallvec, SmallVec};

    use crate::integer::big::BITS_PER_WORD;
    use crate::integer::big::io::from_str_radix;
    use crate::integer::big::ops::non_zero::{both_not_one_non_zero, shl_mut, shl_mut_overflowing, shr, shr_mut};

    #[test]
    fn test_shr() {
        type SV = SmallVec<[usize; 4]>;

        let x: SV = smallvec![1];
        let expected: SV = smallvec![1];
        assert_eq!(shr::<4>(&x, 0, 0), expected);

        let x: SV = smallvec![0, 1];
        let expected: SV = smallvec![1];
        assert_eq!(shr::<4>(&x, 1, 0), expected);

        let x: SV = smallvec![0, 1];
        let expected: SV = smallvec![1 << (mem::size_of::<usize>() * 8 - 1)];
        assert_eq!(shr::<4>(&x, 0, 1), expected);

        let x: SV = smallvec![0, 0, 0, 1];
        let expected: SV = smallvec![0, 1 << (mem::size_of::<usize>() * 8 - 1)];
        assert_eq!(shr::<4>(&x, 1, 1), expected);

        let x: SV = smallvec![0, 1];
        let expected: SV = smallvec![2];
        assert_eq!(shr::<4>(&x, 0, (mem::size_of::<usize>() * 8 - 1) as u32), expected);

        // 15588921427233345156 // 4
        let x: SV = smallvec![15588921427233345156];
        let expected: SV = smallvec![3897230356808336289];
        assert_eq!(shr::<4>(&x, 0, 2), expected);

        // 15588921427233345156 // 4
        let x: SV = smallvec![131522304505784511, 2433];
        assert_eq!(shr::<4>(&x, 0, 0), x);
    }

    #[test]
    fn test_shr_mut() {
        type SV = SmallVec<[usize; 4]>;

        let mut x: SV = smallvec![1];
        shr_mut(&mut x, 0, 0);
        let expected: SV = smallvec![1];
        assert_eq!(x, expected);

        let mut x: SV = smallvec![1];
        shr_mut(&mut x, 0, 1);
        let expected: SV = smallvec![];
        assert_eq!(x, expected);

        let mut x: SV = smallvec![0, 1];
        shr_mut(&mut x, 1, 1);
        let expected: SV = smallvec![];
        assert_eq!(x, expected);

        let mut x: SV = smallvec![6184, 1];
        shr_mut(&mut x, 1, 1);
        let expected: SV = smallvec![];
        assert_eq!(x, expected);

        let mut x: SV = smallvec![0, 1];
        shr_mut(&mut x, 1, 0);
        let expected: SV = smallvec![1];
        assert_eq!(x, expected);

        let mut x: SV = smallvec![0, 1];
        shr_mut(&mut x, 0, 1);
        let expected: SV = smallvec![1 << (mem::size_of::<usize>() * 8 - 1)];
        assert_eq!(x, expected);

        let mut x: SV = smallvec![0, 0, 0, 1];
        shr_mut(&mut x, 1, 1);
        let expected: SV = smallvec![0, 1 << (mem::size_of::<usize>() * 8 - 1)];
        assert_eq!(x, expected);

        let mut x: SV = smallvec![0, 1];
        shr_mut(&mut x, 0, (mem::size_of::<usize>() * 8 - 1) as u32);
        let expected: SV = smallvec![2];
        assert_eq!(x, expected);

        let mut x: SV = smallvec![0, 1 << (80 - 64)];
        shr_mut(&mut x, 0, 0);
        let expected: SV = smallvec![0, 1 << (80 - 64)];
        assert_eq!(x, expected);

        let mut x: SV = smallvec![0, 12511854210725346487, 1932217123071064976, 10302437120704275430, 18852552];
        shr_mut(&mut x, 1, 0);
        let expected: SV = smallvec![12511854210725346487, 1932217123071064976, 10302437120704275430, 18852552];
        assert_eq!(x, expected);
    }

    #[test]
    fn test_shl_mut_overflowing() {
        type SV = SmallVec<[usize; 4]>;

        let mut x: SV = smallvec![0, 1];
        let carry = shl_mut_overflowing(&mut x, 1);
        let expected: SV = smallvec![0, 2];
        assert_eq!(x, expected);
        assert_eq!(carry, None);

        let mut x: SV = smallvec![0, 0, 0, 1];
        let carry = shl_mut_overflowing(&mut x, 1);
        let expected: SV = smallvec![0, 0, 0, 2];
        assert_eq!(x, expected);
        assert_eq!(carry, None);

        let mut x: SV = smallvec![0, 0, 2_usize.pow((mem::size_of::<usize>() * 8 - 2) as u32)];
        let carry = shl_mut_overflowing(&mut x, 1);
        let expected: SV = smallvec![0, 0, 2_usize.pow((mem::size_of::<usize>() * 8 - 1) as u32)];
        assert_eq!(x, expected);
        assert_eq!(carry, None);

        let mut x: SV = smallvec![0, 0, 2_usize.pow((mem::size_of::<usize>() * 8 - 1) as u32)];
        let carry = shl_mut_overflowing(&mut x, 1);
        let expected: SV = smallvec![0, 0, 0];
        assert_eq!(x, expected);
        assert_eq!(carry, Some(NonZeroUsize::new(1).unwrap()));
    }

    #[test]
    fn test_shl_mut() {
        type SV = SmallVec<[usize; 4]>;

        let mut x: SV = smallvec![0, 1];
        shl_mut(&mut x, 0, 1);
        let expected: SV = smallvec![0, 2];
        assert_eq!(x, expected);

        let mut x: SV = smallvec![0, 0, 0, 1];
        shl_mut(&mut x, 2, 1);
        let expected: SV = smallvec![0, 0, 0, 0, 0, 2];
        assert_eq!(x, expected);

        let mut x: SV = smallvec![0, 0, 0, 1];
        shl_mut(&mut x, 2, 0);
        let expected: SV = smallvec![0, 0, 0, 0, 0, 1];
        assert_eq!(x, expected);

        let mut x: SV = smallvec![0, 0, 0, 1];
        shl_mut(&mut x, 0, 0);
        let expected: SV = smallvec![0, 0, 0, 1];
        assert_eq!(x, expected);

        let mut x: SV = smallvec![0, 0, 2_usize.pow(BITS_PER_WORD - 2)];
        shl_mut(&mut x, 1,1);
        let expected: SV = smallvec![0, 0, 0, 2_usize.pow(BITS_PER_WORD - 1)];
        assert_eq!(x, expected);

        let mut x: SV = smallvec![0, 0, 2_usize.pow(BITS_PER_WORD - 1)];
        shl_mut(&mut x, 2,1);
        let expected: SV = smallvec![0, 0, 0, 0, 0, 1];
        assert_eq!(x, expected);

        let mut x: SV = smallvec![0, 1 << (80 - 64)];
        shl_mut(&mut x, 0, 0);
        let expected: SV = smallvec![0, 1 << (80 - 64)];
        assert_eq!(x, expected);

        let mut x: SV = smallvec![8128280416544428961, 8274854478500297142, 3587444469682918757];
        shl_mut(&mut x, 1, 1);
        let expected: SV = from_str_radix::<10, 4>("45037507812500000000000000000000000000000000000000000000000000000000000000000").unwrap();
        assert_eq!(x, expected);
    }

    #[test]
    fn test_both_one() {
        unsafe {
            assert!(!both_not_one_non_zero(&[1], &[1]));
            assert!(!both_not_one_non_zero(&[1], &[2]));
            assert!(!both_not_one_non_zero(&[2], &[1]));
            assert!(!both_not_one_non_zero(&[1], &[1, 2]));
            assert!(!both_not_one_non_zero(&[1, 2], &[1]));
            assert!(both_not_one_non_zero(&[6], &[7]));
            assert!(both_not_one_non_zero(&[6], &[1, 2]));
            assert!(both_not_one_non_zero(&[1, 6], &[3]));
            assert!(both_not_one_non_zero(&[6, 2], &[7, 2]));
            assert!(both_not_one_non_zero(&[1, 2], &[7, 2]));
            assert!(both_not_one_non_zero(&[2, 2], &[1, 2]));
            assert!(both_not_one_non_zero(&[1, 2], &[1, 2]));
        }
    }
}