smallint/src/ops.rs

use crate::smallint::{SmallIntType, SmallUintType};
use crate::{SmallInt, SmallUint};
use core::mem::ManuallyDrop;
use core::ops::{Add, Div, Mul, Neg, Sub};
use core::ops::{AddAssign, DivAssign, MulAssign, SubAssign};

impl Neg for SmallInt {
    type Output = Self;

    fn neg(self) -> Self::Output {
        match self.0 {
            SmallIntType::Inline(i) => {
                if let Some(n) = i.checked_neg() {
                    SmallInt(SmallIntType::Inline(n))
                } else {
                    // -i128::MIN = i128::MAX + 1 = 0x8000...
                    let mut val = ManuallyDrop::new(Box::new([0, 0, 0, 0x80_00_00_00]));
                    SmallInt(SmallIntType::Heap((val.as_mut_ptr(), 4)))
                }
            },
            SmallIntType::Heap((r, s)) => {
                let size = s.unsigned_abs();
                let slice = unsafe { core::slice::from_raw_parts(r, size) };
                let mut ret = vec![0; size];
                ret.clone_from_slice(slice);
                let mut val = ManuallyDrop::new(ret.into_boxed_slice());
                SmallInt(SmallIntType::Heap((val.as_mut_ptr(), -s)))
            }
        }
    }
}

macro_rules! basic_op {
    ($imp:ident, $lower:ident, $typ:ty, $fun:ident) => {
        impl<'a, 'b> $imp<&'a $typ> for &'b $typ {
            type Output = $typ;

            fn $lower(self, rhs: &$typ) -> Self::Output {
                $fun(self, rhs)
            }
        }

        impl<'a> $imp<$typ> for &'a $typ {
            type Output = $typ;

            fn $lower(self, rhs: $typ) -> Self::Output {
                self.$lower(&rhs)
            }
        }

        impl<'a> $imp<&'a $typ> for $typ {
            type Output = $typ;

            fn $lower(self, rhs: &$typ) -> Self::Output {
                (&self).$lower(rhs)
            }
        }

        impl $imp<$typ> for $typ {
            type Output = $typ;

            fn $lower(self, rhs: $typ) -> Self::Output {
                (&self).$lower(&rhs)
            }
        }
    };
}

pub(crate) fn add_two_slices(slice1: &[u32], slice2: &[u32]) -> Vec<u32> {
    let s = slice1.len();
    let j = slice2.len();

    let larger = std::cmp::max(s, j);
    let mut res = Vec::with_capacity(larger + 1);
    let mut carry = false;

    for t in 0..larger {
        let value1 = if t < s { slice1[t] } else { 0 };

        let value2 = if t < j { slice2[t] } else { 0 };

        let (val, overflow) = value1.overflowing_add(value2);
        let (cval, coverflow) = val.overflowing_add(carry as u32);
        res.push(cval);
        carry = overflow | coverflow;
    }

    if carry {
        res.push(1);
    }
    while res.len() != 1 && res[res.len() - 1] == 0 {
        res.pop();
    }
    res
}

fn add(a: &SmallUint, b: &SmallUint) -> SmallUint {
    match (&a.0, &b.0) {
        (&SmallUintType::Inline(i), &SmallUintType::Inline(j)) => match i.overflowing_add(j) {
            (t, false) => SmallUint(SmallUintType::Inline(t)),
            (t, true) => {
                let mut res = [0, 0, 0, 0, 1];

                let mut v = t;
                #[allow(clippy::needless_range_loop)]
                for r in 0..4 {
                    res[r] = v as u32;

                    v >>= 32;
                }

                let mut slice = ManuallyDrop::new(<Box<[u32]>>::from(res));

                SmallUint(SmallUintType::Heap((slice.as_mut_ptr(), 5)))
            }
        },
        (&SmallUintType::Heap((r, s)), &SmallUintType::Inline(i))
        | (&SmallUintType::Inline(i), &SmallUintType::Heap((r, s))) => {
            let slice1 = unsafe { core::slice::from_raw_parts(r, s) };

            let mut res = [0, 0, 0, 0];

            let mut v = i;
            #[allow(clippy::needless_range_loop)]
            for r in 0..4 {
                res[r] = v as u32;

                v >>= 32;
            }

            let result = add_two_slices(slice1, &res[..]);
            let size = result.len();

            let mut slice = ManuallyDrop::new(result.into_boxed_slice());

            SmallUint(SmallUintType::Heap((slice.as_mut_ptr(), size)))
        }
        (&SmallUintType::Heap((r, s)), &SmallUintType::Heap((i, j))) => {
            let slice1 = unsafe { core::slice::from_raw_parts(r, s) };
            let slice2 = unsafe { core::slice::from_raw_parts(i, j) };

            let res = add_two_slices(slice1, slice2);
            let size = res.len();

            let mut slice = ManuallyDrop::new(res.into_boxed_slice());

            SmallUint(SmallUintType::Heap((slice.as_mut_ptr(), size)))
        }
    }
}

