use crate::int::algos::support::limbs::{bit_len, cmp, shl, sub_assign};
use crate::int::types::compute_limbs::MAX_DOUBLE_LIMBS;
pub(crate) fn isqrt_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;
debug_assert!(work <= MAX_DOUBLE_LIMBS, "isqrt_schoolbook scratch overflow");
let mut p = [0u64; MAX_DOUBLE_LIMBS];
let mut r = [0u64; MAX_DOUBLE_LIMBS];
let mut tmp = [0u64; MAX_DOUBLE_LIMBS];
let start: i64 = if bits % 2 == 1 { (bits - 1) as i64 } else { (bits - 2) as i64 };
let mut b = start;
while b >= 0 {
shl(&r[..work], 2, &mut tmp[..work]);
r[..work].copy_from_slice(&tmp[..work]);
let b_u32 = b as u32;
let li = (b_u32 / 64) as usize;
let bo = b_u32 % 64;
let group: u64 = if bo == 63 {
let lo = if li < n.len() { (n[li] >> 63) & 1 } else { 0 };
let hi = if li + 1 < n.len() { n[li + 1] & 1 } else { 0 };
(hi << 1) | lo
} else if li < n.len() {
(n[li] >> bo) & 3
} else { 0 };
r[0] |= group;
shl(&p[..work], 2, &mut tmp[..work]);
let mut t = [0u64; MAX_DOUBLE_LIMBS];
t[..work].copy_from_slice(&tmp[..work]);
t[0] |= 1;
if cmp(&r[..work], &t[..work]) >= 0 {
sub_assign(&mut r[..work], &t[..work]);
shl(&p[..work], 1, &mut tmp[..work]);
p[..work].copy_from_slice(&tmp[..work]);
p[0] |= 1;
} else {
shl(&p[..work], 1, &mut tmp[..work]);
p[..work].copy_from_slice(&tmp[..work]);
}
b -= 2;
}
let copy_len = out.len().min(work);
out[..copy_len].copy_from_slice(&p[..copy_len]);
}
#[cfg(test)]
mod tests {
use super::isqrt_schoolbook;
use crate::int::algos::isqrt::isqrt_newton::isqrt_newton;
fn sb64(n: u64) -> u64 {
let i = [n]; let mut o = [0u64]; isqrt_schoolbook(&i, &mut o); o[0]
}
fn sb128(n: u128) -> u128 {
let i = [n as u64, (n>>64) as u64]; let mut o = [0u64, 0u64];
isqrt_schoolbook(&i, &mut o); (o[0] as u128) | ((o[1] as u128) << 64)
}
fn nt64(n: u64) -> u64 {
let i = [n]; let mut o = [0u64]; isqrt_newton(&i, &mut o); o[0]
}
fn nt128(n: u128) -> u128 {
let i = [n as u64, (n>>64) as u64]; let mut o = [0u64, 0u64];
isqrt_newton(&i, &mut o); (o[0] as u128) | ((o[1] as u128) << 64)
}
#[test]
fn isqrt_schoolbook_known_values_u64() {
let cases: &[(u64,u64)] = &[
(0,0),(1,1),(2,1),(3,1),(4,2),(5,2),(8,2),(9,3),
(15,3),(16,4),(17,4),(24,4),(25,5),(35,5),(36,6),
(99,9),(100,10),(10_000,100),(999_999,999),(1_000_000,1_000),
(2_u64.pow(32)-1,65_535),(2_u64.pow(32),65_536),
(2_u64.pow(32)+1,65_536),
(2_u64.pow(63),3_037_000_499),(u64::MAX,4_294_967_295),
];
for &(n,e) in cases { let g=sb64(n); assert_eq!(g,e,"n={n}: got {g} expected {e}"); }
}
#[test]
fn isqrt_schoolbook_known_values_u128() {
let cases: &[(u128,u128)] = &[
(0,0),(1,1),(4,2),(9,3),
(2_u128.pow(64), 2_u128.pow(32)),
(2_u128.pow(64)-1, u32::MAX as u128),
(2_u128.pow(127), 13_043_817_825_332_782_212_u128),
(u128::MAX, u64::MAX as u128),
];
for &(n,e) in cases { let g=sb128(n); assert_eq!(g,e,"n={n}: got {g} expected {e}"); }
}
#[test]
fn isqrt_schoolbook_matches_newton_u64_range() {
for n in 0u64..=512 { let sb=sb64(n); let nt=nt64(n);
assert_eq!(sb,nt,"n={n}: sb={sb} nt={nt}"); }
for n in [u64::MAX,u64::MAX-1,u64::MAX-2,
2_u64.pow(63),2_u64.pow(63)-1,2_u64.pow(32),2_u64.pow(32)-1] {
let sb=sb64(n); let nt=nt64(n);
assert_eq!(sb,nt,"n={n}: sb={sb} nt={nt}");
}
}
#[test]
fn isqrt_schoolbook_matches_newton_u128_range() {
for n in 0u128..=256 { let sb=sb128(n); let nt=nt128(n);
assert_eq!(sb,nt,"n={n}"); }
for k in [2u128,3,5,10,100,1_000,65_536,1_000_000] {
let n=k*k;
assert_eq!(sb128(n),k,"isqrt({k}^2)");
assert_eq!(nt128(n),k,"newton isqrt({k}^2)");
if n>0 { assert_eq!(sb128(n-1),k-1,"isqrt({k}^2-1)"); }
assert_eq!(sb128(n+1),k,"isqrt({k}^2+1)");
}
}
}