enable rand(::Union{AbstractSet,Associative}) (#22228)

base/random.jl

@@ -257,13 +257,15 @@ globalRNG() = GLOBAL_RNG
 Pick a random element or array of random elements from the set of values specified by `S`; `S` can be
 
 * an indexable collection (for example `1:n` or `['x','y','z']`), or
-* a `Dict`, a `Set` or an `IntSet`, or
+* an `Associative` or `AbstractSet` object, or
 * a type: the set of values to pick from is then equivalent to `typemin(S):typemax(S)` for
   integers (this is not applicable to [`BigInt`](@ref)), and to ``[0, 1)`` for floating
   point numbers;
 
 `S` defaults to [`Float64`](@ref).
 
+# Examples
+
 ```julia-repl
 julia> rand(Int, 2)
 2-element Array{Int64,1}:
@@ -273,6 +275,14 @@ julia> rand(Int, 2)
 julia> rand(MersenneTwister(0), Dict(1=>2, 3=>4))
 1=>2
 ```
+
+!!! note
+    The complexity of `rand(rng, s::Union{Associative,AbstractSet})`
+    is linear in the length of `s`, unless an optimized method with
+    constant complexity is available, which is the case for `Dict`,
+    `Set` and `IntSet`. For more than a few calls, use `rand(rng,
+    collect(s))` instead, or either `rand(rng, Dict(s))` or `rand(rng,
+    Set(s))` as appropriate.
 """
 @inline rand() = rand(GLOBAL_RNG, CloseOpen)
 @inline rand(T::Type) = rand(GLOBAL_RNG, T)
@@ -362,9 +372,8 @@ function rand(r::AbstractRNG, t::Dict)
         Base.isslotfilled(t, i) && @inbounds return (t.keys[i] => t.vals[i])
     end
 end
-rand(t::Dict) = rand(GLOBAL_RNG, t)
+
 rand(r::AbstractRNG, s::Set) = rand(r, s.dict).first
-rand(s::Set) = rand(GLOBAL_RNG, s)
 
 function rand(r::AbstractRNG, s::IntSet)
     isempty(s) && throw(ArgumentError("collection must be non-empty"))
@@ -378,7 +387,16 @@ function rand(r::AbstractRNG, s::IntSet)
     end
 end
 
-rand(s::IntSet) = rand(GLOBAL_RNG, s)
+function nth(iter, n::Integer)::eltype(iter)
+    for (i, x) in enumerate(iter)
+        i == n && return x
+    end
+end
+nth(iter::AbstractArray, n::Integer) = iter[n]
+
+rand(r::AbstractRNG, s::Union{Associative,AbstractSet}) = nth(s, rand(r, 1:length(s)))
+
+rand(s::Union{Associative,AbstractSet}) = rand(GLOBAL_RNG, s)
 
 ## Arrays of random numbers
 
@@ -407,14 +425,16 @@ function rand!(r::AbstractRNG, A::AbstractArray, s::Union{Dict,Set,IntSet})
     A
 end
 
-rand!(A::AbstractArray, s::Union{Dict,Set,IntSet}) = rand!(GLOBAL_RNG, A, s)
+# avoid linear complexity for repeated calls with generic containers
+rand!(r::AbstractRNG, A::AbstractArray, s::Union{Associative,AbstractSet}) = rand!(r, A, collect(s))
+
+rand!(A::AbstractArray, s::Union{Associative,AbstractSet}) = rand!(GLOBAL_RNG, A, s)
 
-rand(r::AbstractRNG, s::Dict{K,V}, dims::Dims) where {K,V} = rand!(r, Array{Pair{K,V}}(dims), s)
-rand(r::AbstractRNG, s::Set{T}, dims::Dims) where {T} = rand!(r, Array{T}(dims), s)
-rand(r::AbstractRNG, s::IntSet, dims::Dims) = rand!(r, Array{Int}(dims), s)
-rand(r::AbstractRNG, s::Union{Dict,Set,IntSet}, dims::Integer...) = rand(r, s, convert(Dims, dims))
-rand(s::Union{Dict,Set,IntSet}, dims::Integer...) = rand(GLOBAL_RNG, s, convert(Dims, dims))
-rand(s::Union{Dict,Set,IntSet}, dims::Dims) = rand(GLOBAL_RNG, s, dims)
+rand(r::AbstractRNG, s::Associative{K,V}, dims::Dims) where {K,V} = rand!(r, Array{Pair{K,V}}(dims), s)
+rand(r::AbstractRNG, s::AbstractSet{T}, dims::Dims) where {T} = rand!(r, Array{T}(dims), s)
+rand(r::AbstractRNG, s::Union{Associative,AbstractSet}, dims::Integer...) = rand(r, s, convert(Dims, dims))
+rand(s::Union{Associative,AbstractSet}, dims::Integer...) = rand(GLOBAL_RNG, s, convert(Dims, dims))
+rand(s::Union{Associative,AbstractSet}, dims::Dims) = rand(GLOBAL_RNG, s, dims)
 
 # MersenneTwister
 

base/test.jl

@@ -11,14 +11,16 @@ test set that throws on the first failure. Users can choose to wrap
 their tests in (possibly nested) test sets that will store results
 and summarize them at the end of the test set with `@testset`.
 """
+:Test # cf. #22288
+
 module Test
 
 export @test, @test_throws, @test_broken, @test_skip, @test_warn, @test_nowarn
 export @testset
 # Legacy approximate testing functions, yet to be included
 export @inferred
 export detect_ambiguities
-export GenericString
+export GenericString, GenericSet, GenericDict
 
 #-----------------------------------------------------------------------
 
@@ -1203,4 +1205,38 @@ Base.convert(::Type{GenericString}, s::AbstractString) = GenericString(s)
 Base.endof(s::GenericString) = endof(s.string)
 Base.next(s::GenericString, i::Int) = next(s.string, i)
 
+"""
+The `GenericSet` can be used to test generic set APIs that program to
+the `AbstractSet` interface, in order to ensure that functions can work
+with set types besides the standard `Set` and `IntSet` types.
+"""
+struct GenericSet{T} <: AbstractSet{T}
+    s::AbstractSet{T}
+end
+
+"""
+The `GenericDict` can be used to test generic dict APIs that program to
+the `Associative` interface, in order to ensure that functions can work
+with associative types besides the standard `Dict` type.
+"""
+struct GenericDict{K,V} <: Associative{K,V}
+    s::Associative{K,V}
+end
+
+for (G, A) in ((GenericSet, AbstractSet),
+               (GenericDict, Associative))
+    @eval begin
+        Base.convert(::Type{$G}, s::$A) = $G(s)
+        Base.done(s::$G, state) = done(s.s, state)
+        Base.next(s::$G, state) = next(s.s, state)
+    end
+    for f in (:eltype, :isempty, :length, :start)
+        @eval begin
+            Base.$f(s::$G) = $f(s.s)
+        end
+    end
+end
+
+Base.get(s::GenericDict, x, y) = get(s.s, x, y)
+
 end # module

test/random.jl

@@ -309,13 +309,18 @@ for rng in ([], [MersenneTwister(0)], [RandomDevice()])
     ftypes = [Float16, Float32, Float64]
     cftypes = [Complex32, Complex64, Complex128, ftypes...]
     types = [Bool, Char, Base.BitInteger_types..., ftypes...]
-    collections = [IntSet(rand(1:100, 20))                => Int,
-                   Set(rand(Int, 20))                     => Int,
-                   Dict(zip(rand(Int,10), rand(Int, 10))) => Pair{Int,Int},
-                   1:100                                  => Int,
-                   rand(Int, 100)                         => Int,
-                   Int                                    => Int,
-                   Float64                                => Float64]
+    randset = Set(rand(Int, 20))
+    randdict = Dict(zip(rand(Int,10), rand(Int, 10)))
+    collections = [IntSet(rand(1:100, 20)) => Int,
+                   randset                 => Int,
+                   GenericSet(randset)     => Int,
+                   randdict                => Pair{Int,Int},
+                   GenericDict(randdict)   => Pair{Int,Int},
+                   1:100                   => Int,
+                   rand(Int, 100)          => Int,
+                   Int                     => Int,
+                   Float64                 => Float64]
+
     b2 = big(2)
     u3 = UInt(3)
     for f in [rand, randn, randexp]
@@ -352,7 +357,8 @@ for rng in ([], [MersenneTwister(0)], [RandomDevice()])
             end
         end
     end
-    for C in [1:0, Dict(), Set(), IntSet(), Int[]]
+    for C in [1:0, Dict(), Set(), IntSet(), Int[],
+              GenericDict(Dict()), GenericSet(Set())]
         @test_throws ArgumentError rand(rng..., C)
         @test_throws ArgumentError rand(rng..., C, 5)
     end