use crate::int::policy::mul::dispatch_slice as mul_slice;
use crate::int::algos::isqrt::isqrt_newton::isqrt_newton;
use crate::int::algos::support::limbs::{cmp_cross, is_zero, sub_assign};
use crate::int::types::compute_limbs::{ComputeLimbs, Limbs};
use crate::int::types::Int;
use crate::support::rounding::RoundingMode;
#[inline]
fn sig_len(a: &[u64]) -> usize {
let mut l = a.len();
while l > 1 && a[l - 1] == 0 {
l -= 1;
}
l
}
#[inline]
#[must_use]
pub(crate) fn sqrt_newton<const N: usize>(raw: Int<N>, scale: u32, mode: RoundingMode) -> Int<N>
where
Limbs<N>: ComputeLimbs,
{
if raw <= Int::<N>::ZERO {
return Int::<N>::ZERO;
}
let mut n_buf = Limbs::<N>::double_buffered_u64();
let n = n_buf.as_mut();
n[..N].copy_from_slice(raw.unsigned_abs().as_limbs());
let mut nl = sig_len(&n[..N]);
{
let mut tmp_buf = Limbs::<N>::double_buffered_u64();
let tmp = tmp_buf.as_mut();
for _ in 0..scale {
let out = nl + 1;
for t in tmp[..out].iter_mut() {
*t = 0;
}
mul_slice(&n[..nl], &[10u64], &mut tmp[..out]);
n[..out].copy_from_slice(&tmp[..out]);
nl = sig_len(&n[..out]);
}
}
let mut q_buf = Limbs::<N>::double_buffered_u64();
let q = q_buf.as_mut();
isqrt_newton(&n[..nl], &mut q[..nl]);
let ql = sig_len(&q[..nl]);
let mut qsq_buf = Limbs::<N>::double_buffered_u64();
let qsq = qsq_buf.as_mut();
let qsq_cap = qsq.len();
mul_slice(&q[..ql], &q[..ql], &mut qsq[..(2 * ql).min(qsq_cap)]);
let mut diff_buf = Limbs::<N>::double_buffered_u64();
let diff = diff_buf.as_mut();
diff[..nl].copy_from_slice(&n[..nl]);
sub_assign(&mut diff[..nl], &qsq[..nl]);
let halfway_round_up = cmp_cross(&diff[..nl], &q[..ql]) > 0;
let diff_nonzero = !is_zero(&diff[..nl]);
let bump = match mode {
RoundingMode::HalfToEven
| RoundingMode::HalfAwayFromZero
| RoundingMode::HalfTowardZero => halfway_round_up,
RoundingMode::Trunc | RoundingMode::Floor => false,
RoundingMode::Ceiling => diff_nonzero,
};
if bump {
let mut i = 0;
loop {
let (v, c) = q[i].overflowing_add(1);
q[i] = v;
if !c {
break;
}
i += 1;
}
}
let mut out = [0u64; N];
out.copy_from_slice(&q[..N]);
Int::<N>::from_limbs(out)
}