unify the 3 rand_lt[eq] functions as composable samplers

Author: rfourquet

base/random/generation.jl

@@ -149,24 +149,6 @@ end
 
 #### helper functions
 
-function rand_lteq(r::AbstractRNG, S, u::U, mask::U) where U<:Integer
-    while true
-        x = rand(r, S) & mask
-        x <= u && return x
-    end
-end
-
-function rand_lteq(rng::AbstractRNG, S, u::T)::T where T
-    while true
-        x = rand(rng, S)
-        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)
-
 uint_sup(::Type{<:Union{Bool,BitInteger}}) = UInt32
 uint_sup(::Type{<:Union{Int64,UInt64}}) = UInt64
 uint_sup(::Type{<:Union{Int128,UInt128}}) = UInt128
@@ -193,22 +175,24 @@ end
 
 function rand(rng::AbstractRNG, sp::SamplerRangeFast{UInt32,T}) where T
     a, bw, m, mask = sp.a, sp.bw, sp.m, sp.mask
-    x = rand_lteq(rng, Val(UInt32), m, mask)
+    x = rand(rng, LessThan(m, Masked(mask, uniform(UInt32))))
     (x + a % UInt32) % T
 end
 
 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, UInt52Raw(), m, mask) :
-                   rand_lteq(rng, Val(UInt64), m, mask)
+    x = bw <= 52 ? rand(rng, LessThan(m, Masked(mask, UInt52Raw()))) :
+                   rand(rng, LessThan(m, Masked(mask, uniform(UInt64))))
     (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, UInt52Raw(), m % UInt64, mask % UInt64) % UInt128 :
-        bw <= 104 ? rand_lteq(rng, UInt104Raw(), m, mask) :
-                    rand_lteq(rng, Val(UInt128), m, mask)
+    x = bw <= 52  ?
+        rand(rng, LessThan(m % UInt64, Masked(mask % UInt64, UInt52Raw()))) % UInt128 :
+    bw <= 104 ?
+        rand(rng, LessThan(m, Masked(mask, UInt104Raw()))) :
+        rand(rng, LessThan(m, Masked(mask, uniform(UInt128))))
     x % T + a
 end
 
@@ -266,19 +250,20 @@ Sampler(::AbstractRNG, r::AbstractUnitRange{T},
         ::Repetition) where {T<:Union{Bool,BitInteger}} = SamplerRangeInt(r)
 
 rand(rng::AbstractRNG, sp::SamplerRangeInt{T,UInt32}) where {T<:Union{Bool,BitInteger}} =
-    (unsigned(sp.a) + rem_knuth(rand_lteq(rng, Val(UInt32), sp.u), sp.k)) % T
+    (unsigned(sp.a) + rem_knuth(rand(rng, LessThan(sp.u, uniform(UInt32))), sp.k)) % T
+
 
 # this function uses 52 bit entropy for small ranges of length <= 2^52
 function rand(rng::AbstractRNG, sp::SamplerRangeInt{T,UInt64}) where T<:BitInteger
-    x = sp.bw <= 52 ? rand_lteq(rng, UInt52(), sp.u) :
-                      rand_lteq(rng, Val(UInt64), sp.u)
+    x = sp.bw <= 52 ? rand(rng, LessThan(sp.u, UInt52())) :
+                      rand(rng, LessThan(sp.u, uniform(UInt64)))
     return ((sp.a % UInt64) + rem_knuth(x, sp.k)) % T
 end
 
 function rand(rng::AbstractRNG, sp::SamplerRangeInt{T,UInt128}) where T<:BitInteger
-    x = sp.bw <= 52  ? rand_lteq(rng, UInt52(UInt128), sp.u) :
-        sp.bw <= 104 ? rand_lteq(rng, UInt104(UInt128), sp.u) :
-                       rand_lteq(rng, Val(UInt128), sp.u)
+    x = sp.bw <= 52  ? rand(rng, LessThan(sp.u, UInt52(UInt128))) :
+        sp.bw <= 104 ? rand(rng, LessThan(sp.u, UInt104(UInt128))) :
+                       rand(rng, LessThan(sp.u, uniform(UInt128)))
     return ((sp.a % UInt128) + rem_knuth(x, sp.k)) % T
 end
 

