Skip to content

Commit

Permalink
Replace dynamic dispatch with runtime branch on rev keyword in sort…
Browse files Browse the repository at this point in the history
…perm (#47966)

(cherry picked from commit 6ac5159)
  • Loading branch information
LilithHafner authored and KristofferC committed Dec 28, 2022
1 parent 6ae7827 commit ef913f8
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 5 deletions.
21 changes: 16 additions & 5 deletions base/sort.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1560,8 +1560,14 @@ function sortperm(A::AbstractArray;
order::Ordering=Forward,
scratch::Union{Vector{<:Integer}, Nothing}=nothing,
dims...) #to optionally specify dims argument
ordr = ord(lt,by,rev,order)
if ordr === Forward && isa(A,Vector) && eltype(A)<:Integer
if rev === true
_sortperm(A; alg, order=ord(lt, by, true, order), scratch, dims...)
else
_sortperm(A; alg, order=ord(lt, by, nothing, order), scratch, dims...)
end
end
function _sortperm(A::AbstractArray; alg, order, scratch, dims...)
if order === Forward && isa(A,Vector) && eltype(A)<:Integer
n = length(A)
if n > 1
min, max = extrema(A)
Expand All @@ -1573,7 +1579,7 @@ function sortperm(A::AbstractArray;
end
end
ix = copymutable(LinearIndices(A))
sort!(ix; alg, order = Perm(ordr, vec(A)), scratch, dims...)
sort!(ix; alg, order = Perm(order, vec(A)), scratch, dims...)
end


Expand Down Expand Up @@ -1615,7 +1621,7 @@ julia> sortperm!(p, A; dims=2); p
2 4
```
"""
function sortperm!(ix::AbstractArray{T}, A::AbstractArray;
@inline function sortperm!(ix::AbstractArray{T}, A::AbstractArray;
alg::Algorithm=DEFAULT_UNSTABLE,
lt=isless,
by=identity,
Expand All @@ -1630,7 +1636,12 @@ function sortperm!(ix::AbstractArray{T}, A::AbstractArray;
if !initialized
ix .= LinearIndices(A)
end
sort!(ix; alg, order = Perm(ord(lt, by, rev, order), vec(A)), scratch, dims...)

if rev === true
sort!(ix; alg, order=Perm(ord(lt, by, true, order), vec(A)), scratch, dims...)
else
sort!(ix; alg, order=Perm(ord(lt, by, nothing, order), vec(A)), scratch, dims...)
end
end

# sortperm for vectors of few unique integers
Expand Down
22 changes: 22 additions & 0 deletions test/sorting.jl
Original file line number Diff line number Diff line change
Expand Up @@ -918,6 +918,28 @@ end
@test bsqs() === bsqs(missing, missing, InsertionSort)
end

function test_allocs()
v = rand(10)
i = randperm(length(v))
@test 1 == @allocations sort(v)
@test 0 == @allocations sortperm!(i, v)
@test 0 == @allocations sort!(i)
@test 0 == @allocations sortperm!(i, v, rev=true)
@test 1 == @allocations sortperm(v, rev=true)
@test 1 == @allocations sortperm(v, rev=false)
@test 0 == @allocations sortperm!(i, v, order=Base.Reverse)
@test 1 == @allocations sortperm(v)
@test 1 == @allocations sortperm(i, by=sqrt)
@test 0 == @allocations sort!(v, lt=(a, b) -> hash(a) < hash(b))
sort!(Int[], rev=false) # compile
@test 0 == @allocations sort!(i, rev=false)
rand!(i)
@test 0 == @allocations sort!(i, order=Base.Reverse)
end
@testset "Small calls do not unnecessarily allocate" begin
test_allocs()
end

# This testset is at the end of the file because it is slow.
@testset "searchsorted" begin
numTypes = [ Int8, Int16, Int32, Int64, Int128,
Expand Down

0 comments on commit ef913f8

Please sign in to comment.