Merge pull request #25004 from JuliaLang/rf/rand/UInt52

RNG: delete rand_ui* functions in favor of multi-dispatch

Author: rfourquet

base/random/RNGs.jl

@@ -42,6 +42,9 @@ for T in (Bool, BitInteger_types...)
     end
 end
 
+# RandomDevice produces natively UInt64
+rng_native_52(::RandomDevice) = UInt64
+
 """
     RandomDevice()
 
@@ -54,10 +57,6 @@ RandomDevice
 RandomDevice(::Void) = RandomDevice()
 srand(rng::RandomDevice) = rng
 
-### generation of floats
-
-rand(r::RandomDevice, sp::SamplerTrivial{<:FloatInterval}) = rand_generic(r, sp[])
-
 
 ## MersenneTwister
 
@@ -209,58 +208,59 @@ const GLOBAL_RNG = MersenneTwister(0)
 
 ### generation
 
+# MersenneTwister produces natively Float64
+rng_native_52(::MersenneTwister) = Float64
+
 #### helper functions
 
 # precondition: !mt_empty(r)
 rand_inbounds(r::MersenneTwister, ::Close1Open2_64) = mt_pop!(r)
-rand_inbounds(r::MersenneTwister, ::CloseOpen_64) =
+rand_inbounds(r::MersenneTwister, ::CloseOpen_64=CloseOpen()) =
     rand_inbounds(r, Close1Open2()) - 1.0
-rand_inbounds(r::MersenneTwister) = rand_inbounds(r, CloseOpen())
 
-rand_ui52_raw_inbounds(r::MersenneTwister) =
-    reinterpret(UInt64, rand_inbounds(r, Close1Open2()))
-rand_ui52_raw(r::MersenneTwister) = (reserve_1(r); rand_ui52_raw_inbounds(r))
+rand_inbounds(r::MersenneTwister, ::UInt52Raw{T}) where {T<:BitInteger} =
+    reinterpret(UInt64, rand_inbounds(r, Close1Open2())) % T
 
-function rand_ui2x52_raw(r::MersenneTwister)
-    reserve(r, 2)
-    rand_ui52_raw_inbounds(r) % UInt128 << 64 | rand_ui52_raw_inbounds(r)
+function rand(r::MersenneTwister, x::SamplerTrivial{UInt52Raw{UInt64}})
+    reserve_1(r)
+    rand_inbounds(r, x[])
 end
 
-function rand_ui104_raw(r::MersenneTwister)
+function rand(r::MersenneTwister, ::SamplerTrivial{UInt2x52Raw{UInt128}})
     reserve(r, 2)
-    rand_ui52_raw_inbounds(r) % UInt128 << 52 ⊻ rand_ui52_raw_inbounds(r)
+    rand_inbounds(r, UInt52Raw(UInt128)) << 64 | rand_inbounds(r, UInt52Raw(UInt128))
 end
 
-rand_ui10_raw(r::MersenneTwister) = rand_ui52_raw(r)
-rand_ui23_raw(r::MersenneTwister) = rand_ui52_raw(r)
+function rand(r::MersenneTwister, ::SamplerTrivial{UInt104Raw{UInt128}})
+    reserve(r, 2)
+    rand_inbounds(r, UInt52Raw(UInt128)) << 52 ⊻ rand_inbounds(r, UInt52Raw(UInt128))
+end
 
 #### floats
 
-rand(r::MersenneTwister, sp::SamplerTrivial{<:FloatInterval_64}) =
+rand(r::MersenneTwister, sp::SamplerTrivial{Close1Open2_64}) =
     (reserve_1(r); rand_inbounds(r, sp[]))
 
-rand(r::MersenneTwister, sp::SamplerTrivial{<:FloatInterval}) = rand_generic(r, sp[])
-
 #### integers
 
 rand(r::MersenneTwister,
      T::SamplerUnion(Union{Bool,Int8,UInt8,Int16,UInt16,Int32,UInt32})) =
-         rand_ui52_raw(r) % T[]
+         rand(r, UInt52Raw()) % T[]
 
 function rand(r::MersenneTwister, ::SamplerType{UInt64})
     reserve(r, 2)
-    rand_ui52_raw_inbounds(r) << 32 ⊻ rand_ui52_raw_inbounds(r)
+    rand_inbounds(r, UInt52Raw()) << 32 ⊻ rand_inbounds(r, UInt52Raw())
 end
 
 function rand(r::MersenneTwister, ::SamplerType{UInt128})
     reserve(r, 3)
-    xor(rand_ui52_raw_inbounds(r) % UInt128 << 96,
-        rand_ui52_raw_inbounds(r) % UInt128 << 48,
-        rand_ui52_raw_inbounds(r))
+    xor(rand_inbounds(r, UInt52Raw(UInt128))  << 96,
+        rand_inbounds(r, UInt52Raw(UInt128))  << 48,
+        rand_inbounds(r, UInt52Raw(UInt128)))
 end
 
-rand(r::MersenneTwister, ::SamplerType{Int64})  = reinterpret(Int64,  rand(r, UInt64))
-rand(r::MersenneTwister, ::SamplerType{Int128}) = reinterpret(Int128, rand(r, UInt128))
+rand(r::MersenneTwister, ::SamplerType{Int64})  = rand(r, UInt64) % Int64
+rand(r::MersenneTwister, ::SamplerType{Int128}) = rand(r, UInt128) % Int128
 
 #### arrays of floats
 
@@ -395,7 +395,7 @@ function rand!(r::MersenneTwister, A::Array{UInt128}, ::SamplerType{UInt128})
         end
     end
     if n > 0
-        u = rand_ui2x52_raw(r)
+        u = rand(r, UInt2x52Raw())
         for i = 1:n
             @inbounds A[i] ⊻= u << (12*i)
         end

base/random/generation.jl

@@ -23,17 +23,17 @@ Sampler(rng::AbstractRNG, ::Type{T}, n::Repetition) where {T<:AbstractFloat} =
 # generic random generation function which can be used by RNG implementors
 # it is not defined as a fallback rand method as this could create ambiguities
 
-rand_generic(r::AbstractRNG, ::CloseOpen{Float16}) =
+rand(r::AbstractRNG, ::SamplerTrivial{CloseOpen{Float16}}) =
     Float16(reinterpret(Float32,
-                        (rand_ui10_raw(r) % UInt32 << 13) & 0x007fe000 | 0x3f800000) - 1)
+                        (rand(r, UInt10(UInt32)) << 13)  | 0x3f800000) - 1)
 
-rand_generic(r::AbstractRNG, ::CloseOpen{Float32}) =
-    reinterpret(Float32, rand_ui23_raw(r) % UInt32 & 0x007fffff | 0x3f800000) - 1
+rand(r::AbstractRNG, ::SamplerTrivial{CloseOpen{Float32}}) =
+    reinterpret(Float32, rand(r, UInt23()) | 0x3f800000) - 1
 
-rand_generic(r::AbstractRNG, ::Close1Open2_64) =
-    reinterpret(Float64, 0x3ff0000000000000 | rand(r, UInt64) & 0x000fffffffffffff)
+rand(r::AbstractRNG, ::SamplerTrivial{Close1Open2_64}) =
+    reinterpret(Float64, 0x3ff0000000000000 | rand(r, UInt52()))
 
-rand_generic(r::AbstractRNG, ::CloseOpen_64) = rand(r, Close1Open2()) - 1.0
+rand(r::AbstractRNG, ::SamplerTrivial{CloseOpen_64}) = rand(r, Close1Open2()) - 1.0
 
 #### BigFloat
 
@@ -101,11 +101,21 @@ rand(rng::AbstractRNG, sp::SamplerBigFloat{T}) where {T<:FloatInterval{BigFloat}
 
 ### random integers
 
-rand_ui10_raw(r::AbstractRNG) = rand(r, UInt16)
-rand_ui23_raw(r::AbstractRNG) = rand(r, UInt32)
+rand(r::AbstractRNG, ::SamplerTrivial{UInt10Raw{UInt16}}) = rand(r, UInt16)
+rand(r::AbstractRNG, ::SamplerTrivial{UInt23Raw{UInt32}}) = rand(r, UInt32)
 
-rand_ui52_raw(r::AbstractRNG) = reinterpret(UInt64, rand(r, Close1Open2()))
-rand_ui52(r::AbstractRNG) = rand_ui52_raw(r) & 0x000fffffffffffff
+rand(r::AbstractRNG, ::SamplerTrivial{UInt52Raw{UInt64}}) =
+    _rand52(r, rng_native_52(r))
+
+_rand52(r::AbstractRNG, ::Type{Float64}) = reinterpret(UInt64, rand(r, Close1Open2()))
+_rand52(r::AbstractRNG, ::Type{UInt64})  = rand(r, UInt64)
+
+rand(r::AbstractRNG, ::SamplerTrivial{UInt10{UInt16}}) = rand(r, UInt10Raw()) & 0x03ff
+rand(r::AbstractRNG, ::SamplerTrivial{UInt23{UInt32}}) = rand(r, UInt23Raw()) & 0x007fffff
+rand(r::AbstractRNG, ::SamplerTrivial{UInt52{UInt64}}) = rand(r, UInt52Raw()) & 0x000fffffffffffff
+
+rand(r::AbstractRNG, sp::SamplerTrivial{<:UniformBits{T}}) where {T} =
+        rand(r, uint_default(sp[])) % T
 
 ### random complex numbers
 
@@ -158,25 +168,28 @@ function SamplerRangeFast(r::AbstractUnitRange{T}) where T<:Union{Int128,UInt128
     SamplerRangeFast{UInt128,T}(first(r), bw, m, mask)
 end
 
-function rand_lteq(r::AbstractRNG, randfun, u::U, mask::U) where U<:Integer
+function rand_lteq(r::AbstractRNG, S, u::U, mask::U) where U<:Integer
     while true
-        x = randfun(r) & mask
+        x = rand(r, S) & mask
         x <= u && return x
     end
 end
 
+# helper function, to turn types to values, should be removed once we can do rand(Uniform(UInt))
+rand(rng::AbstractRNG, ::Val{T}) where {T} = rand(rng, T)
+
 function rand(rng::AbstractRNG, sp::SamplerRangeFast{UInt64,T}) where T
     a, bw, m, mask = sp.a, sp.bw, sp.m, sp.mask
-    x = bw <= 52 ? rand_lteq(rng, rand_ui52_raw, m, mask) :
-                   rand_lteq(rng, rng->rand(rng, UInt64), m, mask)
+    x = bw <= 52 ? rand_lteq(rng, UInt52Raw(), m, mask) :
+                   rand_lteq(rng, Val(UInt64), m, mask)
     (x + a % UInt64) % T
 end
 
 function rand(rng::AbstractRNG, sp::SamplerRangeFast{UInt128,T}) where T
     a, bw, m, mask = sp.a, sp.bw, sp.m, sp.mask
-    x = bw <= 52  ? rand_lteq(rng, rand_ui52_raw, m % UInt64, mask % UInt64) % UInt128 :
-        bw <= 104 ? rand_lteq(rng, rand_ui104_raw, m, mask) :
-                    rand_lteq(rng, rng->rand(rng, UInt128), m, mask)
+    x = bw <= 52  ? rand_lteq(rng, UInt52Raw(), m % UInt64, mask % UInt64) % UInt128 :
+        bw <= 104 ? rand_lteq(rng, UInt104Raw(), m, mask) :
+                    rand_lteq(rng, Val(UInt128), m, mask)
     x % T + a
 end
 

base/random/misc.jl

@@ -147,7 +147,7 @@ randsubseq(A::AbstractArray, p::Real) = randsubseq(GLOBAL_RNG, A, p)
 @inline function rand_lt(r::AbstractRNG, n::Int, mask::Int=nextpow2(n)-1)
     # this duplicates the functionality of rand(1:n), to optimize this special case
     while true
-        x = (rand_ui52_raw(r) % Int) & mask
+        x = rand(r, UInt52Raw(Int)) & mask
         x < n && return x
     end
 end

base/random/normal.jl

@@ -35,7 +35,7 @@ julia> randn(rng, Complex64, (2, 3))
 """
 @inline function randn(rng::AbstractRNG=GLOBAL_RNG)
     @inbounds begin
