use super::super::div::limbs::{div_qr, div_qr_1, div_qr_2, invert_3by2, invert_limb, qr_itch};
use super::super::mul::limbs::{add_1, cmp_n, mul, mul_itch, sqr, sqr_itch, sub_1, sub_n};
use super::{ROUGH_BITS, lshift, mul_1, rough_root, rshift};
use core::cmp::Ordering;
const ESTIMATE_BITS: u64 = 30;
const GUARD_BITS: u64 = 64;
const MAX_STEPS: usize = 128;
struct Plan {
xb: u64,
guard: u64,
cap: usize,
sn: usize,
}
impl Plan {
fn of(ub: u64, un: usize, k: u32) -> Plan {
let xb = (ub - 1) / u64::from(k) + 1;
let guard = if xb > ESTIMATE_BITS && un > 3 * k as usize && logk(k) + 2 < ESTIMATE_BITS {
GUARD_BITS
} else {
0
};
let sn = (xb + guard).div_ceil(64) as usize + 2;
let cap = un + if guard > 0 { k as usize + 2 } else { sn } + 4;
Plan { xb, guard, cap, sn }
}
}
fn logk(k: u32) -> u64 {
u64::from(32 - (k - 1).leading_zeros()) + 2
}
pub(super) fn root_itch(un: usize, k: u32) -> usize {
let plan = Plan::of(64 * un as u64, un, k);
plan.sn + 8 * (plan.cap + 2)
}
#[derive(Clone, Copy)]
struct Buffers {
stride: usize,
t: *mut u64,
tb: *mut u64,
pa: *mut u64,
pb: *mut u64,
pk: *mut u64,
nbuf: *mut u64,
dbuf: *mut u64,
q: *mut u64,
}
pub(super) unsafe fn root(u: *const u64, un: usize, k: u32, w: *mut u64) -> usize {
unsafe {
debug_assert!(k >= 3 && *u.add(un - 1) != 0);
let ub = 64 * un as u64 - u64::from((*u.add(un - 1)).leading_zeros());
let plan = Plan::of(ub, un, k);
debug_assert!(plan.sn + 8 * (plan.cap + 2) <= root_itch(un, k));
let s = w;
let stride = plan.cap + 2;
let base = w.add(plan.sn);
let buffers = Buffers {
stride,
t: base,
tb: base.add(stride),
pa: base.add(2 * stride),
pb: base.add(3 * stride),
pk: base.add(4 * stride),
nbuf: base.add(5 * stride),
dbuf: base.add(6 * stride),
q: base.add(7 * stride),
};
let mut scratch = Vec::new();
let k64 = u64::from(k);
let xb = plan.xb;
if xb <= ESTIMATE_BITS {
*s = settle(
estimate(u, un, ub, 0, k),
k,
u,
un,
ub,
0,
&buffers,
&mut scratch,
);
return 1;
}
let lk = logk(k);
let mut sizes = [0u64; MAX_STEPS];
let mut steps = 0;
let mut p = xb + plan.guard;
loop {
sizes[steps] = p;
steps += 1;
if p <= ESTIMATE_BITS {
break;
}
let q = (p + lk).div_ceil(2);
p = if q < p { q } else { p - 1 };
}
let at = |p: u64| (i128::from(xb) - i128::from(p)) * i128::from(k);
steps -= 1;
let mut p = sizes[steps];
*s = settle(
estimate(u, un, ub, at(p), k),
k,
u,
un,
ub,
at(p),
&buffers,
&mut scratch,
);
let mut sn = 1;
while steps > 0 {
steps -= 1;
let next = sizes[steps];
let b = next - p;
let Buffers {
stride,
t,
tb,
pa,
pb,
pk,
nbuf,
dbuf,
q,
} = buffers;
let tbn = shift_into(tb, u, un, at(p) - i128::from(b));
debug_assert!(tbn < stride);
let tn = shift_into(t, tb, tbn, i128::from(b));
let (mut wp, mut wn, mut pkn);
loop {
(wp, wn) = power_into(s, sn, k64 - 1, pa, pb, stride, &mut scratch);
debug_assert!(wn + sn <= stride);
pkn = mul_into(pk, wp, wn, s, sn, &mut scratch);
if compare(pk, pkn, t, tn) != Ordering::Greater {
break;
}
sub_1(s, sn, 1);
sn = normalized(s, sn);
}
let pbn = shift_into(pk, pk, pkn, -i128::from(b));
debug_assert!(pbn <= tbn && wn < stride && sn + (b.div_ceil(64) as usize) < plan.sn);
let borrow = sub_n(tb, tb, pk, pbn);
sub_1(tb.add(pbn), tbn - pbn, u64::from(borrow));
let nn = normalized(tb, tbn);
let carry = mul_1(wp, wp, wn, k64);
*wp.add(wn) = carry;
let dn = wn + usize::from(carry != 0);
let qn = (nn + 1).saturating_sub(dn);
let cut = dn.saturating_sub(qn + 2).min(nn);
let (nt, nnt, dt, dnt) = if cut > 0 {
*tb.add(nn) = add_1(tb.add(cut), nn - cut, 1);
(
tb.add(cut),
normalized(tb.add(cut), nn - cut + 1),
wp.add(cut),
dn - cut,
)
} else {
(tb, nn, wp, dn)
};
let mut qn = quotient_into(q, nt, nnt, dt, dnt, nbuf, dbuf, &mut scratch);
if bit_len(q, qn) > b {
qn = b.div_ceil(64) as usize;
core::ptr::write_bytes(q, u8::MAX, qn);
if b % 64 != 0 {
*q.add(qn - 1) = (1 << (b % 64)) - 1;
}
}
sn = shift_into(s, s, sn, -i128::from(b));
let low = b.div_ceil(64) as usize;
if sn < low {
core::ptr::write_bytes(s.add(sn), 0, low - sn);
sn = low;
}
for i in 0..qn {
*s.add(i) |= *q.add(i);
}
p = next;
}
if plan.guard == 0 {
return settle_long(s, sn, k, u, un, &buffers, &mut scratch);
}
let low = *s;
core::ptr::copy(s.add(1), s, sn - 1);
sn -= 1;
if low >= 2 {
sn
} else {
settle_long(s, sn, k, u, un, &buffers, &mut scratch)
}
}
}
unsafe fn estimate(u: *const u64, un: usize, ub: u64, sh: i128, k: u32) -> u64 {
let tb = (i128::from(ub) - sh) as u64;
let top = unsafe { top_bits(u, un) };
let root = if tb.div_ceil(u64::from(k)) <= u64::from(ROUGH_BITS) {
rough_root(top, tb, k)
} else {
let log = (top as f64).log2() + tb as f64 - 64.0;
(log / f64::from(k)).exp2()
};
(root as u64).clamp(1, (1 << ESTIMATE_BITS) - 1)
}
unsafe fn top_bits(u: *const u64, un: usize) -> u64 {
unsafe {
let top = *u.add(un - 1);
let z = top.leading_zeros();
if z == 0 || un == 1 {
return top << z;
}
(top << z) | (*u.add(un - 2) >> (64 - z))
}
}
#[allow(clippy::too_many_arguments)]
unsafe fn settle(
mut s: u64,
k: u32,
u: *const u64,
un: usize,
ub: u64,
sh: i128,
buffers: &Buffers,
scratch: &mut Vec<u64>,
) -> u64 {
unsafe {
let k64 = u64::from(k);
let tb = (i128::from(ub) - sh) as u64;
let mut top = top_bits(u, un);
if tb < 64 {
top &= !(u64::MAX >> tb);
}
let mut tn = 0;
let mut above = |x: u64| -> bool {
match compare_power(x, k64, top, tb) {
Some(order) => order == Ordering::Greater,
None => {
if tn == 0 {
tn = shift_into(buffers.t, u, un, sh);
}
let (pp, pn) = power_into(
&raw const x,
1,
k64,
buffers.pa,
buffers.pb,
buffers.stride,
scratch,
);
compare(pp, pn, buffers.t, tn) == Ordering::Greater
}
}
};
let (low, high) = (s > 1 && above(s), above(s + 1));
if !low && high {
return s;
}
if low {
s -= 1;
while s > 1 && above(s) {
s -= 1;
}
return s;
}
s += 1;
while !above(s + 1) {
s += 1;
}
s
}
}
unsafe fn settle_long(
s: *mut u64,
sn: usize,
k: u32,
u: *const u64,
un: usize,
buffers: &Buffers,
scratch: &mut Vec<u64>,
) -> usize {
unsafe {
let k64 = u64::from(k);
let ub = bit_len(u, un);
let sb = bit_len(s, sn);
let top = top_bits(u, un);
let st = top_bits(s, sn);
let mut sn = sn;
let mut over = match compare_power_scaled(st, sb, k64, top, ub) {
Some(order) => order == Ordering::Greater,
None => {
let (pp, pn) =
power_into(s, sn, k64, buffers.pa, buffers.pb, buffers.stride, scratch);
compare(pp, pn, u, un) == Ordering::Greater
}
};
while over {
sub_1(s, sn, 1);
sn = normalized(s, sn);
let sb = bit_len(s, sn);
over = match compare_power_scaled(top_bits(s, sn), sb, k64, top, ub) {
Some(order) => order == Ordering::Greater,
None => {
let (pp, pn) =
power_into(s, sn, k64, buffers.pa, buffers.pb, buffers.stride, scratch);
compare(pp, pn, u, un) == Ordering::Greater
}
};
}
sn
}
}
fn compare_power(x: u64, k: u64, top: u64, tb: u64) -> Option<Ordering> {
let xb = 64 - u64::from(x.leading_zeros());
compare_power_scaled(x << (64 - xb), xb, k, top, tb)
}
fn compare_power_scaled(xm: u64, xb: u64, k: u64, tm: u64, tb: u64) -> Option<Ordering> {
debug_assert!(xm >> 63 == 1 && tm >> 63 == 1 && k >= 1);
let mut m = xm;
let mut e = xb as i64 - 64;
let step = |m: u64, e: i64, y: u64, ey: i64| -> (u64, i64) {
let p = u128::from(m) * u128::from(y);
let s = ((p >> 127) as u32) ^ 1;
(((p << s) >> 64) as u64, e + ey + 64 - i64::from(s))
};
let ex = e;
for i in (0..63 - k.leading_zeros()).rev() {
(m, e) = step(m, e, m, e);
if (k >> i) & 1 == 1 {
(m, e) = step(m, e, xm, ex);
}
}
let et = tb as i64 - 64;
let low = e.min(et);
let (dp, dt) = (e - low, et - low);
if dp > 2 || dt > 2 {
return Some(e.cmp(&et));
}
let margin = 16 * u128::from(k) + 2;
let (p_lo, p_hi) = (u128::from(m) << dp, (u128::from(m) + margin) << dp);
let width = u128::from(tb > 64);
let (t_lo, t_hi) = (u128::from(tm) << dt, (u128::from(tm) + width) << dt);
if p_hi < t_lo {
Some(Ordering::Less)
} else if p_lo >= t_hi && (width == 1 || p_lo > t_lo) {
Some(Ordering::Greater)
} else {
None
}
}
unsafe fn normalized(x: *const u64, mut n: usize) -> usize {
unsafe {
while n > 0 && *x.add(n - 1) == 0 {
n -= 1;
}
}
n
}
unsafe fn bit_len(x: *const u64, n: usize) -> u64 {
if n == 0 {
return 0;
}
64 * n as u64 - u64::from(unsafe { *x.add(n - 1) }.leading_zeros())
}
unsafe fn compare(x: *const u64, xn: usize, y: *const u64, yn: usize) -> Ordering {
match xn.cmp(&yn) {
Ordering::Equal => unsafe { cmp_n(x, y, xn) },
other => other,
}
}
unsafe fn shift_into(dst: *mut u64, x: *const u64, xn: usize, sh: i128) -> usize {
unsafe {
if xn == 0 {
return 0;
}
if sh >= 0 {
let (limbs, bits) = (
usize::try_from(sh / 64).unwrap_or(usize::MAX),
(sh % 64) as u32,
);
if limbs >= xn {
return 0;
}
let n = xn - limbs;
if bits == 0 {
core::ptr::copy(x.add(limbs), dst, n);
} else {
rshift(dst, x.add(limbs), n, bits);
}
normalized(dst, n)
} else {
let sh = -sh;
let (limbs, bits) = ((sh / 64) as usize, (sh % 64) as u32);
let top = lshift(dst.add(limbs), x, xn, bits);
*dst.add(limbs + xn) = top;
core::ptr::write_bytes(dst, 0, limbs);
normalized(dst, limbs + xn + 1)
}
}
}
fn scratch_of(scratch: &mut Vec<u64>, n: usize) -> *mut u64 {
if scratch.len() < n {
scratch.resize(n, 0);
}
scratch.as_mut_ptr()
}
unsafe fn mul_into(
r: *mut u64,
x: *const u64,
xn: usize,
y: *const u64,
yn: usize,
scratch: &mut Vec<u64>,
) -> usize {
let ((a, an), (b, bn)) = if xn >= yn {
((x, xn), (y, yn))
} else {
((y, yn), (x, xn))
};
let t = scratch_of(scratch, mul_itch(an, bn));
unsafe {
mul(r, a, an, b, bn, t);
normalized(r, an + bn)
}
}
#[allow(clippy::too_many_arguments)]
unsafe fn power_into(
x: *const u64,
xn: usize,
e: u64,
a: *mut u64,
b: *mut u64,
cap: usize,
scratch: &mut Vec<u64>,
) -> (*mut u64, usize) {
unsafe {
let (mut cur, mut other) = (a, b);
core::ptr::copy_nonoverlapping(x, cur, xn);
let mut n = xn;
for i in (0..63 - e.leading_zeros()).rev() {
debug_assert!(2 * n <= cap, "a power outgrew its buffer");
let t = scratch_of(scratch, sqr_itch(n));
sqr(other, cur, n, t);
n = normalized(other, 2 * n);
core::mem::swap(&mut cur, &mut other);
if (e >> i) & 1 == 1 {
debug_assert!(n + xn <= cap, "a power outgrew its buffer");
n = mul_into(other, cur, n, x, xn, scratch);
core::mem::swap(&mut cur, &mut other);
}
}
(cur, n)
}
}
#[allow(clippy::too_many_arguments)]
unsafe fn quotient_into(
q: *mut u64,
n: *const u64,
nn: usize,
d: *const u64,
dn: usize,
nbuf: *mut u64,
dbuf: *mut u64,
scratch: &mut Vec<u64>,
) -> usize {
unsafe {
if compare(n, nn, d, dn) == Ordering::Less {
return 0;
}
let s = (*d.add(dn - 1)).leading_zeros();
lshift(dbuf, d, dn, s);
*nbuf.add(nn) = lshift(nbuf, n, nn, s);
let qn = match dn {
1 => {
div_qr_1::<true>(q, nbuf, nn, *nbuf.add(nn), *dbuf, invert_limb(*dbuf));
nn
}
2 => {
let v = invert_3by2(*dbuf.add(1), *dbuf);
div_qr_2(q, nbuf, nn + 1, *dbuf.add(1), *dbuf, v);
nn - 1
}
_ => {
let len = if *nbuf.add(nn) != 0 { nn + 1 } else { nn };
let t = scratch_of(scratch, qr_itch(len, dn));
let v = invert_3by2(*dbuf.add(dn - 1), *dbuf.add(dn - 2));
*q.add(len - dn) = u64::from(div_qr(q, nbuf, len, dbuf, dn, v, t));
len - dn + 1
}
};
normalized(q, qn)
}
}
pub(in super::super) fn quotient(n: &[u64], d: &[u64]) -> Vec<u64> {
let mut q = vec![0u64; n.len().max(1)];
let mut nbuf = vec![0u64; n.len() + 1];
let mut dbuf = vec![0u64; d.len()];
let mut scratch = Vec::new();
let qn = unsafe {
quotient_into(
q.as_mut_ptr(),
n.as_ptr(),
n.len(),
d.as_ptr(),
d.len(),
nbuf.as_mut_ptr(),
dbuf.as_mut_ptr(),
&mut scratch,
)
};
q.truncate(qn);
q
}