use crate::natural::InnerNatural::{Large, Small};
use crate::natural::Natural;
use crate::natural::arithmetic::add::{
limbs_add_greater, limbs_slice_add_greater_in_place_left, limbs_slice_add_limb_in_place,
};
use crate::natural::arithmetic::mul::limb::{limbs_mul_limb_to_out, limbs_slice_mul_limb_in_place};
use crate::natural::arithmetic::mul::{limbs_mul_to_out, limbs_mul_to_out_scratch_len};
use crate::platform::{DoubleLimb, Limb};
use alloc::vec::Vec;
use core::mem::swap;
use malachite_base::num::arithmetic::traits::{AddMul, AddMulAssign, XMulYToZZ};
use malachite_base::num::basic::integers::PrimitiveInt;
use malachite_base::num::basic::traits::{One, Zero};
use malachite_base::num::conversion::traits::{SplitInHalf, WrappingFrom};
pub_test! {limbs_add_mul_limb(xs: &[Limb], ys: &[Limb], limb: Limb) -> Vec<Limb> {
let mut out;
if xs.len() >= ys.len() {
out = xs.to_vec();
limbs_vec_add_mul_limb_greater_in_place_left(&mut out, ys, limb);
} else {
out = ys.to_vec();
limbs_vec_add_mul_limb_smaller_in_place_right(xs, &mut out, limb);
}
out
}}
pub_crate_test! {limbs_slice_add_mul_limb_same_length_in_place_left(
xs: &mut [Limb],
ys: &[Limb],
z: Limb,
) -> Limb {
let len = xs.len();
assert_eq!(ys.len(), len);
let mut carry = 0;
for (x, &y) in xs.iter_mut().zip(ys.iter()) {
let (product_hi, mut product_lo) = XMulYToZZ::x_mul_y_to_zz(y, z);
product_lo = (*x).wrapping_add(product_lo);
let mut add_carry = Limb::from(*x > product_lo);
*x = product_lo.wrapping_add(carry);
add_carry += Limb::from(product_lo > *x);
carry = product_hi.wrapping_add(add_carry);
}
carry
}}
pub(crate) fn limbs_slice_add_mul_two_limbs_matching_length_in_place_left(
xs: &mut [Limb],
ys: &[Limb],
zs: [Limb; 2],
) -> Limb {
let len = ys.len();
assert_eq!(xs.len(), len + 1);
let mut carry_hi: Limb = 0;
let mut carry_lo: Limb = 0;
for (x, &y) in xs.iter_mut().zip(ys.iter()) {
let (mut product_hi, mut product_lo) = XMulYToZZ::x_mul_y_to_zz(y, zs[0]);
product_lo = (*x).wrapping_add(product_lo);
let mut add_carry = Limb::from(*x > product_lo);
*x = product_lo.wrapping_add(carry_lo);
add_carry += Limb::from(product_lo > *x);
carry_lo = product_hi.wrapping_add(add_carry);
carry_lo = carry_hi.wrapping_add(carry_lo);
add_carry = Limb::from(carry_hi > carry_lo);
(product_hi, product_lo) = XMulYToZZ::x_mul_y_to_zz(y, zs[1]);
carry_lo = product_lo.wrapping_add(carry_lo);
add_carry += Limb::from(product_lo > carry_lo);
carry_hi = product_hi.wrapping_add(add_carry);
}
xs[len] = carry_lo;
carry_hi
}
pub_test! {limbs_slice_add_mul_limb_same_length_in_place_right(
xs: &[Limb],
ys: &mut [Limb],
z: Limb,
) -> Limb {
let xs_len = xs.len();
assert_eq!(ys.len(), xs_len);
let mut carry = 0;
let dz = DoubleLimb::from(z);
for (&x, y) in xs.iter().zip(ys.iter_mut()) {
let out = DoubleLimb::from(x) + DoubleLimb::from(*y) * dz + carry;
*y = out.lower_half();
carry = out >> Limb::WIDTH;
}
Limb::wrapping_from(carry)
}}
pub_test! {limbs_vec_add_mul_limb_in_place_left(xs: &mut Vec<Limb>, ys: &[Limb], z: Limb) {
let xs_len = xs.len();
if xs_len >= ys.len() {
limbs_vec_add_mul_limb_greater_in_place_left(xs, ys, z);
} else {
xs.resize(ys.len(), 0);
let (xs_lo, xs_hi) = xs.split_at_mut(xs_len);
let (ys_lo, ys_hi) = ys.split_at(xs_len);
let mut carry = limbs_mul_limb_to_out::<DoubleLimb, Limb>(xs_hi, ys_hi, z);
let inner_carry = limbs_slice_add_mul_limb_same_length_in_place_left(xs_lo, ys_lo, z);
if inner_carry != 0 && limbs_slice_add_limb_in_place(xs_hi, inner_carry) {
carry += 1;
}
if carry != 0 {
xs.push(carry);
}
}
}}
fn limbs_vec_add_mul_limb_greater_in_place_left(xs: &mut Vec<Limb>, ys: &[Limb], z: Limb) {
let ys_len = ys.len();
let carry = limbs_slice_add_mul_limb_same_length_in_place_left(&mut xs[..ys_len], ys, z);
if carry != 0 {
if xs.len() == ys_len {
xs.push(carry);
} else if limbs_slice_add_limb_in_place(&mut xs[ys_len..], carry) {
xs.push(1);
}
}
}
pub_test! {limbs_vec_add_mul_limb_in_place_right(xs: &[Limb], ys: &mut Vec<Limb>, z: Limb) {
let ys_len = ys.len();
if xs.len() >= ys_len {
let carry = limbs_slice_add_mul_limb_same_length_in_place_right(&xs[..ys_len], ys, z);
ys.extend_from_slice(&xs[ys_len..]);
if carry != 0 {
if xs.len() == ys_len {
ys.push(carry);
} else if limbs_slice_add_limb_in_place(&mut ys[ys_len..], carry) {
ys.push(1);
}
}
} else {
limbs_vec_add_mul_limb_smaller_in_place_right(xs, ys, z);
}
}}
fn limbs_vec_add_mul_limb_smaller_in_place_right(xs: &[Limb], ys: &mut Vec<Limb>, z: Limb) {
let (ys_lo, ys_hi) = ys.split_at_mut(xs.len());
let mut carry = limbs_slice_mul_limb_in_place(ys_hi, z);
let inner_carry = limbs_slice_add_mul_limb_same_length_in_place_right(xs, ys_lo, z);
if inner_carry != 0 && limbs_slice_add_limb_in_place(ys_hi, inner_carry) {
carry += 1;
}
if carry != 0 {
ys.push(carry);
}
}
pub_test! {limbs_vec_add_mul_limb_in_place_either(
xs: &mut Vec<Limb>,
ys: &mut Vec<Limb>,
z: Limb,
) -> bool {
if xs.len() >= ys.len() {
limbs_vec_add_mul_limb_greater_in_place_left(xs, ys, z);
false
} else {
limbs_vec_add_mul_limb_smaller_in_place_right(xs, ys, z);
true
}
}}
pub_test! {limbs_add_mul(xs: &[Limb], ys: &[Limb], zs: &[Limb]) -> Vec<Limb> {
let xs_len = xs.len();
let mut out_len = ys.len() + zs.len();
let mut out = vec![0; out_len];
let mut mul_scratch = vec![0; limbs_mul_to_out_scratch_len(ys.len(), zs.len())];
if limbs_mul_to_out(&mut out, ys, zs, &mut mul_scratch) == 0 {
out_len -= 1;
out.pop();
}
assert_ne!(*out.last().unwrap(), 0);
if xs_len >= out_len {
limbs_add_greater(xs, &out)
} else {
if limbs_slice_add_greater_in_place_left(&mut out, xs) {
out.push(1);
}
out
}
}}
pub_test! {limbs_add_mul_in_place_left(xs: &mut Vec<Limb>, ys: &[Limb], zs: &[Limb]) {
let xs_len = xs.len();
let mut out_len = ys.len() + zs.len();
let mut out = vec![0; out_len];
let mut mul_scratch = vec![0; limbs_mul_to_out_scratch_len(ys.len(), zs.len())];
if limbs_mul_to_out(&mut out, ys, zs, &mut mul_scratch) == 0 {
out_len -= 1;
out.pop();
}
assert_ne!(*out.last().unwrap(), 0);
if xs_len < out_len {
swap(xs, &mut out);
}
if limbs_slice_add_greater_in_place_left(xs, &out) {
xs.push(1);
}
}}
impl Natural {
fn add_mul_limb_ref_ref(&self, y: &Self, z: Limb) -> Self {
match (self, y, z) {
(x, _, 0) | (x, &Self::ZERO, _) => x.clone(),
(x, y, 1) => x + y,
(x, &Self::ONE, z) => x + Self::from(z),
(Self(Large(xs)), Self(Large(ys)), z) => Self(Large(limbs_add_mul_limb(xs, ys, z))),
(x, y, z) => x + y * Self::from(z),
}
}
fn add_mul_assign_limb(&mut self, mut y: Self, z: Limb) {
match (&mut *self, &mut y, z) {
(_, _, 0) | (_, &mut Self::ZERO, _) => {}
(x, _, 1) => *x += y,
(x, &mut Self::ONE, z) => *x += Self::from(z),
(Self(Large(xs)), Self(Large(ys)), z) => {
if limbs_vec_add_mul_limb_in_place_either(xs, ys, z) {
*self = y;
}
}
(x, _, z) => *x += y * Self::from(z),
}
}
fn add_mul_assign_limb_ref(&mut self, y: &Self, z: Limb) {
match (&mut *self, y, z) {
(_, _, 0) | (_, &Self::ZERO, _) => {}
(x, y, 1) => *x += y,
(x, &Self::ONE, z) => *x += Self::from(z),
(Self(Large(xs)), Self(Large(ys)), z) => {
limbs_vec_add_mul_limb_in_place_left(xs, ys, z);
}
(x, y, z) => *x += y * Self::from(z),
}
}
}
impl AddMul<Self, Self> for Natural {
type Output = Self;
#[inline]
fn add_mul(mut self, y: Self, z: Self) -> Self {
self.add_mul_assign(y, z);
self
}
}
impl<'a> AddMul<Self, &'a Self> for Natural {
type Output = Self;
#[inline]
fn add_mul(mut self, y: Self, z: &'a Self) -> Self {
self.add_mul_assign(y, z);
self
}
}
impl<'a> AddMul<&'a Self, Self> for Natural {
type Output = Self;
#[inline]
fn add_mul(mut self, y: &'a Self, z: Self) -> Self {
self.add_mul_assign(y, z);
self
}
}
impl<'a, 'b> AddMul<&'a Self, &'b Self> for Natural {
type Output = Self;
#[inline]
fn add_mul(mut self, y: &'a Self, z: &'b Self) -> Self {
self.add_mul_assign(y, z);
self
}
}
impl AddMul<&Natural, &Natural> for &Natural {
type Output = Natural;
fn add_mul(self, y: &Natural, z: &Natural) -> Natural {
match (self, y, z) {
(Natural(Small(x)), y, z) => (y * z).add_limb(*x),
(x, Natural(Small(y)), z) => x.add_mul_limb_ref_ref(z, *y),
(x, y, Natural(Small(z))) => x.add_mul_limb_ref_ref(y, *z),
(Natural(Large(xs)), Natural(Large(ys)), Natural(Large(zs))) => {
Natural(Large(limbs_add_mul(xs, ys, zs)))
}
}
}
}
impl AddMulAssign<Self, Self> for Natural {
fn add_mul_assign(&mut self, mut y: Self, mut z: Self) {
match (&mut *self, &mut y, &mut z) {
(Self(Small(x)), _, _) => *self = (y * z).add_limb(*x),
(_, Self(Small(y)), _) => self.add_mul_assign_limb(z, *y),
(_, _, Self(Small(z))) => self.add_mul_assign_limb(y, *z),
(Self(Large(xs)), Self(Large(ys)), Self(Large(zs))) => {
limbs_add_mul_in_place_left(xs, ys, zs);
}
}
}
}
impl<'a> AddMulAssign<Self, &'a Self> for Natural {
fn add_mul_assign(&mut self, mut y: Self, z: &'a Self) {
match (&mut *self, &mut y, z) {
(Self(Small(x)), _, _) => *self = (y * z).add_limb(*x),
(_, Self(Small(y)), _) => self.add_mul_assign_limb_ref(z, *y),
(_, _, Self(Small(z))) => self.add_mul_assign_limb(y, *z),
(Self(Large(xs)), Self(Large(ys)), Self(Large(zs))) => {
limbs_add_mul_in_place_left(xs, ys, zs);
}
}
}
}
impl<'a> AddMulAssign<&'a Self, Self> for Natural {
fn add_mul_assign(&mut self, y: &'a Self, mut z: Self) {
match (&mut *self, y, &mut z) {
(Self(Small(x)), _, _) => *self = (y * z).add_limb(*x),
(_, Self(Small(y)), _) => self.add_mul_assign_limb(z, *y),
(_, _, Self(Small(z))) => self.add_mul_assign_limb_ref(y, *z),
(Self(Large(xs)), Self(Large(ys)), Self(Large(zs))) => {
limbs_add_mul_in_place_left(xs, ys, zs);
}
}
}
}
impl<'a, 'b> AddMulAssign<&'a Self, &'b Self> for Natural {
fn add_mul_assign(&mut self, y: &'a Self, z: &'b Self) {
match (&mut *self, y, z) {
(Self(Small(x)), _, _) => *self = (y * z).add_limb(*x),
(_, Self(Small(y)), _) => self.add_mul_assign_limb_ref(z, *y),
(_, _, Self(Small(z))) => self.add_mul_assign_limb_ref(y, *z),
(Self(Large(xs)), Self(Large(ys)), Self(Large(zs))) => {
limbs_add_mul_in_place_left(xs, ys, zs);
}
}
}
}