#include #include #include #include #include #include namespace sangi {

// ========================================================================
// main sqrt function
// ========================================================================

Int IntSqrt::sqrt(const Int& value) {
    // Handle special states
    if (value.isNaN()) [[unlikely]] {
        return Int::NaN();
    }

    if (value.isNegative()) [[unlikely]] {
        // The square root of a negative number is NaN
        Int result = Int::NaN();
        result.setState(NumericState::NaN, NumericError::NegativeSqrt);
        return result;
    }

    if (value.getState() == NumericState::PositiveInfinity) [[unlikely]] {
        return Int::PositiveInfinity();
    }

    if (value.isZero()) [[unlikely]] {
        return Int::Zero();
    }

    if (value.isOne()) [[unlikely]] {
        return Int::One();
    }

    // 1 limb fast path: compute isqrt(uint64_t) directly
    size_t an = value.size();
    if (an == 1) {
        uint64_t v = value.word(0);
        uint64_t s = static_cast (std::sqrt(static_cast (v)));
        // Newton correction (double rounding error)
        if (s > 0 && s * s > v) s--;
        if ((s + 1) * (s + 1) <= v) s++;
        return Int(s);
    }

    // 2+ limbs: mpn-level Newton sqrt
    {
        // Reference value.data() directly (no copy needed — sqrtrem does not modify ap)
        const uint64_t* ap = value.data();
        size_t sn = (an + 1) / 2;

        // Use a stack buffer to avoid arena access for small sizes
        constexpr size_t STACK_LIMIT = 64;
        uint64_t sp_stack[STACK_LIMIT];
        size_t scratch_sz = mpn::sqrtrem_scratch_size(an);
        uint64_t scratch_stack[STACK_LIMIT * 4];

        uint64_t* sp;
        uint64_t* scratch;
        ScratchScope scope;
        if (sn <= STACK_LIMIT && scratch_sz <= STACK_LIMIT * 4) {
            sp = sp_stack;
            scratch = scratch_stack;
        } else {
            sp = getThreadArena().alloc_limbs(sn);
            scratch = getThreadArena().alloc_limbs(scratch_sz);
        }

        size_t result_n = mpn::sqrtrem(sp, nullptr, ap, an, scratch);
        if (result_n == 0) return Int::Zero();
        return Int::fromRawWords(std::span (sp, result_n), 1);
    }
}

// ========================================================================
// sqrt_small: square-root computation in double precision
// ========================================================================

Int IntSqrt::sqrt_small(const Int& value) {
    // When MSB < 52 bits, it can be computed exactly with a double
    double x = value.toDouble();
    double sqrt_x = std::sqrt(x);
    uint64_t result = static_cast (sqrt_x);
    return Int(result);
}

// ========================================================================
// sqrtRem: return the square root and the remainder
// ========================================================================

Int IntSqrt::sqrtRem(const Int& value, Int& remainder) {
    // Handle special states
    if (value.isNaN()) {
        remainder = Int::NaN();
        return Int::NaN();
    }

    if (value.isNegative()) {
        Int result = Int::NaN();
        result.setState(NumericState::NaN, NumericError::NegativeSqrt);
        remainder = Int::NaN();
        return result;
    }

    if (value.getState() == NumericState::PositiveInfinity) {
        remainder = Int::NaN();  // remainder is undefined
        return Int::PositiveInfinity();
    }

    if (value.isZero()) {
        remainder = Int::Zero();
        return Int::Zero();
    }

    if (value.isOne()) {
        remainder = Int::Zero();
        return Int::One();
    }

    // Compute the square root and remainder together with mpn-level sqrtrem
    size_t bit_length = value.bitLength();
    if (bit_length < 52) {
        Int s = sqrt_small(value);
        // s is built via double → Normal, and value is normalized at the entry → Unchecked is safe
        Int s_sq; IntOps::squareUnchecked(s, s_sq);
        remainder = value - s_sq;
        return s;
    }

    {
        size_t an = value.size();
        ScratchScope scope;
        uint64_t* ap = getThreadArena().alloc_limbs(an);
        for (size_t i = 0; i < an; i++) ap[i] = value.word(i);

        size_t sn = (an + 1) / 2;
        uint64_t* sp = getThreadArena().alloc_limbs(sn);
        uint64_t* rem_p = getThreadArena().alloc_limbs(an);
        std::memset(rem_p, 0, an * sizeof(uint64_t));
        size_t scratch_sz = mpn::sqrtrem_scratch_size(an);
        uint64_t* scratch = getThreadArena().alloc_limbs(scratch_sz);

        size_t result_n = mpn::sqrtrem(sp, rem_p, ap, an, scratch);

        size_t rem_n = mpn::normalized_size(rem_p, an);
        if (rem_n == 0) {
            remainder = Int::Zero();
        } else {
            remainder = Int::fromRawWords(std::span (rem_p, rem_n), 1);
        }

        if (result_n == 0) return Int::Zero();
        return Int::fromRawWords(std::span (sp, result_n), 1);
    }
}

// ========================================================================
// isSquare: perfect-square detection
// A 3-layer filter learned from GMP perfsqr.c:
//   Filter 1: up[0] % 256 bit table (O(1), 82.81% rejected)
//   Filter 2: mod_34lsub1 + quadratic-residue check with small primes (O(n), ~97.81% additionally rejected)
//   Filter 3: final verification with sqrtRem
//   Total: 99.62% of non-squares rejected before sqrtrem
// ========================================================================

// ================================================================
// Compile-time table generation via consteval
// ================================================================

// Generate the quadratic-residue bitmask mod p at compile time (p < 64)
consteval uint64_t make_qr_mask(unsigned p) {
    uint64_t mask = 0;
    for (unsigned i = 0; i < p; i++)
        mask |= 1ULL << ((i * i) % p);
    return mask;
}

