From c8990483a13cf4ce5b7a77e11a4c737cde53d61b Mon Sep 17 00:00:00 2001 From: Oscar Smith Date: Sat, 11 Jul 2026 20:14:37 -0400 Subject: [PATCH] impliment iroot and friends for native julia. Much faster for native types --- Project.toml | 3 +- README.md | 11 +-- src/IntegerMathUtils.jl | 178 +++++++++++++++++++++++++++++++++----- test/runtests.jl | 185 ++++++++++++++++++++++++++++++++++++++++ 4 files changed, 350 insertions(+), 27 deletions(-) diff --git a/Project.toml b/Project.toml index 44afb05..6f4e352 100644 --- a/Project.toml +++ b/Project.toml @@ -3,10 +3,11 @@ uuid = "18e54dd8-cb9d-406c-a71d-865a43cbb235" version = "0.1.3" [extras] +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Test"] +test = ["Test", "Random"] [compat] julia = "1.6" diff --git a/README.md b/README.md index 8a922cd..02ecb44 100644 --- a/README.md +++ b/README.md @@ -1,11 +1,12 @@ # IntegerMathUtils.jl -This library adds several functions useful for doing math on integers. Most of these are GMP wrappers that may have faster implimentations for smaller integer types. +This library adds several functions useful for doing math on integers. Machine integers (and any other non-`BigInt` `Integer` type) use native Julia algorithms; `BigInt` inputs use GMP. **Functions** -* `iroot(x::Integer, n::integer)` the integer nth root of `x`. Specifically, this is the largest integer `a` such that `a^n <= x`. Note that `n` must fit into an `Int64` (for GMP compatability). -* `ispower(x::Integer)` return if there are integer `base` and `exponent>1` values such that `base^exponent = x`. -* `find_exponent(x::Integer)` returns the largest possible integer `exponent` such that `base^exponent = x` for some `base`. Returns `1` for `x ∈ [0,1]`. +* `iroot(x::Integer, n::Integer)` the integer nth root of `x`, truncated toward zero: for `x >= 0` the largest integer `a` such that `a^n <= x`, and `iroot(-x, n) == -iroot(x, n)` for odd `n` (e.g. `iroot(-9, 3) == -2`). Throws a `DomainError` for `n <= 0` or an even root of a negative number. +* `rootrem(x::Integer, n::Integer)` returns `(iroot(x, n), x - iroot(x, n)^n)`. +* `ispower(x::Integer)` returns whether there are integer `base` and `exponent > 1` values such that `base^exponent == x`. For negative `x` only odd exponents count, so `ispower(-8) == true` but `ispower(-4) == false`. `0` and `±1` are perfect powers. +* `find_exponent(x::Integer)` returns the largest possible integer `exponent` such that `base^exponent == x` for some `base` (only odd exponents for negative `x`). Returns `1` for `x ∈ [-1, 1]`. * `is_probably_prime(x::Integer; reps=25)` returns if `x` is prime. Will be incorrect less than `4^-reps` of the time. -* `kronecker(a::Integer, n::Integer)` Computes the [Kronecker_symbol](https://en.wikipedia.org/wiki/Kronecker_symbol) which is a generalization of the legendre and jacobi symbols. +* `kronecker(a::Integer, n::Integer)` Computes the [Kronecker symbol](https://en.wikipedia.org/wiki/Kronecker_symbol) which is a generalization of the legendre and jacobi symbols. diff --git a/src/IntegerMathUtils.jl b/src/IntegerMathUtils.jl index 136b63f..f0d379c 100644 --- a/src/IntegerMathUtils.jl +++ b/src/IntegerMathUtils.jl @@ -1,45 +1,181 @@ module IntegerMathUtils export iroot, ispower, rootrem, find_exponent, is_probably_prime, kronecker -iroot(x::Integer, n::Integer) = iroot(x, Cint(n)) +using Base: uabs, BitUnsigned64 -# TODO: Add more efficient implimentation -iroot(x::T, n::Cint) where {T<:Integer} = T(iroot(big(x), Cint(n))) +@static if isdefined(Base, :top_set_bit) + using Base: top_set_bit +else + top_set_bit(x::Base.BitInteger) = 8sizeof(x) - leading_zeros(x) + top_set_bit(x::Integer) = ndigits(x, base=2) +end -function iroot(x::BigInt, n::Cint) +function iroot(x::BigInt, n::Integer) n <= 0 && throw(DomainError(n, "Exponent must be > 0")) - x <= 0 && iseven(n) && throw(DomainError(n, "This is a math no-no")) + x < 0 && iseven(n) && throw(DomainError(x, "Even roots require x >= 0")) ans = BigInt() - @ccall :libgmp.__gmpz_root(ans::Ref{BigInt}, x::Ref{BigInt}, n::Cint)::Cint + @ccall :libgmp.__gmpz_root(ans::Ref{BigInt}, x::Ref{BigInt}, n::Culong)::Cint ans end -# TODO: Add more efficient implimentation for smaller types -function rootrem(x::T, n::Integer) where {T<:Integer} - x = big(x) - n = Cint(n) +function rootrem(x::BigInt, n::Integer) n <= 0 && throw(DomainError(n, "Exponent must be > 0")) - x <= 0 && iseven(x) && throw(DomainError(n, "This is a math no-no")) + x < 0 && iseven(n) && throw(DomainError(x, "Even roots require x >= 0")) root = BigInt() rem = BigInt() - @ccall :libgmp.__gmpz_rootrem(root::Ref{BigInt}, rem::Ref{BigInt}, x::Ref{BigInt}, n::Int)::Nothing - return (root, T(rem)) + @ccall :libgmp.__gmpz_rootrem(root::Ref{BigInt}, rem::Ref{BigInt}, x::Ref{BigInt}, n::Culong)::Cvoid + return (root, rem) end -# TODO: Add more efficient implimentation for smaller types -ispower(x::Integer) = ispower(big(x)) - function ispower(x::BigInt) return 0 != @ccall :libgmp.__gmpz_perfect_power_p(x::Ref{BigInt})::Cint end -# TODO: Add more efficient implimentation for smaller types +# a^n in T via power-by-squaring, detecting overflow. Returns (value, overflowed) +function pow_with_overflow(a::T, n::Integer) where T<:Base.BitInteger + p = one(T) + while n > 0 + if isodd(n) + p, o = Base.Checked.mul_with_overflow(p, a) + o && return (p, true) + end + n >>= 1 + n > 0 || break + a, o = Base.Checked.mul_with_overflow(a, a) + o && return (p, true) + end + return (p, false) +end +function _pow_leq(a::T, n::Integer, x::T) where T<:Base.BitInteger + p, o = pow_with_overflow(a, n) + return !o && p <= x +end +_pow_leq(a::T, n::Integer, x::T) where T<:Integer = a^n <= x + +# Float64-accurate estimate of x^(1/n) for x <= typemax(UInt128). +# Because the root must be small, this is within +-1 of the true root when n>2. +function _iroot_approx(x::T, n::Integer) where T<:Base.BitInteger + trunc(T, exp(log(Float64(x)) / n)) +end + +# A Float64-accurate estimate of x^(1/n) for x > typemax(UInt128) +# Split x = hi*2^sh with hi its top 53 bits, estimate log2(root) = (log2(hi)+sh)/n +# and place a 53-bit mantissa at the right exponent. +function _iroot_float(x::T, n::Integer) where T<:Integer + b = top_set_bit(x) + sh = b - 53 + L = (log2(Float64(x >> sh)) + sh) / n + e = floor(Int, L) + m53 = trunc(T, exp2(L - e) * exp2(53)) + return e >= 53 ? (m53 << (e - 53)) : (m53 >> (53 - e)) +end + +# One integer Newton step +@inline function _iroot_newton(x::T, n::Integer, a::T) where T<:Integer + return ((n - 1) * a + x ÷ a^(n - 1)) ÷ n +end + +function _iroot_approx(x::T, n::Integer) where T<:Integer + x <= typemax(UInt128) && return T(_iroot_approx(UInt128(x), n)) + # one unconditional step lifts any seed to >= floor(root)-1 (by weighted AM-GM) + a = _iroot_newton(x, n, _iroot_float(x, n)) + while true # after which the iteration descends until return + b = _iroot_newton(x, n, a) + b >= a && return a + a = b + end +end + +# floor(x^(1/n)) for x >= 0; n >= 2 (n == 2 only for non-bit Integer types, +# machine ints take the Base.isqrt path in iroot) +function _iroot_nonneg(x::T, n::Integer) where T<:Integer + x <= one(T) && return x + n >= top_set_bit(x) && return one(T) # 2^n > x + a = _iroot_approx(x, n) + if _pow_leq(a, n, x) + while _pow_leq(a + one(T), n, x) + a += one(T) + end + else + a -= one(T) + while !_pow_leq(a, n, x) + a -= one(T) + end + end + return a +end + +# Base.isqrt is already optimal for machine ints; its generic Integer +# fallback is Float64-based though, which is wrong for wide values. +_isqrt(x::Base.BitInteger) = Base.isqrt(x) +_isqrt(x::Integer) = _iroot_nonneg(x, 2) + +function iroot(x::T, n::Integer) where T<:Integer + n <= 0 && throw(DomainError(n, "Exponent must be > 0")) + x < 0 && iseven(n) && throw(DomainError(x, "Even roots require x >= 0")) + n == 1 && return x + n == 2 && return T(_isqrt(uabs(x))) # x >= 0 here + r = T(_iroot_nonneg(uabs(x), n)) + return x < 0 ? -r : r +end + +function rootrem(x::Integer, n::Integer) + a = iroot(x, n) + return (a, x - a^n) +end + +# Is m (>= 2) an exact k-th power (k odd, >= 3)? +# Callers handle k == 2 via _isqrt directly +# (its root can be large enough that the float error below exceeds 0.5). +# The float root r = exp(lm/k) has provable worst-case error +# E <= 1.5*r*(|ln r|+1)*2^-52; for k >= 3 and m < 2^128 this is << 0.5 +# so round(r) is the only integer candidate +function _kth_root_exact(m::T, k::Integer, logm::Float64) where T<:Base.BitInteger + lr = logm / k + r = exp(lr) + tol = 2 * r * (abs(lr) + 1) * exp2(-52) + b = round(r) + abs(r - b) > tol && return false + c = unsafe_trunc(T, b) + p, o = pow_with_overflow(c, k) + return !o && p == m +end +# Arbitrary-precision types: no reliable Float64, so verify via the native root. +function _kth_root_exact(m::T, k::Integer, logm::Float64) where T<:Integer + a = iroot(m, k) + return a^k == m # a^k <= m, so no overflow +end + +function ispower(x::Integer) + ux = uabs(x) + ux <= one(ux) && return true + neg = x < 0 + trailing_zeros(ux) == 1 && return false # 2 | x exactly once: not a power + if !neg + a = _isqrt(ux) + a * a == ux && return true + end + lx = log(Float64(ux)) + for k in 3:2:top_set_bit(ux) + _kth_root_exact(ux, k, lx)[1] && return true + end + return false +end + function find_exponent(x::Integer) - x <= 1 && return 1 - for exponent in ndigits(x, base=2):-1:2 - rootrem(x, exponent)[2] == 0 && return exponent + ux = uabs(x) + ux <= one(ux) && return 1 + neg = x < 0 + lx = log(Float64(ux)) + for e in top_set_bit(ux):-1:3 + neg && iseven(e) && continue + _kth_root_exact(ux, e, lx)[1] && return e + end + if !neg + a = _isqrt(ux) + a * a == ux && return 2 end - 1 + return 1 end function is_probably_prime(x::Integer; reps=25) diff --git a/test/runtests.jl b/test/runtests.jl index a945ba2..7e05602 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,4 +1,7 @@ using Test, IntegerMathUtils +using Random + +@testset "IntegerMathUtils" begin @testset "iroot" begin for T in (Int32, Int64, BigInt) @@ -70,3 +73,185 @@ end @test kronecker(-4, -5) == -1 @test kronecker(4, 6) == kronecker(-4, 0) == 0 end + +@testset "iroot(x, 2) matches Base.isqrt" begin + for T in (Int32, UInt32, Int64, UInt64, Int128, UInt128) + @test iroot(typemax(T), 2) == Base.isqrt(typemax(T)) + end + # perfect squares and neighbors, including near the Float64 exactness edge (2^52) + for a in (94906263, 94906264, 94906265, 94906266, 3037000499, 2^26, 2^26 - 1) + @test iroot(Int64(a)^2, 2) == a + @test iroot(Int64(a)^2 + 1, 2) == a + @test iroot(Int64(a)^2 - 1, 2) == a - 1 + end + rng = MersenneTwister(42) + for T in (UInt32, Int64, UInt64, Int128, UInt128), _ in 1:200 + x = rand(rng, T) & typemax(T) # nonnegative + @test iroot(x, 2) === T(Base.isqrt(x)) + end +end + +@testset "iroot extended semantics" begin + for T in (Int32, Int64, Int128, BigInt) + @test iroot(T(0), 2) == 0 + @test iroot(T(0), 3) == 0 + @test iroot(T(1), 100) == 1 + @test iroot(T(-9), 3) == -2 # truncated toward zero, like GMP + @test iroot(T(-1), 5) == -1 + @test_throws DomainError iroot(T(5), 0) + @test_throws DomainError iroot(T(5), -2) + @test_throws DomainError iroot(T(-4), 2) + end + @test iroot(typemin(Int64), 63) == -2 + @test iroot(typemin(Int8), 7) == -2 + @test iroot(typemax(UInt128), 2) == UInt128(2)^64 - 1 + @test iroot(Int128(2)^126, 63) == 4 + @test iroot(typemax(UInt128), 127) == 2 + @test iroot(Int8(127), 100) == 1 +end + +@testset "rootrem" begin + for T in (Int32, Int64, Int128, BigInt) + r = rootrem(T(1001), 3) + @test r == (10, 1) + @test r isa Tuple{T,T} + @test rootrem(T(1000), 3) == (10, 0) + @test rootrem(T(-1001), 3) == (-10, -1) + @test rootrem(T(-1000), 3) == (-10, 0) + @test_throws DomainError rootrem(T(-4), 2) + @test_throws DomainError rootrem(T(5), 0) + end + @test rootrem(typemin(Int64), 63) == (-2, 0) +end + +@testset "negative perfect powers" begin + for T in (Int32, Int64, Int128, BigInt) + @test ispower(T(-8)) == true + @test ispower(T(-4)) == false + @test ispower(T(-1)) == true + @test ispower(T(-16)) == false # even exponents don't count for negatives + @test ispower(T(-64)) == true # (-4)^3 + @test find_exponent(T(-8)) === 3 + @test find_exponent(T(-4)) === 1 + @test find_exponent(T(-1)) === 1 + @test find_exponent(T(-64)) === 3 + end + @test ispower(typemin(Int64)) == true # (-2)^63 + @test find_exponent(typemin(Int64)) === 63 + @test find_exponent(typemin(Int8)) === 7 + @test find_exponent(-big(3)^15) === 15 +end + +@testset "perfect powers near typemax" begin + @test ispower(Int64(3)^39) == true + @test find_exponent(Int64(3)^39) === 39 + @test ispower(UInt64(3)^40) == true + @test find_exponent(UInt64(3)^40) === 40 + @test ispower(Int64(2)^62) == true + @test find_exponent(Int64(2)^62) === 62 + for x in (Int64(3)^39 + 1, Int64(3)^39 - 1, Int64(2)^62 + 1) + @test ispower(x) == ispower(big(x)) + @test find_exponent(x) == find_exponent(big(x)) + end +end + +@testset "cross-check native against BigInt/GMP" begin + rng = MersenneTwister(1234) + types = (Int8, UInt8, Int16, UInt16, Int32, UInt32, Int64, UInt64, Int128, UInt128) + for T in types, _ in 1:200 + x = rand(rng, T) + for n in (1, 2, 3, 5, 7, rand(rng, 2:130)) + if x < 0 && iseven(n) + @test_throws DomainError iroot(x, n) + @test_throws DomainError rootrem(x, n) + else + a = iroot(x, n) + @test a isa T + mx, ma = abs(big(x)), abs(big(a)) + @test ma^n <= mx < (ma + 1)^n + @test iroot(big(x), n) == big(a) + root, rem = rootrem(x, n) + @test root === a + @test rem isa T + @test big(rem) == big(x) - big(a)^n + end + end + @test ispower(x) == ispower(big(x)) + e = find_exponent(x) + if abs(big(x)) > 1 + @test (e > 1) == ispower(big(x)) + @test big(iroot(x, e))^e == big(x) + end + end +end + +@testset "overflow-check band (root == 2 near typemax)" begin + # These are the cases that prove _pow_leq's overflow check is needed: + # 2^n <= typemax < 3^n, so the true root is 2 while (root+1)^n overflows. + @test iroot(typemax(UInt8), 7) == 2 # 3^7 == 139 (mod 256) <= 255 + @test iroot(Int8(127), 6) == 2 + for n in 41:63 + @test iroot(typemax(UInt64), n) == UInt64(2) + @test iroot(typemax(UInt64), n) == UInt64(iroot(big(typemax(UInt64)), n)) + end + for n in 85:127 + @test iroot(typemax(UInt128), n) == UInt128(2) + end +end + +@testset "near-integer reject: bᵏ and neighbors near typemax" begin + for T in (Int32, UInt32, Int64, UInt64, Int128, UInt128) + hi = big(typemax(T)) + for k in (3, 5, 7, 11, 3, 13) + b = Int(floor(hi^(1/k))) + for bb in (b-1, b, b+1) + bb < 2 && continue + p = big(bb)^k + p > hi && continue + x = T(p) + @test ispower(x) == ispower(p) + @test find_exponent(x) == find_exponent(p) + # exact power must be detected, and near-misses rejected + @test ispower(x) == true + if x < typemax(T) + @test ispower(x + one(T)) == ispower(p + 1) + end + if x > T(2) + @test ispower(x - one(T)) == ispower(p - 1) + end + end + end + end +end + +@testset "generic (non-BitInteger) native root path" begin + # Public iroot(::BigInt) uses GMP, but the internal _iroot_nonneg / _iroot_approx + # run the generic T<:Integer code (Newton + rescaled seed / UInt128 shortcut). + # Exercise them on BigInt and cross-check against GMP. + IMU = IntegerMathUtils + rng = MersenneTwister(2024) + for _ in 1:300 + x = rand(rng, big(0):(big(2)^rand(rng, 1:400))) + for n in (2, 3, 5, 7, rand(rng, 2:80)) + @test IMU._iroot_nonneg(x, n) == iroot(x, n) + end + end + # explicitly hit both branches of the generic _iroot_approx and their edges + @test IMU._iroot_nonneg(big(2)^127, 3) == iroot(big(2)^127, 3) # <= typemax(UInt128) + @test IMU._iroot_nonneg(big(2)^128, 3) == iroot(big(2)^128, 3) # just over + @test IMU._iroot_nonneg(big(10)^100, 5) == iroot(big(10)^100, 5) # large, rescaled seed + @test IMU._iroot_nonneg(big(2)^1000, 3) == iroot(big(2)^1000, 3) # root exceeds Float64 range + @test IMU._iroot_nonneg(big(2)^999, 999) == iroot(big(2)^999, 999) # huge x, tiny root + @test IMU._iroot_nonneg((big(7)^200), 200) == 7 # exact wide power +end + +@testset "machine-int paths don't allocate" begin + @test @allocated(iroot(10^15, 2)) == 0 + @test @allocated(iroot(10^15, 3)) == 0 + @test @allocated(iroot(Int128(2)^100, 5)) == 0 + @test @allocated(rootrem(10^15, 3)) == 0 + @test @allocated(ispower(10^15)) == 0 + @test @allocated(find_exponent(Int64(3)^39)) == 0 +end + +end