Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add shuffle(::NTuple) to Random #56906

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

add shuffle(::NTuple) to Random #56906

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 30 additions & 0 deletions stdlib/Random/src/misc.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -183,6 +183,36 @@ ltm52(n::Int, mask::Int=nextpow(2, n)-1) = LessThan(n-1, Masked(mask, UInt52Raw(

## shuffle & shuffle!

function _tuple_without_element_at_index(tup::(Tuple{T, Vararg{T, N}} where {T}), i::Int) where {N}
clo = let tup = tup, i = i
function closure(j::Int)
k = if j < i
j
else
j + 1
end
tup[k]
end
end
ntuple(clo, Val{N}())
end

function shuffle(rng::AbstractRNG, tup::(Tuple{Vararg{T, N}} where {T})) where {N}
if tup isa Tuple{Any, Vararg}
let i = rand(rng, Base.OneTo(N)) # Fisher–Yates shuffle
s = _tuple_without_element_at_index(tup, i)
t = shuffle(rng, s)
(tup[i], t...)
end
Comment on lines +186 to +206
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's also possible to rely on Memory getting stack-allocated, avoiding recursion:

Suggested change
function _tuple_without_element_at_index(tup::(Tuple{T, Vararg{T, N}} where {T}), i::Int) where {N}
clo = let tup = tup, i = i
function closure(j::Int)
k = if j < i
j
else
j + 1
end
tup[k]
end
end
ntuple(clo, Val{N}())
end
function shuffle(rng::AbstractRNG, tup::(Tuple{Vararg{T, N}} where {T})) where {N}
if tup isa Tuple{Any, Vararg}
let i = rand(rng, Base.OneTo(N)) # Fisher–Yates shuffle
s = _tuple_without_element_at_index(tup, i)
t = shuffle(rng, s)
(tup[i], t...)
end
function shuffle(rng::AbstractRNG, tup::(Tuple{Vararg{T, N}} where {T})) where {N}
if tup isa Tuple{Any, Vararg}
@inline let # `@inline` and `@inbounds` are here to help escape analysis
clo = @inbounds let mem = Memory{UInt16}(undef, N) # use `UInt16` to save stack space/prevent heap allocation
copyto!(mem, Base.OneTo(N))
shuffle!(mem)
let mem = mem, tup = tup
function closure(i::Int)
@inbounds tup[mem[i]]
end
end
end
ntuple(clo, Val{N}())
end

This suggestion depends on #56847 getting merged first.

else
tup
end::typeof(tup)
end

function shuffle(tup::NTuple)
shuffle(default_rng(), tup)
end

"""
shuffle!([rng=default_rng(),] v::AbstractArray)

Expand Down
28 changes: 28 additions & 0 deletions stdlib/Random/test/runtests.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1032,6 +1032,34 @@ end
@test maximum(m) <= 0.106
end

@testset "`shuffle(::NTuple)`" begin
@testset "sorted" begin
for n ∈ 0:20
tup = ntuple(identity, n)
@test tup === sort(@inferred shuffle(tup))
end
end
@testset "not identity" begin
function is_not_identity_at_least_once()
function f(::Any)
tup = ntuple(identity, 9)
tup !== shuffle(tup)
end
@test any(f, 1:1000000)
end
is_not_identity_at_least_once()
end
@testset "no heap allocation" begin
function no_heap_allocation(tup)
@test iszero(@allocated shuffle(tup))
end
for n ∈ 0:9
tup = ntuple(identity, n)
no_heap_allocation(tup)
end
end
end

# issue #42752
# test that running finalizers that launch tasks doesn't change RNG stream
function f42752(do_gc::Bool, cell = (()->Any[[]])())
Expand Down
Loading