// Generate the quadratic-residue bitmask mod p at compile time (p >= 64, 2 limb)
struct QRMask128 { uint64_t lo; uint64_t hi; };
consteval QRMask128 make_qr_mask_wide(unsigned p) {
    uint64_t lo = 0, hi = 0;
    for (unsigned i = 0; i < p; i++) {
        unsigned r = (i * i) % p;
        if (r < 64) lo |= 1ULL << r;
        else        hi |= 1ULL << (r - 64);
    }
    return {lo, hi};
}

// Quadratic-residue bit table for up[0] % 256 (equivalent to sq_res_0x100 in GMP perfsqr.h)
// A bit of 1 means it may be a quadratic residue; 0 means it is definitely not a square
consteval std::array make_sq_res_256() {
    std::array t{};
    for (unsigned i = 0; i < 256; i++) {
        unsigned r = (i * i) & 0xFF;
        t[r / 64] |= 1ULL << (r % 64);
    }
    return t;
}
static constexpr auto SQ_RES_256 = make_sq_res_256();

// Quadratic-residue mask for each prime (compile-time verifiable)
static constexpr uint64_t QR_9  = make_qr_mask(9);   // rejection rate 55.56%
static constexpr uint64_t QR_5  = make_qr_mask(5);   // rejection rate 40%
static constexpr uint64_t QR_7  = make_qr_mask(7);   // rejection rate 42.86%
static constexpr uint64_t QR_13 = make_qr_mask(13);  // rejection rate 46.15%
static constexpr uint64_t QR_17 = make_qr_mask(17);  // rejection rate 47.06%
static constexpr auto QR_97 = make_qr_mask_wide(97);  // rejection rate 49.48%

// ★ GMP-style packed composite table (fewer % operations)
// Check the quadratic residues of 7*13=91 and 5*17=85 at once with a packed bit table
static constexpr auto QR_91 = make_qr_mask_wide(91);  // 7*13
static constexpr auto QR_85 = make_qr_mask_wide(85);  // 5*17

// ★ Additional filter: also test small primes other than mod_34lsub1
// Add 11, 19, 23, 29, 31 (further improves the overall rejection rate)
static constexpr uint64_t QR_11 = make_qr_mask(11);
static constexpr uint64_t QR_19 = make_qr_mask(19);
static constexpr uint64_t QR_23 = make_qr_mask(23);
static constexpr uint64_t QR_29 = make_qr_mask(29);
static constexpr uint64_t QR_31 = make_qr_mask(31);

// Quadratic-residue check of the mod_34lsub1 result r with small primes
// ★ GMP-style packed test: % from 6 times → 4 times + 5 additional primes
static bool check_sq_residue_mod34(uint64_t r) {
    constexpr uint64_t M48 = (1ULL << 48) - 1;
    r = (r & M48) + (r >> 48);

    // GMP-style packed: 91=7*13, 85=5*17 covers 4 primes with 2 % operations
    if (((QR_9  >> (r % 9))  & 1) == 0) return false;
    { uint64_t r91 = r % 91;
      if (r91 < 64) { if (((QR_91.lo >> r91) & 1) == 0) return false; }
      else { if (((QR_91.hi >> (r91 - 64)) & 1) == 0) return false; }
    }
    { uint64_t r85 = r % 85;
      if (r85 < 64) { if (((QR_85.lo >> r85) & 1) == 0) return false; }
      else { if (((QR_85.hi >> (r85 - 64)) & 1) == 0) return false; }
    }
    { uint64_t r97 = r % 97;
      if (r97 < 64) { if (((QR_97.lo >> r97) & 1) == 0) return false; }
      else { if (((QR_97.hi >> (r97 - 64)) & 1) == 0) return false; }
    }
    return true;
}

// ★ Additional filter: check small primes directly from word(0) (mod_34lsub1 not needed)
[[maybe_unused]] static bool check_sq_residue_extra(uint64_t w0) {
    if (((QR_11 >> (w0 % 11)) & 1) == 0) return false;
    if (((QR_19 >> (w0 % 19)) & 1) == 0) return false;
    if (((QR_23 >> (w0 % 23)) & 1) == 0) return false;
    if (((QR_29 >> (w0 % 29)) & 1) == 0) return false;
    if (((QR_31 >> (w0 % 31)) & 1) == 0) return false;
    return true;
}

