Fix overflow in searchsorted* (#286)

* partly fix overflow in searchsorte*

* version bump to v1.11.2

* fix searchsorted insertion indices

* forward searchsorted* to parent

* Add searchsorted test for typemin(Int)

* Tests for nested offset arrays

Author: jishnub

Project.toml

@@ -1,6 +1,6 @@
 name = "OffsetArrays"
 uuid = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
-version = "1.11.1"
+version = "1.11.2"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/OffsetArrays.jl

@@ -771,74 +771,25 @@ centered(A::AbstractArray, cp::Dims=center(A)) = OffsetArray(A, .-cp)
 
 centered(A::AbstractArray, i::CartesianIndex) = centered(A, Tuple(i))
 
-####
-# work around for segfault in searchsorted*
-#  https://github.com/JuliaLang/julia/issues/33977
-####
-
-function _safe_searchsorted(v::OffsetArray, x, ilo::T, ihi::T, o::Base.Ordering) where T<:Integer
-    u = T(1)
-    lo = ilo - u
-    hi = ihi + u
-    @inbounds while lo < hi - u
-        m = (lo + hi) ÷ 2
-        if Base.lt(o, v[m], x)
-            lo = m
-        elseif Base.lt(o, x, v[m])
-            hi = m
-        else
-            a = searchsortedfirst(v, x, max(lo,ilo), m, o)
-            b = searchsortedlast(v, x, m, min(hi,ihi), o)
-            return a : b
-        end
+# we may pass the searchsorted* functions to the parent, and wrap the offset
+for f in [:searchsortedfirst, :searchsortedlast, :searchsorted]
+    _safe_f = Symbol("_safe_" * String(f))
+    @eval function $_safe_f(v::OffsetVector, x, ilo, ihi, o::Base.Ordering)
+        offset = v.offsets[1]
+        $f(parent(v), x, ilo - offset, ihi - offset, o) .+ offset
     end
-    return (lo + 1) : (hi - 1)
-end
-function _safe_searchsortedfirst(v::OffsetArray, x, lo::T, hi::T, o::Base.Ordering) where T<:Integer
-    u = T(1)
-    lo = lo - u
-    hi = hi + u
-    @inbounds while lo < hi - u
-        m = (lo + hi) ÷ 2
-        if Base.lt(o, v[m], x)
-            lo = m
-        else
-            hi = m
-        end
-    end
-    return hi
-end
-function _safe_searchsortedlast(v::OffsetArray, x, lo::T, hi::T, o::Base.Ordering) where T<:Integer
-    u = T(1)
-    lo = lo - u
-    hi = hi + u
-    @inbounds while lo < hi - u
-        m = (lo + hi) ÷ 2
-        if Base.lt(o, x, v[m])
-            hi = m
-        else
-            lo = m
-        end
-    end
-    return lo
+    @eval Base.$f(v::OffsetVector, x, ilo::T, ihi::T, o::Base.Ordering) where T<:Integer =
+        $_safe_f(v, x, ilo, ihi, o)
 end
 
-if VERSION ≤ v"1.2"
+if VERSION <= v"1.2"
     # ambiguity warnings in earlier versions
-    Base.searchsorted(v::OffsetArray, x, ilo::Int, ihi::Int, o::Base.Ordering) =
-        _safe_searchsorted(v, x, ilo, ihi, o)
-    Base.searchsortedfirst(v::OffsetArray, x, lo::Int, hi::Int, o::Base.Ordering) =
-        _safe_searchsortedfirst(v, x, lo, hi, o)
-    Base.searchsortedlast(v::OffsetArray, x, lo::Int, hi::Int, o::Base.Ordering) =
-        _safe_searchsortedlast(v, x, lo, hi, o)
-end
-
-Base.searchsorted(v::OffsetArray, x, ilo::T, ihi::T, o::Base.Ordering) where T<:Integer =
-    _safe_searchsorted(v, x, ilo, ihi, o)
-Base.searchsortedfirst(v::OffsetArray, x, lo::T, hi::T, o::Base.Ordering) where T<:Integer =
-    _safe_searchsortedfirst(v, x, lo, hi, o)
-Base.searchsortedlast(v::OffsetArray, x, lo::T, hi::T, o::Base.Ordering) where T<:Integer =
-    _safe_searchsortedlast(v, x, lo, hi, o)
+    for f in [:searchsortedfirst, :searchsortedlast, :searchsorted]
+        _safe_f = Symbol("_safe_" * String(f))
+        @eval Base.$f(v::OffsetVector, x, ilo::Int, ihi::Int, o::Base.Ordering) =
+            $_safe_f(v, x, ilo, ihi, o)
+    end
+end
 
 if VERSION < v"1.1.0-DEV.783"
     Base.copyfirst!(dest::OffsetArray, src::OffsetArray) = (maximum!(parent(dest), parent(src)); return dest)

test/runtests.jl

@@ -2320,6 +2320,73 @@ end
     @test searchsorted(o,  2) ==  0:-1
     @test searchsorted(o,  5) ==  2:2
     @test searchsorted(o,  6) ==  3:2
+
+    if VERSION > v"1.2"
+        # OffsetVector of another offset vector
+        v = OffsetVector(Base.IdentityUnitRange(4:10),-2)
+        @test searchsortedfirst(v, first(v)-1) == firstindex(v)
+        for i in axes(v,1)
+            @test searchsortedfirst(v, v[i]) == i
+        end
+        @test searchsortedfirst(v, last(v)+1) == lastindex(v)+1
+        @test searchsortedlast(v, first(v)-1) == firstindex(v)-1
+        for i in axes(v,1)
+            @test searchsortedlast(v, v[i]) == i
+        end
+        @test searchsortedlast(v, last(v)+1) == lastindex(v)
+        @test searchsorted(v, first(v)-1) === firstindex(v) .+ (0:-1)
+        for i in axes(v,1)
+            @test searchsorted(v, v[i]) == i:i
+        end
+        @test searchsorted(v, last(v)+1) === lastindex(v) .+ (1:0)
+    end
+
+    v = OffsetVector{Float64, OffsetVector{Float64, Vector{Float64}}}(OffsetVector([2,2,3,3,3,4], 3), 4)
+    @test searchsortedfirst(v, minimum(v)-1) == firstindex(v)
+    for el in unique(v)
+        @test searchsortedfirst(v, el) == findfirst(isequal(el), v)
+    end
+    @test searchsortedfirst(v, maximum(v)+1) == lastindex(v)+1
+
+    @test searchsortedlast(v, minimum(v)-1) == firstindex(v)-1
+    for el in unique(v)
+        @test searchsortedlast(v, el) == findlast(isequal(el), v)
+    end
+    @test searchsortedlast(v, maximum(v)+1) == lastindex(v)
+
+    @test searchsorted(v, minimum(v)-1) === firstindex(v) .+ (0:-1)
+    for el in unique(v)
+        @test searchsorted(v, el) == findfirst(isequal(el), v):findlast(isequal(el), v)
+    end
+    @test searchsorted(v, maximum(v)+1) === lastindex(v) .+ (1:0)
+
+    soa = OffsetArray([2,2,3], typemax(Int)-3)
+    @test searchsortedfirst(soa, 1) == firstindex(soa) == typemax(Int)-2
+    @test searchsortedfirst(soa, 2) == firstindex(soa) == typemax(Int)-2
+    @test searchsortedfirst(soa, 3) == lastindex(soa) == typemax(Int)
+
+    soa = OffsetArray([2,2,3], typemin(Int))
+    @test searchsortedlast(soa, 2) == firstindex(soa) + 1 == typemin(Int) + 2
+    @test searchsortedlast(soa, 3) == lastindex(soa) == typemin(Int) + 3
+    @test searchsortedlast(soa, 1) == typemin(Int)
+
+    soa = OffsetArray([2,2,3], typemax(Int)-4)
+    @test searchsorted(soa, 1) === firstindex(soa) .+ (0:-1)
+    @test searchsorted(soa, 2) == firstindex(soa) .+ (0:1) == typemax(Int) .+ (-3:-2)
+    @test searchsorted(soa, 3) == lastindex(soa) .+ (0:0) == typemax(Int) .+ (-1:-1)
+    @test searchsorted(soa, 4) === lastindex(soa) .+ (1:0)
+
+    soa = OffsetArray([2,2,3], typemax(Int)-3)
+    @test searchsorted(soa, 1) === firstindex(soa) .+ (0:-1)
+    @test searchsorted(soa, 2) == firstindex(soa) .+ (0:1) == typemax(Int) .+ (-2:-1)
+    @test searchsorted(soa, 3) == lastindex(soa) .+ (0:0) == typemax(Int) .+ (0:0)
+    @test searchsorted(soa, 4) === lastindex(soa) .+ (1:0)
+
+    soa = OffsetArray([2,2,3], typemin(Int))
+    @test searchsorted(soa, 1) === firstindex(soa) .+ (0:-1)
+    @test searchsorted(soa, 2) == firstindex(soa) .+ (0:1) == typemin(Int) .+ (1:2)
+    @test searchsorted(soa, 3) == lastindex(soa) .+ (0:0) == typemin(Int) .+ (3:3)
+    @test searchsorted(soa, 4) === lastindex(soa) .+ (1:0)
 end
 
 @testset "Adapt" begin