Skip to content

Commit ef913f8

Browse files
LilithHafnerKristofferC
authored andcommitted
Replace dynamic dispatch with runtime branch on rev keyword in sortperm (#47966)
(cherry picked from commit 6ac5159)
1 parent 6ae7827 commit ef913f8

File tree

2 files changed

+38
-5
lines changed

2 files changed

+38
-5
lines changed

base/sort.jl

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1560,8 +1560,14 @@ function sortperm(A::AbstractArray;
15601560
order::Ordering=Forward,
15611561
scratch::Union{Vector{<:Integer}, Nothing}=nothing,
15621562
dims...) #to optionally specify dims argument
1563-
ordr = ord(lt,by,rev,order)
1564-
if ordr === Forward && isa(A,Vector) && eltype(A)<:Integer
1563+
if rev === true
1564+
_sortperm(A; alg, order=ord(lt, by, true, order), scratch, dims...)
1565+
else
1566+
_sortperm(A; alg, order=ord(lt, by, nothing, order), scratch, dims...)
1567+
end
1568+
end
1569+
function _sortperm(A::AbstractArray; alg, order, scratch, dims...)
1570+
if order === Forward && isa(A,Vector) && eltype(A)<:Integer
15651571
n = length(A)
15661572
if n > 1
15671573
min, max = extrema(A)
@@ -1573,7 +1579,7 @@ function sortperm(A::AbstractArray;
15731579
end
15741580
end
15751581
ix = copymutable(LinearIndices(A))
1576-
sort!(ix; alg, order = Perm(ordr, vec(A)), scratch, dims...)
1582+
sort!(ix; alg, order = Perm(order, vec(A)), scratch, dims...)
15771583
end
15781584

15791585

@@ -1615,7 +1621,7 @@ julia> sortperm!(p, A; dims=2); p
16151621
2 4
16161622
```
16171623
"""
1618-
function sortperm!(ix::AbstractArray{T}, A::AbstractArray;
1624+
@inline function sortperm!(ix::AbstractArray{T}, A::AbstractArray;
16191625
alg::Algorithm=DEFAULT_UNSTABLE,
16201626
lt=isless,
16211627
by=identity,
@@ -1630,7 +1636,12 @@ function sortperm!(ix::AbstractArray{T}, A::AbstractArray;
16301636
if !initialized
16311637
ix .= LinearIndices(A)
16321638
end
1633-
sort!(ix; alg, order = Perm(ord(lt, by, rev, order), vec(A)), scratch, dims...)
1639+
1640+
if rev === true
1641+
sort!(ix; alg, order=Perm(ord(lt, by, true, order), vec(A)), scratch, dims...)
1642+
else
1643+
sort!(ix; alg, order=Perm(ord(lt, by, nothing, order), vec(A)), scratch, dims...)
1644+
end
16341645
end
16351646

16361647
# sortperm for vectors of few unique integers

test/sorting.jl

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -918,6 +918,28 @@ end
918918
@test bsqs() === bsqs(missing, missing, InsertionSort)
919919
end
920920

921+
function test_allocs()
922+
v = rand(10)
923+
i = randperm(length(v))
924+
@test 1 == @allocations sort(v)
925+
@test 0 == @allocations sortperm!(i, v)
926+
@test 0 == @allocations sort!(i)
927+
@test 0 == @allocations sortperm!(i, v, rev=true)
928+
@test 1 == @allocations sortperm(v, rev=true)
929+
@test 1 == @allocations sortperm(v, rev=false)
930+
@test 0 == @allocations sortperm!(i, v, order=Base.Reverse)
931+
@test 1 == @allocations sortperm(v)
932+
@test 1 == @allocations sortperm(i, by=sqrt)
933+
@test 0 == @allocations sort!(v, lt=(a, b) -> hash(a) < hash(b))
934+
sort!(Int[], rev=false) # compile
935+
@test 0 == @allocations sort!(i, rev=false)
936+
rand!(i)
937+
@test 0 == @allocations sort!(i, order=Base.Reverse)
938+
end
939+
@testset "Small calls do not unnecessarily allocate" begin
940+
test_allocs()
941+
end
942+
921943
# This testset is at the end of the file because it is slow.
922944
@testset "searchsorted" begin
923945
numTypes = [ Int8, Int16, Int32, Int64, Int128,

0 commit comments

Comments
 (0)