Sign-aware computation of midpoint for sorting (fixes #33977) (#34106)

Author: timholy

base/sort.jl

@@ -164,6 +164,10 @@ same thing as `partialsort!` but leaving `v` unmodified.
 partialsort(v::AbstractVector, k::Union{Integer,OrdinalRange}; kws...) =
     partialsort!(copymutable(v), k; kws...)
 
+# This implementation of `midpoint` is performance-optimized but safe
+# only if `lo <= hi`.
+midpoint(lo::T, hi::T) where T<:Integer = lo + ((hi - lo) >>> 0x01)
+midpoint(lo::Integer, hi::Integer) = midpoint(promote(lo, hi)...)
 
 # reference on sorted binary search:
 #   http://www.tbray.org/ongoing/When/200x/2003/03/22/Binary
@@ -175,7 +179,7 @@ function searchsortedfirst(v::AbstractVector, x, lo::T, hi::T, o::Ordering) wher
     lo = lo - u
     hi = hi + u
     @inbounds while lo < hi - u
-        m = (lo + hi) >>> 1
+        m = midpoint(lo, hi)
         if lt(o, v[m], x)
             lo = m
         else
@@ -192,7 +196,7 @@ function searchsortedlast(v::AbstractVector, x, lo::T, hi::T, o::Ordering) where
     lo = lo - u
     hi = hi + u
     @inbounds while lo < hi - u
-        m = (lo + hi) >>> 1
+        m = midpoint(lo, hi)
         if lt(o, x, v[m])
             hi = m
         else
@@ -210,7 +214,7 @@ function searchsorted(v::AbstractVector, x, ilo::T, ihi::T, o::Ordering) where T
     lo = ilo - u
     hi = ihi + u
     @inbounds while lo < hi - u
-        m = (lo + hi) >>> 1
+        m = midpoint(lo, hi)
         if lt(o, v[m], x)
             lo = m
         elseif lt(o, x, v[m])
@@ -487,7 +491,7 @@ end
 
 @inline function selectpivot!(v::AbstractVector, lo::Integer, hi::Integer, o::Ordering)
     @inbounds begin
-        mi = (lo+hi)>>>1
+        mi = midpoint(lo, hi)
 
         # sort v[mi] <= v[lo] <= v[hi] such that the pivot is immediately in place
         if lt(o, v[lo], v[mi])
@@ -552,7 +556,7 @@ function sort!(v::AbstractVector, lo::Integer, hi::Integer, a::MergeSortAlg, o::
     @inbounds if lo < hi
         hi-lo <= SMALL_THRESHOLD && return sort!(v, lo, hi, SMALL_ALGORITHM, o)
 
-        m = (lo+hi)>>>1
+        m = midpoint(lo, hi)
         (length(t) < m-lo+1) && resize!(t, m-lo+1)
 
         sort!(v, lo,  m,  a, o, t)

test/offsetarray.jl

@@ -475,6 +475,19 @@ v = OffsetArray(rand(8), (-2,))
 @test sortslices(A, dims=2) == OffsetArray(sortslices(parent(A), dims=2), A.offsets)
 @test sort(A, dims=1) == OffsetArray(sort(parent(A), dims=1), A.offsets)
 @test sort(A, dims=2) == OffsetArray(sort(parent(A), dims=2), A.offsets)
+# Issue #33977
+soa = OffsetArray([2,2,3], -2)
+@test searchsorted(soa, 1) == -1:-2
+@test searchsortedfirst(soa, 1) == -1
+@test searchsortedlast(soa, 1) == -2
+@test first(sort!(soa; alg=QuickSort)) == 2
+@test first(sort!(soa; alg=MergeSort)) == 2
+soa = OffsetArray([2,2,3], typemax(Int)-4)
+@test searchsorted(soa, 1) == typemax(Int)-3:typemax(Int)-4
+@test searchsortedfirst(soa, 2) == typemax(Int) - 3
+@test searchsortedlast(soa, 2) == typemax(Int) - 2
+@test first(sort!(soa; alg=QuickSort)) == 2
+@test first(sort!(soa; alg=MergeSort)) == 2
 
 @test mapslices(sort, A, dims=1) == OffsetArray(mapslices(sort, parent(A), dims=1), A.offsets)
 @test mapslices(sort, A, dims=2) == OffsetArray(mapslices(sort, parent(A), dims=2), A.offsets)

test/sorting.jl

@@ -11,6 +11,21 @@ using Test
     @test ReverseOrdering(Forward) == ReverseOrdering() == Reverse
 end
 
+@testset "midpoint" begin
+    @test Base.Sort.midpoint(1, 3) === 2
+    @test Base.Sort.midpoint(2, 4) === 3
+    @test 2 <= Base.Sort.midpoint(1, 4) <= 3
+    @test Base.Sort.midpoint(-3, -1) === -2
+    @test Base.Sort.midpoint(-4, -2) === -3
+    @test -3 <= Base.Sort.midpoint(-4, -1) <= -2
+    @test Base.Sort.midpoint(-1, 1) ===  0
+    @test -1 <= Base.Sort.midpoint(-2, 1) <= 0
+    @test 0 <= Base.Sort.midpoint(-1, 2) <= 1
+    @test Base.Sort.midpoint(-2, 2) ===  0
+    @test Base.Sort.midpoint(typemax(Int)-2, typemax(Int)) === typemax(Int)-1
+    @test Base.Sort.midpoint(typemin(Int), typemin(Int)+2) === typemin(Int)+1
+    @test -1 <= Base.Sort.midpoint(typemin(Int), typemax(Int)) <= 0
+end
 
 @testset "sort" begin
     @test sort([2,3,1]) == [1,2,3] == sort([2,3,1]; order=Forward)