base/random/misc.jl

@@ -141,17 +141,10 @@ large.) Technically, this process is known as "Bernoulli sampling" of `A`.
 randsubseq(A::AbstractArray, p::Real) = randsubseq(GLOBAL_RNG, A, p)
 
 
-## rand_lt (helper function)
-
-"Return a random `Int` (masked with `mask`) in ``[0, n)``, when `n <= 2^52`."
-@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(r, UInt52Raw(Int)) & mask
-        x < n && return x
-    end
-end
+## rand Less Than Masked 52 bits (helper function)
 
+"Return a sampler generating a random `Int` (masked with `mask`) in ``[0, n)``, when `n <= 2^52`."
+ltm52(n::Int, mask::Int=nextpow2(n)-1) = LessThan(n-1, Masked(mask, UInt52Raw(Int)))
 
 ## shuffle & shuffle!
 
@@ -191,7 +184,7 @@ function shuffle!(r::AbstractRNG, a::AbstractArray)
     mask = nextpow2(n) - 1
     for i = n:-1:2
         (mask >> 1) == i && (mask >>= 1)
-        j = 1 + rand_lt(r, i, mask)
+        j = 1 + rand(r, ltm52(i, mask))
         a[i], a[j] = a[j], a[i]
     end
     return a
@@ -277,7 +270,7 @@ function randperm!(r::AbstractRNG, a::Array{<:Integer})
     a[1] = 1
     mask = 3
     @inbounds for i = 2:n
-        j = 1 + rand_lt(r, i, mask)
+        j = 1 + rand(r, ltm52(i, mask))
         if i != j # a[i] is uninitialized (and could be #undef)
             a[i] = a[j]
         end
@@ -339,7 +332,7 @@ function randcycle!(r::AbstractRNG, a::Array{<:Integer})
     a[1] = 1
     mask = 3
     @inbounds for i = 2:n
-        j = 1 + rand_lt(r, i-1, mask)
+        j = 1 + rand(r, ltm52(i-1, mask))
         a[i] = a[j]
         a[j] = i
         i == 1+mask && (mask = 2mask + 1)

base/random/random.jl

@@ -105,7 +105,7 @@ Sampler(rng::AbstractRNG, sp::Sampler, ::Repetition) =
 Sampler(rng::AbstractRNG, X) = Sampler(rng, X, Val(Inf))
 Sampler(rng::AbstractRNG, ::Type{X}) where {X} = Sampler(rng, X, Val(Inf))
 
-#### pre-defined useful Sampler subtypes
+#### pre-defined useful Sampler types
 
 # default fall-back for types
 struct SamplerType{T} <: Sampler end
@@ -138,6 +138,38 @@ struct SamplerTag{T,S} <: Sampler
 end
 
 
+#### helper samplers
+
+##### Adapter to generate a randome value in [0, n]
+
+struct LessThan{T<:Integer,S} <: Sampler
+    sup::T
+    s::S    # the scalar specification/sampler to feed to rand
+end
+
+function rand(rng::AbstractRNG, sp::LessThan)
+    while true
+        x = rand(rng, sp.s)
+        x <= sp.sup && return x
+    end
+end
+
+struct Masked{T<:Integer,S} <: Sampler
+    mask::T
+    s::S
+end
+
+rand(rng::AbstractRNG, sp::Masked) = rand(rng, sp.s) & sp.mask
+
+##### Uniform
+
+struct UniformT{T} <: Sampler end
+
+uniform(::Type{T}) where {T} = UniformT{T}()
+
+rand(rng::AbstractRNG, ::UniformT{T}) where {T} = rand(rng, T)
+
+
 ### machinery for generation with Sampler
 
 # This describes how to generate random scalars or arrays, by generating a Sampler