basic_op!(Add, add, SmallUint, add);

impl<'a> AddAssign<&'a SmallUint> for SmallUint {
    fn add_assign(&mut self, rhs: &'a SmallUint) {
        *self = self.clone() + rhs;
    }
}

impl AddAssign<SmallUint> for SmallUint {
    fn add_assign(&mut self, rhs: SmallUint) {
        *self = self.clone() + rhs;
    }
}

fn add_signed(a: &SmallInt, b: &SmallInt) -> SmallInt {
    let a_sign;
    match &a.0 {
        SmallIntType::Inline(i) => a_sign = i.signum() as i8,
        SmallIntType::Heap((_, s)) => a_sign = s.signum() as i8,
    }

    let b_sign;
    match &b.0 {
        SmallIntType::Inline(i) => b_sign = i.signum() as i8,
        SmallIntType::Heap((_, s)) => b_sign = s.signum() as i8,
    }

    match (a_sign, b_sign) {
        x if (x.0 >= 0 && x.1 >= 0) => SmallInt::from(
            SmallUint::from_smallint_unsigned(a.clone())
                + SmallUint::from_smallint_unsigned(b.clone()),
        ),
        x if (x.0 < 0 && x.1 < 0) => -SmallInt::from(
            SmallUint::from_smallint_unsigned(a.clone())
                + SmallUint::from_smallint_unsigned(b.clone()),
        ),

        x if (x.0 >= 0 && x.1 < 0) => {
            let s = SmallUint::from_smallint_unsigned(a.clone());
            let b = SmallUint::from_smallint_unsigned(b.clone());
            if b <= s {
                SmallInt::from(s - b)
            } else {
                -SmallInt::from(b - s)
            }
        }

        x if (x.0 < 0 && x.1 >= 0) => {
            let s = SmallUint::from_smallint_unsigned(a.clone());
            let b = SmallUint::from_smallint_unsigned(b.clone());
            if s <= b {
                SmallInt::from(b - s)
            } else {
                -SmallInt::from(s - b)
            }
        }
        (_, _) => {
            panic!("This shouldn't happen. ");
        }
    }
}

basic_op!(Add, add, SmallInt, add_signed);

impl<'a> AddAssign<&'a SmallInt> for SmallInt {
    fn add_assign(&mut self, rhs: &'a SmallInt) {
        *self = self.clone() + rhs;
    }
}

impl AddAssign<SmallInt> for SmallInt {
    fn add_assign(&mut self, rhs: SmallInt) {
        *self = self.clone() + rhs;
    }
}

pub(crate) fn sub_two_slices(slice1: &[u32], slice2: &[u32]) -> Vec<u32> {
    let b = slice1.len();
    let s = slice2.len();

    if b < s {
        panic!("First number is smaller than second.");
    }

    let mut res = Vec::with_capacity(std::cmp::max(s, b));
    let mut borrow = false;

    for i in 0..b {
        let mut value1 = slice1[i];

        let value2 = if i < s { slice2[i] } else { 0 };

        if borrow {
            let (temp, b) = value1.overflowing_sub(1);
            value1 = temp;
            borrow = b;
        }

        if value2 > value1 {
            borrow = true;
        }

        let val = value1.wrapping_sub(value2);
        res.push(val);
    }

    if borrow {
        panic!("First number is smaller than second. ");
    }

    res
}

