fix #34157, searchsorted corner case (e.g. searching for Inf) (#34224)

JuliaLang · Apr 11, 2020 · 6a63e29 · 6a63e29
1 parent 12e83af
commit 6a63e29
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 4 deletions.
diff --git a/base/sort.jl b/base/sort.jl
@@ -250,19 +250,37 @@ end
 
 function searchsortedlast(a::AbstractRange{<:Integer}, x::Real, o::DirectOrdering)
     require_one_based_indexing(a)
-    if step(a) == 0
+    h = step(a)
+    if h == 0
         lt(o, x, first(a)) ? 0 : length(a)
+    elseif h > 0 && x < first(a)
+        firstindex(a) - 1
+    elseif h > 0 && x >= last(a)
+        lastindex(a)
+    elseif h < 0 && x > first(a)
+        firstindex(a) - 1
+    elseif h < 0 && x <= last(a)
+        lastindex(a)
     else
-        clamp( fld(floor(Integer, x) - first(a), step(a)) + 1, 0, length(a))
+        fld(floor(Integer, x) - first(a), h) + 1
     end
 end
 
 function searchsortedfirst(a::AbstractRange{<:Integer}, x::Real, o::DirectOrdering)
     require_one_based_indexing(a)
-    if step(a) == 0
+    h = step(a)
+    if h == 0
         lt(o, first(a), x) ? length(a)+1 : 1
+    elseif h > 0 && x <= first(a)
+        firstindex(a)
+    elseif h > 0 && x > last(a)
+        lastindex(a) + 1
+    elseif h < 0 && x >= first(a)
+        firstindex(a)
+    elseif h < 0 && x < last(a)
+        lastindex(a) + 1
     else
-        clamp(-fld(floor(Integer, -x) + first(a), step(a)) + 1, 1, length(a) + 1)
+        -fld(floor(Integer, -x) + first(a), h) + 1
     end
 end
 

diff --git a/test/sorting.jl b/test/sorting.jl
@@ -138,6 +138,42 @@ end
         @test searchsortedlast(500:1.0:600, -1.0e20) == 0
         @test searchsortedlast(500:1.0:600, 1.0e20) == 101
     end
+    @testset "issue #34157" begin
+        @test searchsorted(1:2.0, -Inf) === 1:0
+        @test searchsorted([1,2], -Inf) === 1:0
+        @test searchsorted(1:2,   -Inf) === 1:0
+
+        @test searchsorted(1:2.0, Inf) === 3:2
+        @test searchsorted([1,2], Inf) === 3:2
+        @test searchsorted(1:2,   Inf) === 3:2
+
+        for coll in [
+                Base.OneTo(10),
+                1:2,
+                -4:6,
+                5:2:10,
+                [1,2],
+                1.0:4,
+                [10.0,20.0],
+            ]
+            for huge in [Inf, 1e300]
+                @test searchsortedfirst(coll, huge) === lastindex(coll) + 1
+                @test searchsortedfirst(coll, -huge)=== firstindex(coll)
+                @test searchsortedlast(coll, huge)  === lastindex(coll)
+                @test searchsortedlast(coll, -huge) === firstindex(coll) - 1
+                @test searchsorted(coll, huge)      === lastindex(coll)+1 : lastindex(coll)
+                @test searchsorted(coll, -huge)     === firstindex(coll) : firstindex(coll) - 1
+
+                @test searchsortedfirst(reverse(coll), huge, rev=true) === firstindex(coll)
+                @test searchsortedfirst(reverse(coll), -huge, rev=true) === lastindex(coll) + 1
+                @test searchsortedlast(reverse(coll), huge, rev=true) === firstindex(coll) - 1
+                @test searchsortedlast(reverse(coll), -huge, rev=true) === lastindex(coll)
+                @test searchsorted(reverse(coll), huge, rev=true) === firstindex(coll):firstindex(coll) - 1
+                @test searchsorted(reverse(coll), -huge, rev=true) === lastindex(coll)+1:lastindex(coll)
+
+            end
+        end
+    end
 end
 # exercise the codepath in searchsorted* methods for ranges that check for zero step range
 struct ConstantRange{T} <: AbstractRange{T}