bool IntSqrt::isSquare(const Int& value, Int* pSqrt) {
    // Handle special states
    if (value.isNaN()) [[unlikely]] {
        if (pSqrt) *pSqrt = Int::NaN();
        return false;
    }

    if (value.isInfinite()) [[unlikely]] {
        if (pSqrt) *pSqrt = Int::NaN();
        return false;
    }

    // Negative numbers are not squares
    if (value.isNegative()) [[unlikely]] {
        return false;
    }

    // 0 and 1 are squares
    if (value.isZero()) [[unlikely]] {
        if (pSqrt) *pSqrt = Int::Zero();
        return true;
    }

    if (value.isOne()) [[unlikely]] {
        if (pSqrt) *pSqrt = Int::One();
        return true;
    }

    // Filter 1: up[0] % 256 bit table (O(1), 82.81% rejected)
    {
        unsigned idx = static_cast (value.word(0) & 0xFF);
        if (((SQ_RES_256[idx / 64] >> (idx % 64)) & 1) == 0)
            return false;
    }

    // Filter 1b: removed — check_sq_residue_extra uses word(0) % p, but
    // for multi-word numbers value % p ≠ word(0) % p, so false negatives occur

    // Filter 2: mod_34lsub1-based quadratic-residue filter (O(n), GMP-style packed test)
    {
        size_t vn = value.size();
        constexpr size_t BUF = 128;
        uint64_t stack_buf[BUF];
        uint64_t* wp = (vn <= BUF) ? stack_buf : new uint64_t[vn];
        for (size_t i = 0; i < vn; i++) wp[i] = value.word(i);
        uint64_t r = mpn::mod_34lsub1(wp, vn);
        if (wp != stack_buf) delete[] wp;
        if (!check_sq_residue_mod34(r))
            return false;
    }

    // Filter 3: final verification with mpn::sqrtrem_check_exact
    // ★ Codex note: use the bool-only path instead of sqrtRem (Int construction)
    {
        size_t an = value.size();
        ScratchScope scope;
        uint64_t* ap = getThreadArena().alloc_limbs(an);
        for (size_t i = 0; i < an; i++) ap[i] = value.word(i);
        size_t sn = (an + 1) / 2;
        uint64_t* sp = getThreadArena().alloc_limbs(sn);
        size_t scratch_sz = mpn::sqrtrem_scratch_size(an);
        uint64_t* scratch = getThreadArena().alloc_limbs(scratch_sz);
        auto [result_n, is_exact] = mpn::sqrtrem_check_exact(sp, ap, an, scratch);
        if (pSqrt && is_exact && result_n > 0) {
            *pSqrt = Int::fromRawWords(std::span (sp, result_n), 1);
        }
        return is_exact;
    }
}


// ★ Direct mpn power: compute pow(x, n-1) faster than Int::pow
// n=3: x^2, n=5: x^4, n=7: x^6 built with an addition chain
// Return value: {x^(n-1), x^n} (x^n is optional)
// ★ Direct-limb-return version: returns the pointer from the allocator as-is (no copy)
// The caller is responsible for the buffer's lifetime
struct MpnPowResult {
    const uint64_t* nm1; size_t nm1_n;  // x^(n-1)
    const uint64_t* pn;  size_t pn_n;   // x^n (need_xn only)
};

// ★ Tier 3: templatize the allocator — works with either an arena or a local bump
// final_keep > 0: compute the final x^(n-1) with square_high truncation (only the top final_keep limbs are accuracy-guaranteed).
//   Since divide_q_approx in nthRoot Newton only references the top qn+3 limbs of the divisor,
//   final_keep ≥ qn+3+safety can be much faster than a full square.
//   final_keep = 0: ordinary full square (when called with need_xn=true or from the verification loop).
template static MpnPowResult mpn_pow_for_root_impl(
    const uint64_t* xp, size_t xn_sz, uint32_t n, bool need_xn,
    AllocFn alloc_fn, size_t final_keep = 0)
{
    auto alloc_and_square = [&](const uint64_t* a, size_t an) -> std::pair {
        size_t rn = 2 * an;
        uint64_t* r = alloc_fn(rn);
        uint64_t* sc = alloc_fn(mpn::square_scratch_size(an));
        mpn::square(r, a, an, sc);
        return {r, mpn::normalized_size(r, rn)};
    };
    // Final square: square_high if final_keep > 0 and truncation is profitable.
    // The output is the ordinary full-square buffer (size 2*an); only the top final_keep limbs are accuracy-guaranteed,
    // and the lower part is zero-filled (so divide_q_approx's fallback path does not read garbage).
    auto alloc_and_square_high = [&](const uint64_t* a, size_t an, size_t keep)
        -> std::pair {
        size_t rn = 2 * an;
        uint64_t* r = alloc_fn(rn);
        if (keep == 0 || keep >= rn || an < mpn::BZ_THRESHOLD / 2) {
            // truncation inefficient → full square
            uint64_t* sc = alloc_fn(mpn::square_scratch_size(an));
            mpn::square(r, a, an, sc);
            return {r, mpn::normalized_size(r, rn)};
        }
        // square_high: zero-fill the lower (rn-keep) limbs + compute the top keep limbs with square_high
        std::memset(r, 0, (rn - keep) * sizeof(uint64_t));
        uint64_t* sh_sc = alloc_fn(mpn::square_high_scratch_size(an, keep));
        mpn::square_high(r + (rn - keep), a, an, keep, sh_sc);
        return {r, mpn::normalized_size(r, rn)};
    };
    auto alloc_and_mul = [&](const uint64_t* a, size_t an, const uint64_t* b, size_t bn) -> std::pair {
        size_t rn = an + bn;
        uint64_t* r = alloc_fn(rn);
        uint64_t* sc = alloc_fn(mpn::multiply_scratch_size(an, bn));
        mpn::multiply(r, a, an, b, bn, sc);
        return {r, mpn::normalized_size(r, rn)};
    };

    MpnPowResult res = {nullptr, 0, nullptr, 0};
    if (n == 3) {
        // x² is the final nm1 — full if need_xn, otherwise square_high truncation
        auto [x2, x2n] = need_xn
            ? alloc_and_square(xp, xn_sz)
            : alloc_and_square_high(xp, xn_sz, final_keep);
        res.nm1 = x2; res.nm1_n = x2n;
        if (need_xn) { auto [x3, x3n] = alloc_and_mul(x2, x2n, xp, xn_sz); res.pn = x3; res.pn_n = x3n; }
    } else if (n == 5) {
        auto [x2, x2n] = alloc_and_square(xp, xn_sz);
        // x⁴ is the final nm1 — full if need_xn
        auto [x4, x4n] = need_xn
            ? alloc_and_square(x2, x2n)
            : alloc_and_square_high(x2, x2n, final_keep);
        res.nm1 = x4; res.nm1_n = x4n;
        if (need_xn) { auto [x5, x5n] = alloc_and_mul(x4, x4n, xp, xn_sz); res.pn = x5; res.pn_n = x5n; }
    } else if (n == 7) {
        auto [x2, x2n] = alloc_and_square(xp, xn_sz);
        auto [x3, x3n] = alloc_and_mul(x2, x2n, xp, xn_sz);
        // x⁶ = (x³)² is the final nm1 — full if need_xn
        auto [x6, x6n] = need_xn
            ? alloc_and_square(x3, x3n)
            : alloc_and_square_high(x3, x3n, final_keep);
        res.nm1 = x6; res.nm1_n = x6n;
        if (need_xn) { auto [x7, x7n] = alloc_and_mul(x6, x6n, xp, xn_sz); res.pn = x7; res.pn_n = x7n; }
    } else {
        // General n: compute x^(n-1) with binary exponentiation
        uint32_t e = n - 1;

        // result = x (initial value)
        size_t rn = xn_sz;
        uint64_t* rp = alloc_fn(xn_sz);
        std::memcpy(rp, xp, xn_sz * sizeof(uint64_t));

        int top_bit = 31 - std::countl_zero(e);
        for (int bit = top_bit - 1; bit >= 0; bit--) {
            auto [sq, sqn] = alloc_and_square(rp, rn);
            rp = sq; rn = sqn;

            if (e & (1u << bit)) {
                auto [prod, prodn] = alloc_and_mul(rp, rn, xp, xn_sz);
                rp = prod; rn = prodn;
            }
        }
        res.nm1 = rp; res.nm1_n = rn;
        if (need_xn) {
            auto [xn_p, xn_n] = alloc_and_mul(rp, rn, xp, xn_sz);
            res.pn = xn_p; res.pn_n = xn_n;
        }
    }
    return res;
}

