smallint/src/ops.rs

use crate::smallint::SmallUintType;
use crate::SmallUint;
use core::mem::ManuallyDrop;
use core::ops::{Add, Mul, Sub};

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

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

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

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

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

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

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

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

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, SmallUint, add);

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, SmallUint, sub);

// 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, SmallUint, mul);