simplify rand(::UnitRange) a bit (#24989)

Author: rfourquet

base/random/RNGs.jl

@@ -4,7 +4,7 @@
 
 # SamplerUnion(Union{X,Y,...}) == Union{SamplerType{X},SamplerType{Y},...}
 SamplerUnion(U::Union) = Union{map(T->SamplerType{T}, Base.uniontypes(U))...}
-const SamplerBoolBitInteger = SamplerUnion(Union{Bool, Base.BitInteger})
+const SamplerBoolBitInteger = SamplerUnion(Union{Bool, BitInteger})
 
 if Sys.iswindows()
     struct RandomDevice <: AbstractRNG
@@ -30,7 +30,7 @@ else # !windows
 end # os-test
 
 # NOTE: this can't be put within the if-else block above
-for T in (Bool, Base.BitInteger_types...)
+for T in (Bool, BitInteger_types...)
     if Sys.iswindows()
         @eval function rand!(rd::RandomDevice, A::Array{$T}, ::SamplerType{$T})
             ccall((:SystemFunction036, :Advapi32), stdcall, UInt8, (Ptr{Void}, UInt32),
@@ -403,7 +403,7 @@ function rand!(r::MersenneTwister, A::Array{UInt128}, ::SamplerType{UInt128})
     A
 end
 
-for T in Base.BitInteger_types
+for T in BitInteger_types
     T === UInt128 && continue
     @eval function rand!(r::MersenneTwister, A::Array{$T}, ::SamplerType{$T})
         n = length(A)
@@ -451,7 +451,7 @@ function rand(rng::MersenneTwister,
     x % T + first(r)
 end
 
-for T in (Bool, Base.BitInteger_types...) # eval because of ambiguity otherwise
+for T in (Bool, BitInteger_types...) # eval because of ambiguity otherwise
     @eval Sampler(rng::MersenneTwister, r::UnitRange{$T}, ::Val{1}) =
         SamplerTrivial(r)
 end

base/random/generation.jl

@@ -133,68 +133,57 @@ rem_knuth(a::T, b::T) where {T<:Unsigned} = b != 0 ? a % b : a
 # that is 0xFFFF...FFFF if k = typemax(T) - typemin(T) with intentional underflow
 # see http://stackoverflow.com/questions/29182036/integer-arithmetic-add-1-to-uint-max-and-divide-by-n-without-overflow
 maxmultiple(k::T) where {T<:Unsigned} =
-    (div(typemax(T) - k + oneunit(k), k + (k == 0))*k + k - oneunit(k))::T
+    (div(typemax(T) - k + one(k), k + (k == 0))*k + k - one(k))::T
+
+# serves as rejection threshold
+_maxmultiple(k)  = maxmultiple(k)
 
 # maximum multiple of k within 1:2^32 or 1:2^64 decremented by one, depending on size
-maxmultiplemix(k::UInt64) = k >> 32 != 0 ?
+_maxmultiple(k::UInt64)::UInt64 = k >> 32 != 0 ?
     maxmultiple(k) :
-    (div(0x0000000100000000, k + (k == 0))*k - oneunit(k))::UInt64
+    div(0x0000000100000000, k + (k == 0))*k - one(k)
 
-struct SamplerRangeInt{T<:Integer,U<:Unsigned} <: Sampler
+struct SamplerRangeInt{T<:Union{Bool,Integer},U<:Unsigned} <: Sampler
     a::T   # first element of the range
     k::U   # range length or zero for full range
     u::U   # rejection threshold
 end
 
-# generators with 32, 128 bits entropy
-SamplerRangeInt(a::T, k::U) where {T,U<:Union{UInt32,UInt128}} =
-    SamplerRangeInt{T,U}(a, k, maxmultiple(k))
+function SamplerRangeInt(a::T, diff::U) where {T<:Union{Bool,Integer},U<:Unsigned}
+    k = diff+one(U)
+    SamplerRangeInt{T,U}(a, k, _maxmultiple(k)) # overflow ok
+end
 
-# mixed 32/64 bits entropy generator
-SamplerRangeInt(a::T, k::UInt64) where {T} =
-    SamplerRangeInt{T,UInt64}(a, k, maxmultiplemix(k))
+uint_sup(::Type{<:Union{Bool,BitInteger}}) = UInt32
+uint_sup(::Type{<:Union{Int64,UInt64}}) = UInt64
+uint_sup(::Type{<:Union{Int128,UInt128}}) = UInt128
 
-function Sampler(::AbstractRNG, r::AbstractUnitRange{T}, ::Repetition) where T<:Unsigned
+function SamplerRangeInt(r::AbstractUnitRange{T}) where T<:Union{Bool,BitInteger}
     isempty(r) && throw(ArgumentError("range must be non-empty"))
-    SamplerRangeInt(first(r), last(r) - first(r) + oneunit(T))
+    SamplerRangeInt(first(r), (last(r) - first(r)) % uint_sup(T))
 end
 
-for (T, U) in [(UInt8, UInt32), (UInt16, UInt32),
-               (Int8, UInt32), (Int16, UInt32), (Int32, UInt32),
-               (Int64, UInt64), (Int128, UInt128), (Bool, UInt32)]
+Sampler(::AbstractRNG, r::AbstractUnitRange{T}, ::Repetition) where {T<:Union{Bool,BitInteger}} =
+    SamplerRangeInt(r)
 
-    @eval Sampler(::AbstractRNG, r::AbstractUnitRange{$T}, ::Repetition) = begin
-        isempty(r) && throw(ArgumentError("range must be non-empty"))
-        # overflow ok:
-        SamplerRangeInt(first(r), convert($U, unsigned(last(r) - first(r)) + one($U)))
+function rand_lteq(rng::AbstractRNG, u::T)::T where T
+    while true
+        x = rand(rng, T)
+        x <= u && return x
     end
 end
 
 # this function uses 32 bit entropy for small ranges of length <= typemax(UInt32) + 1
-# SamplerRangeInt is responsible for providing the right value of k
-function rand(rng::AbstractRNG, sp::SamplerRangeInt{T,UInt64}) where T<:Union{UInt64,Int64}
-    local x::UInt64
-    if (sp.k - 1) >> 32 == 0
-        x = rand(rng, UInt32)
-        while x > sp.u
-            x = rand(rng, UInt32)
-        end
-    else
-        x = rand(rng, UInt64)
-        while x > sp.u
-            x = rand(rng, UInt64)
-        end
-    end
-    return reinterpret(T, reinterpret(UInt64, sp.a) + rem_knuth(x, sp.k))
+function rand(rng::AbstractRNG, sp::SamplerRangeInt{T,UInt64}) where T<:BitInteger
+    x::UInt64 = (sp.k - 1) >> 32 == 0 ?
+        rand_lteq(rng, sp.u % UInt32) % UInt64 :
+        rand_lteq(rng, sp.u)
+    return ((sp.a % UInt64) + rem_knuth(x, sp.k)) % T
 end
 
-function rand(rng::AbstractRNG, sp::SamplerRangeInt{T,U}) where {T<:Integer,U<:Unsigned}
-    x = rand(rng, U)
-    while x > sp.u
-        x = rand(rng, U)
-    end
-    (unsigned(sp.a) + rem_knuth(x, sp.k)) % T
-end
+rand(rng::AbstractRNG, sp::SamplerRangeInt{T,U}) where {T<:Union{Bool,BitInteger},U} =
+    (unsigned(sp.a) + rem_knuth(rand_lteq(rng, sp.u), sp.k)) % T
+
 
 ### BigInt
 

base/random/random.jl

@@ -4,6 +4,7 @@ module Random
 
 using Base.dSFMT
 using Base.GMP: Limb, MPZ
+using Base: BitInteger, BitInteger_types
 import Base: copymutable, copy, copy!, ==, hash
 
 export srand,

test/random.jl

@@ -103,9 +103,9 @@ if sizeof(Int32) < sizeof(Int)
     @test all([div(0x00010000000000000000,k)*k - 1 == Base.Random.RangeGenerator(map(UInt64,1:k)).u for k in 13 .+ Int64(2).^(32:62)])
     @test all([div(0x00010000000000000000,k)*k - 1 == Base.Random.RangeGenerator(map(Int64,1:k)).u for k in 13 .+ Int64(2).^(32:61)])
 
-    @test Base.Random.maxmultiplemix(0x000100000000) === 0xffffffffffffffff
-    @test Base.Random.maxmultiplemix(0x0000FFFFFFFF) === 0x00000000fffffffe
-    @test Base.Random.maxmultiplemix(0x000000000000) === 0xffffffffffffffff
+    @test Base.Random._maxmultiple(0x000100000000) === 0xffffffffffffffff
+    @test Base.Random._maxmultiple(0x0000FFFFFFFF) === 0x00000000fffffffe
+    @test Base.Random._maxmultiple(0x000000000000) === 0xffffffffffffffff
 end
 
 # BigInt specific