-        r = rand_ui52(rng)
+        r = rand(rng, UInt52())
         rabs = Int64(r>>1) # One bit for the sign
         idx = rabs & 0xFF
         x = ifelse(r % Bool, -rabs, rabs)*wi[idx+1]
@@ -95,7 +95,7 @@ julia> randexp(rng, 3, 3)
 """
 function randexp(rng::AbstractRNG=GLOBAL_RNG)
     @inbounds begin
-        ri = rand_ui52(rng)
+        ri = rand(rng, UInt52())
         idx = ri & 0xFF
         x = ri*we[idx+1]
         ri < ke[idx+1] && return x # 98.9% of the time we return here 1st try

base/random/random.jl

@@ -25,6 +25,43 @@ export srand,
 
 abstract type AbstractRNG end
 
+### integers
+
+# we define types which encode the generation of a specific number of bits
+# the "raw" version means that the unused bits are not zeroed
+
+abstract type UniformBits{T<:BitInteger} end
+
+struct UInt10{T}    <: UniformBits{T} end
+struct UInt10Raw{T} <: UniformBits{T} end
+
+struct UInt23{T}    <: UniformBits{T} end
+struct UInt23Raw{T} <: UniformBits{T} end
+
+struct UInt52{T}    <: UniformBits{T} end
+struct UInt52Raw{T} <: UniformBits{T} end
+
+struct UInt104{T}    <: UniformBits{T} end
+struct UInt104Raw{T} <: UniformBits{T} end
+
+struct UInt2x52{T}    <: UniformBits{T} end
+struct UInt2x52Raw{T} <: UniformBits{T} end
+
+uint_sup(::Type{<:Union{UInt10,UInt10Raw}}) = UInt16
+uint_sup(::Type{<:Union{UInt23,UInt23Raw}}) = UInt32
+uint_sup(::Type{<:Union{UInt52,UInt52Raw}}) = UInt64
+uint_sup(::Type{<:Union{UInt104,UInt104Raw}}) = UInt128
+uint_sup(::Type{<:Union{UInt2x52,UInt2x52Raw}}) = UInt128
+
+for UI = (:UInt10, :UInt10Raw, :UInt23, :UInt23Raw, :UInt52, :UInt52Raw,
+          :UInt104, :UInt104Raw, :UInt2x52, :UInt2x52Raw)
+    @eval begin
+        $UI(::Type{T}=uint_sup($UI)) where {T} = $UI{T}()
+        # useful for defining rand generically:
+        uint_default(::$UI) = $UI{uint_sup($UI)}()
+    end
+end
+
 ### floats
 
 abstract type FloatInterval{T<:AbstractFloat} end