use super::super::mul::fft::{Cyclic, Transform, Twiddles, cyclic, mul_cyclic, transform};
use super::super::mul::limbs::{add_1, add_n, cmp_n, mul, mul_itch, sub_1, sub_n};
use super::limbs::{DC_DIV, dc_itch, div_qr_classic, invert_3by2};
use core::cmp::Ordering;
use core::ptr;
pub(in super::super) const MU_DIV: usize = if cfg!(miri) { 24 } else { 128 };
pub(in super::super) const MU_QUOTIENT: usize = if cfg!(miri) { 40 } else { 768 };
const MU_PREPARED: usize = if cfg!(miri) { 40 } else { 256 };
const INV_NEWTON: usize = if cfg!(miri) { 8 } else { 80 };
const CYCLIC_MIN: usize = if cfg!(miri) { 40 } else { 300 };
const CYCLIC_MIN_SHARED: usize = if cfg!(miri) { 30 } else { 80 };
#[inline]
pub(in super::super) fn mu_prepared(nn: usize, dn: usize) -> bool {
dn >= MU_PREPARED && nn >= dn + MU_PREPARED.max(3)
}
#[inline]
pub(in super::super) fn mu(nn: usize, dn: usize) -> bool {
dn >= MU_DIV && nn >= dn + MU_QUOTIENT.max(3) && 8 * (nn - dn) >= dn
}
#[derive(Clone, Copy)]
enum Modulus {
Cyclic(Cyclic),
Fold(usize),
}
impl Modulus {
fn of(min: usize, shorter: usize, shared: bool) -> Modulus {
let from = if shared {
CYCLIC_MIN_SHARED
} else {
CYCLIC_MIN
};
if min >= from
&& let Some(p) = cyclic(min, shorter)
{
Modulus::Cyclic(p)
} else {
Modulus::Fold(min)
}
}
fn limbs(self) -> usize {
match self {
Modulus::Cyclic(p) => p.limbs(),
Modulus::Fold(m) => m,
}
}
}
struct Ready {
x: *const u64,
n: usize,
t: Option<Transform>,
}
impl Ready {
unsafe fn of(x: *const u64, n: usize, m: Modulus, tw: Option<&Twiddles>) -> Ready {
let t = match m {
Modulus::Cyclic(p) => Some(unsafe { transform(x, n, p, tw) }),
Modulus::Fold(_) => None,
};
Ready { x, n, t }
}
fn side(&self) -> Side<'_> {
Side {
x: self.x,
n: self.n,
t: self.t.as_ref(),
}
}
}
#[derive(Clone, Copy)]
struct Side<'a> {
x: *const u64,
n: usize,
t: Option<&'a Transform>,
}
unsafe fn mul_mod(
r: *mut u64,
m: Modulus,
x: *const u64,
nx: usize,
y: Side<'_>,
tw: Option<&Twiddles>,
) {
unsafe {
match m {
Modulus::Cyclic(p) => {
let t = y.t.expect("an operand transformed for its modulus");
mul_cyclic(r, p, x, nx, t, tw);
}
Modulus::Fold(mm) => {
let (a, na, b, nb) = if nx >= y.n {
(x, nx, y.x, y.n)
} else {
(y.x, y.n, x, nx)
};
let pn = na + nb;
let mut t = Vec::<u64>::with_capacity(pn + mul_itch(na, nb));
let t = t.as_mut_ptr();
mul(t, a, na, b, nb, t.add(pn));
if pn <= mm {
ptr::copy_nonoverlapping(t, r, pn);
ptr::write_bytes(r.add(pn), 0, mm - pn);
} else {
ptr::copy_nonoverlapping(t, r, mm);
add_mod(r, mm, t.add(mm), pn - mm, 0);
}
}
}
}
}
unsafe fn add_mod(r: *mut u64, m: usize, y: *const u64, ny: usize, at: usize) {
debug_assert!(at < m && ny <= m);
unsafe {
let low = ny.min(m - at);
let mut carry = u64::from(add_n(r.add(at), r.add(at), y, low));
carry = add_1(r.add(at + low), m - at - low, carry);
if ny > low {
let c = u64::from(add_n(r, r, y.add(low), ny - low));
carry += add_1(r.add(ny - low), m - (ny - low), c);
}
while carry != 0 {
carry = add_1(r, m, carry);
}
}
}
unsafe fn sub_power(r: *mut u64, m: usize, e: usize) {
unsafe {
if sub_1(r.add(e), m - e, 1) != 0 {
sub_1(r, m, 1);
}
}
}
unsafe fn all_ones(x: *const u64, n: usize) -> bool {
(0..n).all(|i| unsafe { *x.add(i) } == u64::MAX)
}
pub(super) unsafe fn invert(x: *mut u64, a: *const u64, n: usize, tw: Option<&Twiddles>) {
debug_assert!(n >= 3 && unsafe { *a.add(n - 1) } >> 63 == 1);
unsafe {
if n < INV_NEWTON.max(6) {
invert_base(x, a, n);
return;
}
let l = (n - 1) / 2;
let h = n - l;
invert(x.add(l), a.add(l), h, tw);
newton_step(x, a, n, h, tw);
}
}
fn step_modulus(n: usize) -> Modulus {
let h = n - (n - 1) / 2;
Modulus::of((n + 1).max(2 * h + 1), h, false)
}
unsafe fn invert_base(x: *mut u64, a: *const u64, n: usize) {
let mut u = vec![u64::MAX; 2 * n];
let mut t = Vec::<u64>::with_capacity(if n >= DC_DIV { dc_itch(2 * n, n) } else { 0 });
unsafe {
let v = invert_3by2(*a.add(n - 1), *a.add(n - 2));
let qh = div_qr_classic(x, u.as_mut_ptr(), 2 * n, a, n, v, t.as_mut_ptr());
debug_assert!(qh, "the inverse is above B^n");
}
}
unsafe fn newton_step(x: *mut u64, a: *const u64, n: usize, h: usize, tw: Option<&Twiddles>) {
let l = n - h;
debug_assert!(h > l);
let m = step_modulus(n);
let mm = m.limbs();
let mut t = Vec::<u64>::with_capacity(mm);
let mut e = vec![0u64; n + 1];
unsafe {
let t = t.as_mut_ptr();
let e = e.as_mut_ptr();
let ih = x.add(l);
let ready = Ready::of(ih, h, m, tw);
mul_mod(t, m, a, n, ready.side(), tw);
add_mod(t, mm, a, n, h);
sub_power(t, mm, (n + h) % mm);
let positive = ((n + 1..mm).all(|i| *t.add(i) == 0) && *t.add(n) <= 1) || all_ones(t, mm);
let mut c = 0;
if positive {
if !all_ones(t, mm) {
ptr::copy_nonoverlapping(t, e, n + 1);
}
loop {
sub_1(ih, h, 1);
c += 1;
let borrow = sub_n(e, e, a, n);
if sub_1(e.add(n), 1, u64::from(borrow)) != 0 {
break;
}
}
debug_assert!(c <= 4);
for i in 0..=n {
*e.add(i) = !*e.add(i);
}
add_1(e, n + 1, 1);
} else {
for i in 0..=n {
*e.add(i) = !*t.add(i);
}
debug_assert!((n + 1..mm).all(|i| *t.add(i) == u64::MAX));
}
debug_assert!(*e.add(n) <= 1);
let tm = e.add(l);
mul_mod(t, m, tm, h + 1, ready.side(), tw);
if matches!(m, Modulus::Fold(_)) {
c = 0;
}
for _ in 0..c {
let borrow = sub_n(t, t, tm, h + 1);
sub_1(t.add(h + 1), h, u64::from(borrow));
}
let carry = add_n(t.add(h), t.add(h), tm, h + 1);
debug_assert!(!carry && (2 * h + 1..mm).all(|i| *t.add(i) == 0));
ptr::copy_nonoverlapping(t.add(2 * h - l), x, l);
let out = add_1(ih, h, *t.add(2 * h));
debug_assert!(out == 0, "the inverse is below 2 B^n");
}
}
unsafe fn top_inverse(x: *mut u64, d: *const u64, dn: usize, k: usize, tw: Option<&Twiddles>) {
let mut a = Vec::<u64>::with_capacity(k);
unsafe {
let a = a.as_mut_ptr();
ptr::copy_nonoverlapping(d.add(dn - k), a, k);
if add_1(a, k, 1) != 0 {
ptr::write_bytes(x, 0, k);
} else {
invert(x, a, k, tw);
}
}
}
fn blocks(qn: usize, dn: usize, rem: bool) -> (usize, usize) {
let shape = |b: usize| {
let inn = if rem {
qn.div_ceil(b)
} else {
(qn + 1).div_ceil(b)
};
let inn = inn.max(3).min(dn);
(inn, qn.div_ceil(inn))
};
let first = qn.div_ceil(dn).max(1);
let mut best: Option<((usize, usize), u64)> = None;
for b in first..first + 4 {
let (inn, b) = shape(b);
if !rem && b * inn == qn && best.is_some() {
continue;
}
let Some(cost) = model(dn, inn, b, rem) else {
return shape(first.max(if 3 * qn > dn { 2 } else { 1 }));
};
if best.is_none_or(|(_, c)| cost < c) {
best = Some(((inn, b), cost));
}
}
best.map(|(s, _)| s).unwrap_or_else(|| shape(first))
}
fn model(dn: usize, inn: usize, b: usize, rem: bool) -> Option<u64> {
let plan = |min: usize, shorter: usize| cyclic(min, shorter);
let mut cost = 0;
let mut n = inn;
while n >= INV_NEWTON.max(6) {
let h = n - (n - 1) / 2;
let p = plan((n + 1).max(2 * h + 1), h)?;
cost += p.transform_cost() + 2 * p.cost(true);
n = h;
}
let q = plan(2 * inn, inn)?;
cost += q.transform_cost() + b as u64 * q.cost(true);
let remainders = if rem { b } else { b - 1 };
if remainders > 0 {
let r = plan(dn + 1, inn)?;
cost += r.transform_cost() + remainders as u64 * r.cost(true);
}
Some(cost)
}
pub(super) unsafe fn mu_div_qr(
q: *mut u64,
n: *mut u64,
nn: usize,
d: *const u64,
dn: usize,
) -> bool {
unsafe { mu_div(q, n, nn, d, dn, true) }
}
pub(super) unsafe fn mu_div_q(
q: *mut u64,
n: *mut u64,
nn: usize,
d: *const u64,
dn: usize,
) -> bool {
unsafe { mu_div(q, n, nn, d, dn, false) }
}
fn twiddles_for(moduli: &[Modulus]) -> Option<Twiddles> {
let plans: Vec<Cyclic> = moduli
.iter()
.filter_map(|m| match m {
Modulus::Cyclic(p) => Some(*p),
Modulus::Fold(_) => None,
})
.collect();
(!plans.is_empty()).then(|| Twiddles::of(&plans))
}
unsafe fn mu_div(q: *mut u64, n: *mut u64, nn: usize, d: *const u64, dn: usize, rem: bool) -> bool {
let qn = nn - dn;
let (inn, b) = blocks(qn, dn, rem);
let mq = Modulus::of(2 * inn, inn, b > 1);
let mr = Modulus::of(dn + 1, inn, b > 2 || (rem && b > 1));
let twiddles = twiddles_for(&[mq, mr, step_modulus(inn)]);
let tw = twiddles.as_ref();
let mut inverse = Vec::<u64>::with_capacity(inn);
unsafe {
let inv = inverse.as_mut_ptr();
top_inverse(inv, d, dn, inn, tw);
let ready_inv = Ready::of(inv, inn, mq, tw);
let ready_d = (rem || b > 1).then(|| Ready::of(d, dn, mr, tw));
let blocks = Blocks {
inn,
b,
mq,
inv: ready_inv.side(),
mr,
d: ready_d.as_ref().map(Ready::side),
tw,
};
run(q, n, nn, d, dn, &blocks, rem)
}
}
#[derive(Clone)]
pub(in super::super) struct Prepared {
inv: Vec<u64>,
mq: Modulus,
t_inv: Option<Transform>,
mr: Modulus,
t_d: Option<Transform>,
}
pub(in super::super) unsafe fn prepare(d: *const u64, dn: usize) -> Prepared {
let mq = Modulus::of(2 * dn, dn, true);
let mr = Modulus::of(dn + 1, dn, true);
let twiddles = twiddles_for(&[mq, mr, step_modulus(dn)]);
let tw = twiddles.as_ref();
let mut inv = vec![0; dn];
unsafe {
top_inverse(inv.as_mut_ptr(), d, dn, dn, tw);
let t_inv = Ready::of(inv.as_ptr(), dn, mq, tw).t;
let t_d = Ready::of(d, dn, mr, tw).t;
Prepared {
inv,
mq,
t_inv,
mr,
t_d,
}
}
}
#[allow(clippy::too_many_arguments)]
pub(in super::super) unsafe fn mu_div_prepared(
q: *mut u64,
n: *mut u64,
nn: usize,
d: *const u64,
dn: usize,
v: u64,
pre: &Prepared,
rem: bool,
) -> bool {
let qn = nn - dn;
debug_assert!(pre.inv.len() == dn);
let d_side = Side {
x: d,
n: dn,
t: pre.t_d.as_ref(),
};
unsafe {
if 3 * qn >= 2 * dn {
let twiddles = twiddles_for(&[pre.mq, pre.mr]);
let r = qn % dn;
let r = if 4 * r < dn { r } else { 0 };
let blocks = Blocks {
inn: dn,
b: (qn - r).div_ceil(dn),
mq: pre.mq,
inv: Side {
x: pre.inv.as_ptr(),
n: dn,
t: pre.t_inv.as_ref(),
},
mr: pre.mr,
d: Some(d_side),
tw: twiddles.as_ref(),
};
let qh = run(q.add(r), n.add(r), nn - r, d, dn, &blocks, rem || r > 0);
if r > 0 {
let mut t = Vec::<u64>::with_capacity(dc_itch(dn + r, dn));
let high = div_qr_classic(q, n, dn + r, d, dn, v, t.as_mut_ptr());
debug_assert!(!high, "a partial remainder is below the divisor");
}
return qh;
}
let inn = if rem { qn } else { qn + 1 };
let mq = Modulus::of(2 * inn, inn, false);
let twiddles = twiddles_for(&[mq, pre.mr]);
let tw = twiddles.as_ref();
let ready_inv = Ready::of(pre.inv.as_ptr().add(dn - inn), inn, mq, tw);
let blocks = Blocks {
inn,
b: 1,
mq,
inv: ready_inv.side(),
mr: pre.mr,
d: Some(d_side),
tw,
};
run(q, n, nn, d, dn, &blocks, rem)
}
}
struct Blocks<'a> {
inn: usize,
b: usize,
mq: Modulus,
inv: Side<'a>,
mr: Modulus,
d: Option<Side<'a>>,
tw: Option<&'a Twiddles>,
}
unsafe fn run(
q: *mut u64,
n: *mut u64,
nn: usize,
d: *const u64,
dn: usize,
blocks: &Blocks<'_>,
rem: bool,
) -> bool {
let qn = nn - dn;
let (inn, b, mq, mr, tw) = (blocks.inn, blocks.b, blocks.mq, blocks.mr, blocks.tw);
debug_assert!(inn <= dn && (b - 1) * inn < qn && qn <= b * inn);
unsafe {
let top = n.add(qn);
let qh = cmp_n(top, d, dn) != Ordering::Less;
if qh {
sub_n(top, top, d, dn);
}
let mut p = Vec::<u64>::with_capacity(mq.limbs().max(mr.limbs()));
let p = p.as_mut_ptr();
let mut w = Vec::<u64>::with_capacity(mr.limbs());
let w = w.as_mut_ptr();
let mut fresh = None;
let mut pos = qn;
for block in 0..b {
let k = if block + 1 < b {
inn
} else {
qn - (b - 1) * inn
};
pos -= k;
let wt = n.add(pos + dn + k - inn);
mul_mod(p, mq, wt, inn, blocks.inv, tw);
let qe = p.add(inn);
let carry = add_n(qe, qe, wt, inn);
debug_assert!(!carry, "the estimate is below B^in");
ptr::copy_nonoverlapping(qe.add(inn - k), q.add(pos), k);
if !rem && block + 1 == b && k < inn {
let g = inn - k;
if !(all_ones(qe.add(1), g - 1) && *qe > u64::MAX - 9) {
break;
}
}
let side = match blocks.d {
Some(side) => side,
None => fresh.get_or_insert_with(|| Ready::of(d, dn, mr, tw)).side(),
};
remainder(n.add(pos), dn, k, q.add(pos), d, mr, side, w, p, tw);
}
debug_assert!(pos == 0 || !rem);
qh
}
}
#[allow(clippy::too_many_arguments)]
unsafe fn remainder(
wp: *mut u64,
dn: usize,
k: usize,
qk: *mut u64,
d: *const u64,
m: Modulus,
side: Side<'_>,
w: *mut u64,
p: *mut u64,
tw: Option<&Twiddles>,
) {
let mm = m.limbs();
debug_assert!(mm > dn);
unsafe {
mul_mod(p, m, qk, k, side, tw);
let wn = dn + k;
if wn <= mm {
ptr::copy_nonoverlapping(wp, w, wn);
ptr::write_bytes(w.add(wn), 0, mm - wn);
} else {
ptr::copy_nonoverlapping(wp, w, mm);
add_mod(w, mm, wp.add(mm), wn - mm, 0);
}
if sub_n(w, w, p, mm) {
sub_1(w, mm, 1);
}
if all_ones(w, mm) {
ptr::write_bytes(w, 0, mm);
}
debug_assert!((dn + 1..mm).all(|i| *w.add(i) == 0) && *w.add(dn) < 10);
let mut top = *w.add(dn);
while top != 0 || cmp_n(w, d, dn) != Ordering::Less {
top -= u64::from(sub_n(w, w, d, dn));
let out = add_1(qk, k, 1);
debug_assert!(out == 0, "the block's quotient is below B^k");
}
ptr::copy_nonoverlapping(w, wp, dn);
}
}