fn sub(a: &SmallUint, b: &SmallUint) -> SmallUint {
    match (&a.0, &b.0) {
        (&SmallUintType::Inline(i), &SmallUintType::Inline(j)) => {
            if let (t, false) = i.overflowing_sub(j) {
                SmallUint(SmallUintType::Inline(t))
            } else {
                panic!("First number is smaller than second. ");
            }
        }
        (&SmallUintType::Heap((r, s)), &SmallUintType::Inline(i)) => {
            let slice1 = unsafe { core::slice::from_raw_parts(r, s) };

            let mut res = [0, 0, 0, 0];

            let mut v = i;
            #[allow(clippy::needless_range_loop)]
            for r in 0..4 {
                res[r] = v as u32;

                v >>= 32;
            }

            let result = sub_two_slices(slice1, &res[..]);
            let size = result.len();

            let mut slice = ManuallyDrop::new(result.into_boxed_slice());

            SmallUint(SmallUintType::Heap((slice.as_mut_ptr(), size)))
        }
        (&SmallUintType::Inline(_), &SmallUintType::Heap((_, _))) => {
            panic!("First number is smaller than second. ");
        }
        (&SmallUintType::Heap((r, s)), &SmallUintType::Heap((i, j))) => {
            let slice1 = unsafe { core::slice::from_raw_parts(r, s) };
            let slice2 = unsafe { core::slice::from_raw_parts(i, j) };

            let res = sub_two_slices(slice1, slice2);
            let size = res.len();

            let mut slice = ManuallyDrop::new(res.into_boxed_slice());

            SmallUint(SmallUintType::Heap((slice.as_mut_ptr(), size)))
        }
    }
}

basic_op!(Sub, sub, SmallUint, sub);

impl<'a> SubAssign<&'a SmallUint> for SmallUint {
    fn sub_assign(&mut self, rhs: &'a SmallUint) {
        *self = self.clone() - rhs;
    }
}

impl SubAssign<SmallUint> for SmallUint {
    fn sub_assign(&mut self, rhs: SmallUint) {
        *self = self.clone() - rhs;
    }
}

fn sub_signed(a: &SmallInt, b: &SmallInt) -> SmallInt {
    a + (-b.clone())
}

basic_op!(Sub, sub, SmallInt, sub_signed);

impl<'a> SubAssign<&'a SmallInt> for SmallInt {
    fn sub_assign(&mut self, rhs: &'a SmallInt) {
        *self = self.clone() - rhs;
    }
}

impl SubAssign<SmallInt> for SmallInt {
    fn sub_assign(&mut self, rhs: SmallInt) {
        *self = self.clone() - rhs;
    }
}

// Taken from https://github.com/rust-lang/rust/issues/85532#issuecomment-916309635. Credit to
// AaronKutch.
const fn carrying_mul_u128(lhs: u128, rhs: u128, carry: u128) -> (u128, u128) {
    //                       [rhs_hi]  [rhs_lo]
    //                       [lhs_hi]  [lhs_lo]
    //                     X___________________
    //                       [------tmp0------]
    //             [------tmp1------]
    //             [------tmp2------]
    //     [------tmp3------]
    //                       [-------add------]
    // +_______________________________________
    //                       [------sum0------]
    //     [------sum1------]

    let lhs_lo = lhs as u64;
    let rhs_lo = rhs as u64;
    let lhs_hi = (lhs.wrapping_shr(64)) as u64;
    let rhs_hi = (rhs.wrapping_shr(64)) as u64;
    let tmp0 = (lhs_lo as u128).wrapping_mul(rhs_lo as u128);
    let tmp1 = (lhs_lo as u128).wrapping_mul(rhs_hi as u128);
    let tmp2 = (lhs_hi as u128).wrapping_mul(rhs_lo as u128);
    let tmp3 = (lhs_hi as u128).wrapping_mul(rhs_hi as u128);
    // tmp1 and tmp2 straddle the boundary. We have to handle three carries
    let (sum0, carry0) = tmp0.overflowing_add(tmp1.wrapping_shl(64));
    let (sum0, carry1) = sum0.overflowing_add(tmp2.wrapping_shl(64));
    let (sum0, carry2) = sum0.overflowing_add(carry);
    let sum1 = tmp3
        .wrapping_add(tmp1.wrapping_shr(64))
        .wrapping_add(tmp2.wrapping_shr(64))
        .wrapping_add(carry0 as u128)
        .wrapping_add(carry1 as u128)
        .wrapping_add(carry2 as u128);
    (sum0, sum1)
}