// ★ arena version (outside the PD loop, used from mpn_pow_for_root)
static MpnPowResult mpn_pow_for_root_raw(
    const uint64_t* xp, size_t xn_sz, uint32_t n, bool need_xn)
{
    return mpn_pow_for_root_impl(xp, xn_sz, n, need_xn,
        [](size_t s) { return getThreadArena().alloc_limbs(s); });
}

// ★ Tier 3: compute the maximum scratch size for one PD-loop step
// Includes all buffers for pow(need_xn=false) + division + addition
static size_t pd_step_scratch_size(uint32_t n, size_t max_xn, size_t val_n) {
    auto a8 = [](size_t s) -> size_t { return (s + 7) & ~size_t(7); };
    size_t total = 0;

    // ap_buf (value >> shift)
    total += a8(val_n + 1);

    // ★ 2026-04-28: since square_high (Karatsuba-high truncation) is adopted,
    //    bound the final square's scratch by max(square, square_high fallback).
    //    square_high's fallback path requires 2n + square_scratch_size(n), so
    //    add a 2n margin on the safe side.
    auto sq_or_sqhigh_sc = [](size_t n) -> size_t {
        return 2 * n + mpn::square_scratch_size(n);
    };
    if (n == 3) {
        // x² (square_high or full square)
        total += a8(2 * max_xn);
        total += a8(sq_or_sqhigh_sc(max_xn));
        // quotient
        total += a8(val_n + 2);
        // divide_q_approx scratch (no remainder, no correction multiply)
        if (2 * max_xn >= 2)
            total += a8(mpn::divide_q_approx_scratch_size(val_n, 2 * max_xn));
        // tp (2*x + quot)
        total += a8(val_n + 4);
    } else {
        // pow(x, n-1) — all intermediate results expanded on the bump allocator
        // n=3,5,7 use dedicated paths, general n uses binary exp
        if (n == 5) {
            total += a8(2 * max_xn) + a8(mpn::square_scratch_size(max_xn));     // x² (full)
            total += a8(4 * max_xn) + a8(sq_or_sqhigh_sc(2 * max_xn));          // x⁴ (final, sq_high)
        } else if (n == 7) {
            total += a8(2 * max_xn) + a8(mpn::square_scratch_size(max_xn));                    // x² (full)
            total += a8(3 * max_xn) + a8(mpn::multiply_scratch_size(2 * max_xn, max_xn));      // x³ (full)
            total += a8(6 * max_xn) + a8(sq_or_sqhigh_sc(3 * max_xn));                          // x⁶ (final, sq_high)
        } else {
            // binary exp of (n-1)
            uint32_t e = n - 1;
            int top_bit = 31 - std::countl_zero(e);
            total += a8(max_xn);  // initial copy
            size_t cur = max_xn;
            for (int bit = top_bit - 1; bit >= 0; bit--) {
                total += a8(2 * cur) + a8(mpn::square_scratch_size(cur));
                cur = 2 * cur;
                if (e & (1u << bit)) {
                    total += a8(cur + max_xn) + a8(mpn::multiply_scratch_size(cur, max_xn));
                    cur = cur + max_xn;
                }
            }
        }
        size_t max_pow_sz = (n - 1) * max_xn;
        // quotient
        total += a8(val_n + 2);
        // divide_q_approx scratch (no remainder, no correction multiply)
        if (max_pow_sz >= 2)
            total += a8(mpn::divide_q_approx_scratch_size(val_n, max_pow_sz));
        // tp ((n-1)*x + quot)
        total += a8(val_n + 4);
    }
    return total;
}

