use crate::int::algos::mul::mul_schoolbook::mul_schoolbook;
use crate::int::algos::support::limbs::{
add_assign, bit_len, cmp, shl, sub_assign,
};
use crate::int::types::compute_limbs::MAX_QUADRUPLE_LIMBS;
pub(crate) fn icbrt_schoolbook(n: &[u64], out: &mut [u64]) {
for o in out.iter_mut() {
*o = 0;
}
let bits = bit_len(n);
if bits == 0 {
return;
}
if bits <= 1 {
out[0] = 1;
return;
}
let work = n.len() + 1;
let sq_work = (work * 2).min(MAX_QUADRUPLE_LIMBS);
debug_assert!(work <= MAX_QUADRUPLE_LIMBS, "icbrt_schoolbook scratch overflow");
let mut p = [0u64; MAX_QUADRUPLE_LIMBS];
let mut p_sq = [0u64; MAX_QUADRUPLE_LIMBS];
let mut rem = [0u64; MAX_QUADRUPLE_LIMBS];
rem[..n.len()].copy_from_slice(n);
let result_bits = bits.div_ceil(3);
let mut k = result_bits as i64 - 1;
while k >= 0 {
let bit_pos = k as u32;
let d_limb = (bit_pos / 64) as usize;
let d_off = bit_pos % 64;
let mut d = [0u64; MAX_QUADRUPLE_LIMBS];
if d_limb < MAX_QUADRUPLE_LIMBS {
d[d_limb] = 1u64 << d_off;
}
let mut t1 = [0u64; MAX_QUADRUPLE_LIMBS];
{
let mut shifted = [0u64; MAX_QUADRUPLE_LIMBS];
shl(&p_sq[..sq_work], 1, &mut shifted[..sq_work]);
t1[..sq_work].copy_from_slice(&shifted[..sq_work]);
add_assign(&mut t1[..sq_work], &p_sq[..sq_work]);
}
let pd_work = (work + d_limb + 1).min(MAX_QUADRUPLE_LIMBS);
let mut p_d = [0u64; MAX_QUADRUPLE_LIMBS];
mul_schoolbook(&p[..work], &d[..d_limb + 1], &mut p_d[..pd_work]);
let mut t2 = [0u64; MAX_QUADRUPLE_LIMBS];
{
let mut shifted = [0u64; MAX_QUADRUPLE_LIMBS];
shl(&p_d[..pd_work], 1, &mut shifted[..pd_work]);
t2[..pd_work].copy_from_slice(&shifted[..pd_work]);
add_assign(&mut t2[..pd_work], &p_d[..pd_work]);
}
let d_sq_pos = (k as u32) * 2;
let d_sq_limb = (d_sq_pos / 64) as usize;
let d_sq_off = d_sq_pos % 64;
let mut d_sq = [0u64; MAX_QUADRUPLE_LIMBS];
if d_sq_limb < MAX_QUADRUPLE_LIMBS {
d_sq[d_sq_limb] = 1u64 << d_sq_off;
}
let inner_work = sq_work.max(pd_work).max(d_sq_limb + 1) + 1;
let inner_work = inner_work.min(MAX_QUADRUPLE_LIMBS);
let mut inner = [0u64; MAX_QUADRUPLE_LIMBS];
inner[..sq_work].copy_from_slice(&t1[..sq_work]);
add_assign(&mut inner[..inner_work], &t2[..pd_work.min(inner_work)]);
add_assign(&mut inner[..inner_work], &d_sq[..(d_sq_limb + 1).min(inner_work)]);
let delta_work = (d_limb + 1 + inner_work).min(MAX_QUADRUPLE_LIMBS);
let mut delta = [0u64; MAX_QUADRUPLE_LIMBS];
mul_schoolbook(&d[..d_limb + 1], &inner[..inner_work], &mut delta[..delta_work]);
if cmp(&rem[..work], &delta[..delta_work.min(work)]) >= 0 {
sub_assign(&mut rem[..work], &delta[..delta_work.min(work)]);
add_assign(&mut p[..work], &d[..d_limb + 1]);
let mut new_p_sq = [0u64; MAX_QUADRUPLE_LIMBS];
mul_schoolbook(&p[..work], &p[..work], &mut new_p_sq[..sq_work]);
p_sq[..sq_work].copy_from_slice(&new_p_sq[..sq_work]);
}
k -= 1;
}
let copy_len = out.len().min(work);
out[..copy_len].copy_from_slice(&p[..copy_len]);
}
#[cfg(test)]
mod tests {
use super::icbrt_schoolbook;
use crate::int::algos::icbrt::icbrt_newton::icbrt_newton;
fn schoolbook_u64(n: u64) -> u64 {
let input = [n];
let mut out = [0u64];
icbrt_schoolbook(&input, &mut out);
out[0]
}
fn schoolbook_u128(n: u128) -> u128 {
let input = [n as u64, (n >> 64) as u64];
let mut out = [0u64, 0u64];
icbrt_schoolbook(&input, &mut out);
(out[0] as u128) | ((out[1] as u128) << 64)
}
fn newton_u64(n: u64) -> u64 {
let input = [n];
let mut out = [0u64];
icbrt_newton(&input, &mut out);
out[0]
}
fn newton_u128(n: u128) -> u128 {
let input = [n as u64, (n >> 64) as u64];
let mut out = [0u64, 0u64];
icbrt_newton(&input, &mut out);
(out[0] as u128) | ((out[1] as u128) << 64)
}
#[test]
fn icbrt_schoolbook_known_values_u64() {
let cases: &[(u64, u64)] = &[
(0, 0),
(1, 1),
(2, 1),
(7, 1),
(8, 2), (9, 2),
(26, 2),
(27, 3), (28, 3),
(63, 3),
(64, 4), (65, 4),
(125, 5), (126, 5),
(999, 9),
(1_000, 10), (1_001, 10),
(2_u64.pow(63), 2_097_152), (u64::MAX, 2_642_245), ];
for &(n, expected) in cases {
let got = schoolbook_u64(n);
assert_eq!(got, expected,
"icbrt_schoolbook({n}) = {got}, expected {expected}");
}
}
#[test]
fn icbrt_schoolbook_known_values_u128() {
let cases: &[(u128, u128)] = &[
(0, 0),
(1, 1),
(7, 1),
(8, 2),
(27, 3),
(64, 4),
(125, 5),
(2_u128.pow(64), 2_642_245), (2_u128.pow(127), 5_541_191_377_756), (1_000_000_000_u128, 1_000), ];
for &(n, expected) in cases {
let got = schoolbook_u128(n);
assert_eq!(got, expected,
"icbrt_schoolbook({n}) = {got}, expected {expected}");
}
}
#[test]
fn icbrt_schoolbook_matches_newton_u64_range() {
for n in 0u64..=512 {
let sb = schoolbook_u64(n);
let nt = newton_u64(n);
assert_eq!(sb, nt, "mismatch at n={n}: schoolbook={sb}, newton={nt}");
}
for n in [u64::MAX, u64::MAX - 1, 2_u64.pow(63),
2_u64.pow(63) - 1, 2_u64.pow(32), 2_u64.pow(21).pow(3)] {
let sb = schoolbook_u64(n);
let nt = newton_u64(n);
assert_eq!(sb, nt, "mismatch at n={n}: schoolbook={sb}, newton={nt}");
}
}
#[test]
fn icbrt_schoolbook_matches_newton_u128_range() {
for n in 0u128..=256 {
let sb = schoolbook_u128(n);
let nt = newton_u128(n);
assert_eq!(sb, nt, "mismatch at n={n}");
}
for k in [2u128, 3, 5, 10, 100, 1_000, 10_000] {
let n = k * k * k;
let sb = schoolbook_u128(n);
let nt = newton_u128(n);
assert_eq!(sb, k, "icbrt({n}) schoolbook = {sb}, expected {k}");
assert_eq!(nt, k, "icbrt({n}) newton = {nt}, expected {k}");
if n > 0 {
let sb_below = schoolbook_u128(n - 1);
assert_eq!(sb_below, k - 1,
"icbrt({}) schoolbook should be {}", n - 1, k - 1);
}
let sb_above = schoolbook_u128(n + 1);
assert_eq!(sb_above, k,
"icbrt({}) schoolbook should be {}", n + 1, k);
}
}
}