fn mul_two_slices(slice1: &[u32], slice2: &[u32]) -> Vec<u32> {
    // https://en.wikipedia.org/wiki/Karatsuba_algorithm

    let l1 = slice1.len();
    let l2 = slice2.len();

    if l1 == 0 || l2 == 0 {
        return vec![];
    } else if l1 == 1 {
        let mut overflow = 0;
        let mut res: Vec<u32> = Vec::with_capacity(l2 + 1);

        #[allow(clippy::needless_range_loop)]
        for i in 0..l2 {
            let mut r = (slice2[i] as u64) * (slice1[0] as u64);
            r += overflow as u64;
            let m = r as u32;
            overflow = (r >> 32) as u32;
            res.push(m);
        }

        if overflow != 0 {
            res.push(overflow);
        }

        return res;
    } else if l2 == 1 {
        let mut overflow = 0;
        let mut res: Vec<u32> = Vec::with_capacity(l2 + 1);

        #[allow(clippy::needless_range_loop)]
        for i in 0..l1 {
            let mut r = (slice1[i] as u64) * (slice2[0] as u64);
            r += overflow as u64;
            let m = r as u32;
            overflow = (r >> 32) as u32;
            res.push(m);
        }

        if overflow != 0 {
            res.push(overflow);
        }

        return res;
    }

    let m = std::cmp::min(l1, l2);
    let m2 = (m as u32) / 2;

    let (low1, high1) = slice1.split_at(m2 as usize);
    let (low2, high2) = slice2.split_at(m2 as usize);

    let z0 = mul_two_slices(low1, low2);
    let z1 = mul_two_slices(&add_two_slices(low1, high1), &add_two_slices(low2, high2));
    let z2 = mul_two_slices(high1, high2);

    let mut op0 = z2.clone();

    op0.reverse();

    op0.resize(op0.len() + (m2 as usize * 2), 0);

    op0.reverse();

    let mut op1 = sub_two_slices(&sub_two_slices(&z1, &z2), &z0);

    op1.reverse();

    op1.resize(op1.len() + (m2 as usize), 0);

    op1.reverse();

    add_two_slices(&add_two_slices(&op0, &op1), &z0)
}

fn mul(a: &SmallUint, b: &SmallUint) -> SmallUint {
    match (&a.0, &b.0) {
        (&SmallUintType::Inline(i), &SmallUintType::Inline(j)) => {
            match carrying_mul_u128(i, j, 0) {
                (t, 0) => SmallUint(SmallUintType::Inline(t)),
                (t, o) => {
                    let mut res = Vec::with_capacity(8);

                    let mut v = t;
                    #[allow(clippy::needless_range_loop)]
                    for _ in 0..4 {
                        res.push(v as u32);

                        v >>= 32;
                    }

                    let mut v = o;
                    for _ in 4..8 {
                        res.push(v as u32);

                        v >>= 32;
                    }

                    while res[res.len() - 1] == 0 {
                        res.pop();
                    }

                    let size = res.len();

                    let mut slice = ManuallyDrop::new(res.into_boxed_slice());

                    SmallUint(SmallUintType::Heap((slice.as_mut_ptr(), size)))
                }
            }
        }

        (&SmallUintType::Heap((r, s)), &SmallUintType::Inline(i))
        | (&SmallUintType::Inline(i), &SmallUintType::Heap((r, s))) => {
            let slice1 = unsafe { core::slice::from_raw_parts(r, s) };

            let mut res = [0, 0, 0, 0];

            let mut v = i;
            #[allow(clippy::needless_range_loop)]
            for r in 0..4 {
                res[r] = v as u32;

                v >>= 32;
            }

            let result = mul_two_slices(slice1, &res[..]);
            let size = result.len();

            let mut slice = ManuallyDrop::new(result.into_boxed_slice());

            SmallUint(SmallUintType::Heap((slice.as_mut_ptr(), size)))
        }

        (&SmallUintType::Heap((r, s)), &SmallUintType::Heap((i, j))) => {
            let slice1 = unsafe { core::slice::from_raw_parts(r, s) };
            let slice2 = unsafe { core::slice::from_raw_parts(i, j) };

            let res = mul_two_slices(slice1, slice2);
            let size = res.len();

            let mut slice = ManuallyDrop::new(res.into_boxed_slice());

            SmallUint(SmallUintType::Heap((slice.as_mut_ptr(), size)))
        }
    }
}