// Int-version wrapper (used outside the PD loop)
// Uses mpn_pow_for_root_raw for all n (no Int::pow fallback needed)
static std::pair mpn_pow_for_root(const Int& x, uint32_t n, bool need_xn) {
    if (x.isZero() || x.isOne()) return {x, x};
    ScratchScope scope;
    size_t xn_sz = x.size();
    uint64_t* xbuf = getThreadArena().alloc_limbs(xn_sz);
    for (size_t i = 0; i < xn_sz; i++) xbuf[i] = x.word(i);
    auto r = mpn_pow_for_root_raw(xbuf, xn_sz, n, need_xn);
    Int pm1 = Int::fromRawWords(std::span (r.nm1, r.nm1_n), 1);
    Int pn = (need_xn && r.pn) ? Int::fromRawWords(std::span (r.pn, r.pn_n), 1) : Int(0);
    return {pm1, pn};
}

Int IntSqrt::nthRoot(const Int& value, uint32_t n) {
    return nthRoot_internal(value, n, nullptr);
}

Int IntSqrt::nthRoot_internal(const Int& value, uint32_t n, Int* pRemainder) {
    // Special cases n = 0 or n = 1
    if (n == 0) {
        // n = 0 is an invalid argument
        if (pRemainder) *pRemainder = Int::NaN();
        return Int::NaN();
    }

    if (n == 1) {
        if (pRemainder) *pRemainder = Int::Zero();
        return value;
    }

    // Use sqrt when n = 2
    if (n == 2) {
        if (pRemainder) {
            // Using sqrtRem avoids the redundant pow(root,2)
            return sqrtRem(value, *pRemainder);
        }
        return sqrt(value);
    }

    // Handle special states
    if (value.isNaN()) {
        if (pRemainder) *pRemainder = Int::NaN();
        return Int::NaN();
    }

    // Handle negative values
    if (value.isNegative()) {
        // When n is odd, return the negative root: -root(|value|, n)
        if (n & 1) {
            Int pos_rem;
            Int result = nthRoot_internal(-value, n, pRemainder ? &pos_rem : nullptr);
            if (pRemainder) {
                // value < 0, root = -|root|, root^n = -|root|^n (n odd)
                // remainder = value - root^n = value + |root|^n = -(|value| - |root|^n) = -pos_rem
                *pRemainder = -pos_rem;
            }
            return -result;
        } else {
            // When n is even, return NaN (a negative even root does not exist in the reals)
            Int result = Int::NaN();
            result.setState(NumericState::NaN, NumericError::NegativeSqrt);
            if (pRemainder) *pRemainder = Int::NaN();
            return result;
        }
    }

    if (value.getState() == NumericState::PositiveInfinity) {
        if (pRemainder) *pRemainder = Int::NaN();
        return Int::PositiveInfinity();
    }

    if (value.isZero()) {
        if (pRemainder) *pRemainder = Int::Zero();
        return Int::Zero();
    }

    if (value.isOne()) {
        if (pRemainder) *pRemainder = Int::Zero();
        return Int::One();
    }

    // Note: tried integrating mpn::cbrtrem for n=3 + remainder, but measurements showed
    // 2.9-6.0x slower than GMP (slower than PD Newton + remainder tracking).
    // bench: 2026-04-29 session 16, see PLAN_NTHROOT_FOLLOWUP.md.
    // The integration was not adopted; the PD Newton path is continued.

    // Precision Doubling Newton + floor verification
    //
    // Algorithm overview (Brent & Zimmermann, Modern Computer Arithmetic §1.5.2):
    //   1. Get an initial value of ~53/n bits with a double approximation
    //   2. Precision-doubling Newton: roughly double the number of correct bits each step
    //      precision schedule: c_prev = ceil((c + L) / 2), L = 2*ceil(log2(n)) + 2
    //   3. At the final step, Newton correction with remainder tracking + floor guarantee
    //
    // Complexity:
    //   PD loop: compute pow(x, n-1) at each step. By the property of geometric series,
    //   the total cost of all steps is about twice the final step (≈ 2·M(P)·log(n)).
    //   Post-processing: reuse the final step's pow for Newton correction and floor verification.
    //   Add one verification pow only when the Newton correction Q > 0.

    int B = static_cast (value.bitLength());
    int ni = static_cast (n);

    // When value < 2^n, the root is 1
    if (B <= ni) {
        if (pRemainder) *pRemainder = value - Int::One();
        return Int::One();
    }

    // --- Initial approximation (double-based) ---
    // For a perfect-power input (e.g., 2^603 / n=3), if std::pow/std::exp2 double rounding
    // returns a value slightly smaller than the true value (e.g., 262143.99 vs 262144),
    // truncation leaves ri short by 1, and the subsequent PD Newton can only converge
    // sub-quadratically, causing the final result to overshoot greatly (2026-05-08 fix).
    // Countermeasure: use ceil so the initial value is always at or above the true value. Newton
    // converges quadratically and exactly on the x > root side, so the -1 step of the verify loop resolves the terminal off-by-one.
    Int x;
    {
        int shift = (B > 53) ? (B - 53) : 0;
        double m_d = (shift > 0) ? (value >> shift).toDouble() : value.toDouble();
        int q = shift / ni;
        int r = shift % ni;
        double root_d = std::pow(m_d, 1.0 / n) * std::exp2(static_cast (r) / n);
        // ceil tolerates a slight overestimate (a 1 ULP overshoot is absorbed by Newton).
        int64_t ri = std::max(static_cast (1), static_cast (std::ceil(root_d)));
        x = (q > 0) ? (Int(ri) << q) : Int(ri);
    }
    if (x.isZero()) x = Int::One();

    const Int n_minus_1(static_cast (n - 1));
    const Int n_int(static_cast (n));

    int P = (B + ni - 1) / ni;     // target bit precision of root
    int x_bits = static_cast (x.bitLength());

    // Number of correct bits of the double approximation
    int P0 = std::min(x_bits, std::max(53 / ni - 1, 4));

    if (P > P0 + 10) {
        // === Precision Doubling Newton (Brent-Zimmermann style) ===

        // ceil(log2(n))
        int logk = 0;
        { unsigned tmp = n - 1; while (tmp > 0) { logk++; tmp >>= 1; } }
        int L = 2 * logk + 2;

        // Build the precision schedule backward from P
        int sizes[64];
        int ns = 0;
        sizes[0] = P;
        while (sizes[ns] > P0 && ns < 60) {
            int next = (sizes[ns] + L + 1) / 2;
            if (next >= sizes[ns]) break;
            sizes[ns + 1] = next;
            ns++;
        }
        // Reverse to forward order: sizes[0]=smallest, sizes[ns]=P
        for (int i = 0, j = ns; i < j; i++, j--) {
            std::swap(sizes[i], sizes[j]);
        }

        // Truncate to the initial precision
        if (x_bits > sizes[0]) {
            x = x >> (x_bits - sizes[0]);
        }
        int cur_prec = static_cast (x.bitLength());

        // --- PD Newton loop: step 1 .. ns-1 ---
        // ★ Tier 3: pre-allocate scratch — completely removes the TLS lookup and mark/rewind inside the loop
        {
            ScratchScope pd_scope;
            // Expand x into a limb array
            size_t max_limbs = (P + 63) / 64 + 4;
            uint64_t* xp = getThreadArena().alloc_limbs(max_limbs);
            size_t xn = x.size();
            for (size_t i = 0; i < xn; i++) xp[i] = x.word(i);

            // ★ Tier 3: allocate scratch for all PD-loop steps at once
            size_t val_n = value.size();
            size_t pd_scratch_sz = pd_step_scratch_size(n, max_limbs, val_n);
            uint64_t* pd_buf = getThreadArena().alloc_limbs(pd_scratch_sz);

            // ★ Tier 5: includes step ns (the final Newton step also runs inside the PD loop)
            // → no need to compute a separate Newton correction Q; after the loop, only floor verification
            for (int step = 1; step <= ns; step++) {
                // ★ Local bump allocator: reset to the start each step
                size_t buf_off = 0;
                auto balloc = [&](size_t s) -> uint64_t* {
                    if (s == 0) return pd_buf;  // size 0 is a no-op
                    uint64_t* p = pd_buf + buf_off;
                    buf_off += (s + 7) & ~size_t(7);
                    return p;
                };

                int new_prec = sizes[step];
                int extend = new_prec - cur_prec;
                if (extend > 0) {
                    size_t limb_sh = extend / 64;
                    unsigned bit_sh = extend % 64;
                    size_t new_xn = xn + limb_sh + (bit_sh > 0 ? 1 : 0);
                    if (new_xn > max_limbs) new_xn = max_limbs;
                    if (bit_sh > 0) {
                        uint64_t carry = mpn::lshift(xp + limb_sh, xp, xn, bit_sh);
                        if (xn + limb_sh < max_limbs) xp[xn + limb_sh] = carry;
                    } else if (limb_sh > 0) {
                        for (size_t i = xn; i > 0; i--) xp[i - 1 + limb_sh] = xp[i - 1];
                    }
                    for (size_t i = 0; i < limb_sh; i++) xp[i] = 0;
                    xn = mpn::normalized_size(xp, new_xn);
                }

                // ★ Tier 1: trailing zero stripping
                // Detect the zero limbs added by the precision extension, and
                // run pow and division only on the non-zero part (halving the operands)
                size_t x_tz = 0;
                while (x_tz < xn && xp[x_tz] == 0) x_tz++;
                size_t xn_nz = xn - x_tz;
                if (xn_nz == 0) { cur_prec = 0; continue; }

                // a = value >> a_shift (direct mpn shift)
                int new_scale = P - new_prec;
                int a_shift = ni * new_scale;
                size_t a_limb_off = static_cast (a_shift) / 64;
                unsigned a_bit_off = static_cast (a_shift) % 64;
                size_t a_n = (a_limb_off < val_n) ? (val_n - a_limb_off) : 0;
                // Build a on the local buffer
                uint64_t* ap_buf = nullptr;
                if (a_n > 0 && a_shift > 0) {
                    ap_buf = balloc(a_n + 1);
                    for (size_t i = 0; i < a_n; i++) ap_buf[i] = value.word(i + a_limb_off);
                    if (a_bit_off > 0) mpn::rshift(ap_buf, ap_buf, a_n, a_bit_off);
                    a_n = mpn::normalized_size(ap_buf, a_n);
                } else if (a_n > 0) {
                    ap_buf = balloc(a_n);
                    for (size_t i = 0; i < a_n; i++) ap_buf[i] = value.word(i);
                }
                if (a_n == 0) { cur_prec = 0; continue; }

                // ★ Tier 1: remove the lower part of a within the range that does not affect division accuracy
                // The trailing zeros of x^(n-1) = (n-1)*x_tz limbs are also removed from a
                size_t div_tz = (n == 3) ? 2 * x_tz : static_cast (n - 1) * x_tz;
                if (div_tz >= a_n) div_tz = (a_n > 1) ? a_n - 1 : 0;

              if (n == 3) {
                // ★ n=3 dedicated: x_new = (2*x + a/x²) / 3
                // x² is computed only on the non-zero part (trailing zeros implicit)
                // ★ 2026-04-28: divide_q_approx only references the top qn+3 limbs of the divisor.
                //    Switch to square_high (Karatsuba-high) with final_keep = qn_est + 5;
                //    for an extreme non-balanced case where drop ≫ keep, a0² is fully skipped → much faster.
                size_t x2n_full = 2 * xn_nz;
                uint64_t* x2p = balloc(x2n_full);
                size_t a_div_n_pre = a_n - div_tz;
                size_t qn_est = (a_div_n_pre > x2n_full) ? (a_div_n_pre - x2n_full + 1) : 0;
                // SQHIGH_SAFETY breakdown: divide_q_approx EXTRA(3) + x² leading-zero margin(δ≤1) + safety(4)
                constexpr size_t SQHIGH_SAFETY = 8;
                size_t keep = qn_est + SQHIGH_SAFETY;
                // square_high's Karatsuba-high recursion is adopted only when drop ≥ n (deep truncation)
                // The usual nthRoot case has qn ≈ xn, so keep ≈ xn → shallow truncation → automatic fallback
                // To avoid extra dispatch cost, call mpn::square directly for non-deep-truncation cases
                bool use_sqhigh = (x2n_full >= mpn::BZ_THRESHOLD)
                                  && (qn_est > 0)
                                  && (keep < xn_nz)  // deep truncation: drop > xn_nz
                                  && (xn_nz >= mpn::BZ_THRESHOLD / 2);

                if (use_sqhigh) {
                    // Zero-fill the lower part (insurance against divide_q_approx fallback reading garbage)
                    std::memset(x2p, 0, (x2n_full - keep) * sizeof(uint64_t));
                    uint64_t* sh_sc = balloc(mpn::square_high_scratch_size(xn_nz, keep));
                    mpn::square_high(x2p + (x2n_full - keep), xp + x_tz, xn_nz, keep, sh_sc);
                } else {
                    uint64_t* sq_sc = balloc(mpn::square_scratch_size(xn_nz));
                    mpn::square(x2p, xp + x_tz, xn_nz, sq_sc);
                }
                size_t x2n = mpn::normalized_size(x2p, x2n_full);

                // quot = a_eff / x²_nz (after trailing-zero removal)
                const uint64_t* a_div = ap_buf + div_tz;
                size_t a_div_n = a_n - div_tz;
                uint64_t* qp_div = nullptr; size_t qn_div = 0;
                if (a_div_n >= x2n && x2n >= 2) {
                    size_t q_sz = a_div_n - x2n + 1;
                    qp_div = balloc(q_sz + 2);
                    size_t dsz = mpn::divide_q_approx_scratch_size(a_div_n, x2n);
                    uint64_t* dsc = balloc(dsz);
                    qn_div = mpn::divide_q_approx(qp_div, a_div, a_div_n, x2p, x2n, dsc);
                    qn_div = mpn::normalized_size(qp_div, qn_div);
                } else if (x2n == 1 && a_div_n > 0) {
                    qp_div = balloc(a_div_n);
                    mpn::divmod_1(qp_div, a_div, a_div_n, x2p[0]);
                    qn_div = mpn::normalized_size(qp_div, a_div_n);
                }

                // tp = 2*x + quot (x stays at full size xn)
                size_t tn = xn + 1;
                uint64_t* tp = balloc(std::max(tn, qn_div) + 2);
                tp[xn] = mpn::mul_1(tp, xp, xn, 2ULL); // 2*x
                tn = mpn::normalized_size(tp, tn);
                if (qn_div > 0 && qp_div) {
                    if (qn_div > tn) { for (size_t i = tn; i < qn_div; i++) tp[i] = 0; tn = qn_div; }
                    uint64_t carry = mpn::add(tp, tp, tn, qp_div, qn_div);
                    if (carry) { tp[tn] = carry; tn++; }
                }

                // xp = tp / 3
                mpn::divmod_1(xp, tp, tn, 3ULL);
                xn = mpn::normalized_size(xp, tn);
              } else {
                // General path (n=5,7,11,...) — run pow only on the non-zero part
                // ★ 2026-04-28: for n=5/7, switch the final square to square_high (final_keep).
                //    bn_est = (n-1)*xn_nz, qn_est = a_div_n - bn_est + 1.
                //    With final_keep > 0, mpn_pow_for_root_impl truncates the final square.
                size_t bn_est = static_cast (n - 1) * xn_nz;
                size_t a_div_n_pre = a_n - div_tz;
                size_t qn_est = (a_div_n_pre > bn_est) ? (a_div_n_pre - bn_est + 1) : 0;
                // Safety margin of the final square_high:
                //   EXTRA(3) + δ_max + safety(2)
                //   δ_max is the upper bound of leading-zero accumulation in the pow chain (n=5: 3, n=7: 5)
                size_t final_keep = 0;
                size_t fk_safety = (n == 5) ? 8 : (n == 7 ? 12 : 0);
                if (fk_safety > 0 && qn_est > 0
                    && bn_est >= mpn::BZ_THRESHOLD
                    && (qn_est + fk_safety) < bn_est)
                {
                    final_keep = qn_est + fk_safety;
                }
                MpnPowResult pw = mpn_pow_for_root_impl(xp + x_tz, xn_nz, n, false, balloc, final_keep);

                // quot = a_eff / xpm1_nz (after trailing-zero removal)
                const uint64_t* a_div = ap_buf + div_tz;
                size_t a_div_n = a_n - div_tz;
                size_t p_sz = pw.nm1_n;
                uint64_t* qp_div = nullptr; size_t qn_div = 0;
                if (a_div_n >= p_sz && p_sz >= 2) {
                    size_t q_sz = a_div_n - p_sz + 1;
                    qp_div = balloc(q_sz + 2);
                    size_t dsz = mpn::divide_q_approx_scratch_size(a_div_n, p_sz);
                    uint64_t* dsc = balloc(dsz);
                    qn_div = mpn::divide_q_approx(qp_div, a_div, a_div_n, pw.nm1, p_sz, dsc);
                    qn_div = mpn::normalized_size(qp_div, qn_div);
                } else if (p_sz == 1 && a_div_n > 0) {
                    qp_div = balloc(a_div_n);
                    mpn::divmod_1(qp_div, a_div, a_div_n, pw.nm1[0]);
                    qn_div = mpn::normalized_size(qp_div, a_div_n);
                }

                uint64_t nm1_word = static_cast (n - 1);
                size_t tn = xn + 1;
                uint64_t* tp = balloc(std::max(tn, qn_div) + 2);
                tp[xn] = mpn::mul_1(tp, xp, xn, nm1_word);
                tn = mpn::normalized_size(tp, tn);

                // tp += quot (direct limb)
                if (qn_div > 0 && qp_div) {
                    if (qn_div > tn) {
                        for (size_t i = tn; i < qn_div; i++) tp[i] = 0;
                        tn = qn_div;
                    }
                    uint64_t carry = mpn::add(tp, tp, tn, qp_div, qn_div);
                    if (carry) { tp[tn] = carry; tn++; }
                }

                // xp = tp / n (1-word division)
                mpn::divmod_1(xp, tp, tn, static_cast (n));
                xn = mpn::normalized_size(xp, tn);
              } // end else (general n)

                cur_prec = static_cast (xn * 64 - (xn > 0 ? std::countl_zero(xp[xn-1]) : 64));
            }

            // Restore xp → x (after step ns completes, full P bits)
            x = Int::fromRawWords(std::span (xp, xn), 1);
        }
        // ★ Tier 5: since the PD loop ran through step ns,
        // no separate Newton correction is needed. Delegate directly to the floor verification below.
    } else {
        // Small root: standard Newton (no PD needed)
        constexpr int MAX_ITER = 200;
        for (int iter = 0; iter < MAX_ITER; iter++) {
            auto [xpm1, _xpn_unused] = mpn_pow_for_root(x, n, false);
            Int quot;
            IntOps::divUnchecked(value, xpm1, quot);
            // x_new = (n_minus_1 * x + quot) / n_int (Unchecked: no special states)
            Int tmp;
            IntOps::mulUnchecked(n_minus_1, x, tmp);
            IntOps::addUnchecked(tmp, quot, tmp);
            Int x_new;
            IntOps::divUnchecked(tmp, n_int, x_new);
            if (x_new == x) break;
            Int diff = (x_new > x) ? (x_new - x) : (x - x_new);
            if (diff <= Int::One()) {
                if (x_new < x) x = x_new;
                break;
            }
            x = x_new;
        }
    }

    // === Final verification for the small root ===
    {
        auto [xpm1, _xpn_unused] = mpn_pow_for_root(x, n, false);
        Int R = value - xpm1 * x;

        while (R.isNegative()) {
            x = x - Int::One();
            if (x.isZero()) {
                if (pRemainder) *pRemainder = value;  // value - 0^n = value
                return x;
            }
            xpm1 = mpn_pow_for_root(x, n, false).first;
            R = value - xpm1 * x;
        }

        // ★ Codex improvement #6: replace pow(x+1,n) with a difference test
        // (x+1)^n - x^n = Σ_{k=0}^{n-1} C(n,k) * x^k
        // linear term = n*x^(n-1); if R >= linear term, an undershoot is possible
        // Safe test: if R >= n*xpm1, check with mpn whether (x+1)^n ≤ value
        if (R >= n_int * xpm1) {
            auto [_, xp1_n] = mpn_pow_for_root(x + Int::One(), n, true);
            if (xp1_n <= value) {
                x = x + Int::One();
                R = value - xp1_n;  // bumped: new x = old x+1, R uses (x+1)^n
            }
        }

        // Candidate A: pass R to pRemainder (avoids the redundant pow(root,n) in nthRootRem)
        if (pRemainder) *pRemainder = R;
    }

    return x;
}


