Repository: JuliaLang/julia
File: stdlib/Random/src/misc.jl (Julia)

Code before:
large.) Technically, this process is known as "Bernoulli sampling" of `A`.

# Examples
```jldoctest
julia> randsubseq(Xoshiro(123), 1:8, 0.3)
2-element Vector{Int64}:
 4
 7
```
"""
randsubseq(A::AbstractArray, p::Real) = randsubseq(default_rng(), A, p)


## 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=nextpow(2, n)-1) = LessThan(n-1, Masked(mask, UInt52Raw(Int)))

## shuffle & shuffle!

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
    else
        tup
    end::typeof(tup)
end

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

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

In-place version of [`shuffle`](@ref): randomly permute `v` in-place,
optionally supplying the random-number generator `rng`.

# Examples
```jldoctest

Review comment (suggestion):
```suggestion
                shuffle!(rng, mem)
```

Code after:
large.) Technically, this process is known as "Bernoulli sampling" of `A`.

# Examples
```jldoctest
julia> randsubseq(Xoshiro(123), 1:8, 0.3)
2-element Vector{Int64}:
 4
 7
```
"""
randsubseq(A::AbstractArray, p::Real) = randsubseq(default_rng(), A, p)


## 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=nextpow(2, n)-1) = LessThan(n-1, Masked(mask, UInt52Raw(Int)))

## shuffle & shuffle!

function shuffle(rng::AbstractRNG, tup::NTuple{N}) where {N}
    # `@inline` and `@inbounds` are here to help escape analysis eliminate the `Memory` allocation
    #
    # * `@inline` might be necessary because escape analysis relies on everything
    #   touching the `Memory` being inlined because there's no interprocedural escape
    #   analysis yet, relevant WIP PR: https://github.com/JuliaLang/julia/pull/56849
    #
    # * `@inbounds` might be necessary because escape analysis requires any throws of
    #   `BoundsError` to be eliminated as dead code, because `BoundsError` stores the
    #   array itself, making the throw escape the array from the function, relevant
    #   WIP PR: https://github.com/JuliaLang/julia/pull/56167
    @inline let
        # use a narrow integer type to save stack space and prevent heap allocation
        Ind = if N ≤ typemax(UInt8)
            UInt8
        elseif N ≤ typemax(UInt16)
            UInt16
        else
            UInt
        end
        mem = @inbounds randperm!(rng, Memory{Ind}(undef, N))
        function closure(i::Int)
            @inbounds tup[mem[i]]
        end
        ntuple(closure, Val{N}())
    end
end

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