basic_op!(Mul, mul, SmallUint, mul);

impl<'a> MulAssign<&'a SmallUint> for SmallUint {
    fn mul_assign(&mut self, rhs: &'a SmallUint) {
        *self = self.clone() * rhs;
    }
}

impl MulAssign<SmallUint> for SmallUint {
    fn mul_assign(&mut self, rhs: SmallUint) {
        *self = self.clone() * rhs;
    }
}

fn mul_signed(a: &SmallInt, b: &SmallInt) -> SmallInt {
    let a_sign;
    match &a.0 {
        SmallIntType::Inline(i) => a_sign = i.signum() as i8,
        SmallIntType::Heap((_, s)) => a_sign = s.signum() as i8,
    }

    let b_sign;
    match &b.0 {
        SmallIntType::Inline(i) => b_sign = i.signum() as i8,
        SmallIntType::Heap((_, s)) => b_sign = s.signum() as i8,
    }

    match (a_sign, b_sign) {
        x if (x.0 >= 0 && x.1 >= 0) => SmallInt::from(
            SmallUint::from_smallint_unsigned(a.clone())
                * SmallUint::from_smallint_unsigned(b.clone()),
        ),
        x if (x.0 < 0 && x.1 < 0) => SmallInt::from(
            SmallUint::from_smallint_unsigned(a.clone())
                * SmallUint::from_smallint_unsigned(b.clone()),
        ),

        x if (x.0 >= 0 && x.1 < 0) => -SmallInt::from(
            SmallUint::from_smallint_unsigned(a.clone())
                * SmallUint::from_smallint_unsigned(b.clone()),
        ),

        x if (x.0 < 0 && x.1 >= 0) => -SmallInt::from(
            SmallUint::from_smallint_unsigned(a.clone())
                * SmallUint::from_smallint_unsigned(b.clone()),
        ),

        (_, _) => {
            panic!("This shouldn't happen. ");
        }
    }
}

basic_op!(Mul, mul, SmallInt, mul_signed);

impl<'a> MulAssign<&'a SmallInt> for SmallInt {
    fn mul_assign(&mut self, rhs: &'a SmallInt) {
        *self = self.clone() * rhs;
    }
}

impl MulAssign<SmallInt> for SmallInt {
    fn mul_assign(&mut self, rhs: SmallInt) {
        *self = self.clone() * rhs;
    }
}

fn div(_a: &SmallUint, _b: &SmallUint) -> SmallUint {
    todo!()
}

basic_op!(Div, div, SmallUint, div);

impl<'a> DivAssign<&'a SmallUint> for SmallUint {
    fn div_assign(&mut self, rhs: &'a SmallUint) {
        *self = self.clone() / rhs;
    }
}

impl DivAssign<SmallUint> for SmallUint {
    fn div_assign(&mut self, rhs: SmallUint) {
        *self = self.clone() / rhs;
    }
}

fn div_signed(_a: &SmallInt, _b: &SmallInt) -> SmallInt {
    todo!()
}

basic_op!(Div, div, SmallInt, div_signed);

impl<'a> DivAssign<&'a SmallInt> for SmallInt {
    fn div_assign(&mut self, rhs: &'a SmallInt) {
        *self = self.clone() / rhs;
    }
}

impl DivAssign<SmallInt> for SmallInt {
    fn div_assign(&mut self, rhs: SmallInt) {
        *self = self.clone() / rhs;
    }
}