// ========================================================================
// nthRootRem: compute the n-th root and the remainder together
// ========================================================================

Int IntSqrt::nthRootRem(const Int& value, uint32_t n, Int& remainder) {
    // Candidate A (PLAN_NTHROOT_FOLLOWUP): obtain root and remainder together
    // via nthRoot_internal. Since the verification block has already computed
    // R = value - x^n, eliminate the old implementation's redundant pow(root, n) call (~33% speedup expected).
    return nthRoot_internal(value, n, &remainder);
}

// =============================================================================
// isPerfectPower — determine whether a k with value = b^k (k >= 2) exists
// =============================================================================

bool IntSqrt::isPerfectPower(const Int& value, Int* pBase, uint32_t* pExp) {
    if (value.isNaN() || value.isInfinite()) return false;
    if (value.isNegative()) return false;  // negative perfect powers are not supported
    if (value.isZero()) {
        if (pBase) *pBase = Int(0);
        if (pExp) *pExp = 2;
        return true;  // 0 = 0^2
    }
    if (value == Int::One()) {
        if (pBase) *pBase = Int(1);
        if (pExp) *pExp = 2;
        return true;  // 1 = 1^2
    }

    size_t bits = value.bitLength();

    // Upper bound of k: since 2^k <= value, k <= bitLength
    uint32_t max_k = static_cast (bits);

    // Trying only prime exponents is sufficient (e.g., a^6 = (a^2)^3 = (a^3)^2)
    // Try in order starting from the smallest prime
    static const uint32_t primes[] = {
        2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61
    };

    for (uint32_t p : primes) {
        if (p > max_k) break;

        Int root = nthRoot(value, p);
        if (pow(root, p) == value) {
            if (pBase) *pBase = root;
            if (pExp) *pExp = p;
            return true;
        }
    }

    return false;
}

} // namespace sangi