Sort RangeCumsum for AbstractUnitRange (#224)

Author: jishnub

Project.toml

@@ -1,7 +1,7 @@
 name = "ArrayLayouts"
 uuid = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
 authors = ["Sheehan Olver <[email protected]>"]
-version = "1.9.0"
+version = "1.9.1"
 
 [deps]
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"

src/cumsum.jl

@@ -64,6 +64,16 @@ union(a::RangeCumsum{<:Any,<:OneTo}, b::RangeCumsum{<:Any,<:OneTo}) =
 
 sort!(a::RangeCumsum{<:Any,<:Base.OneTo}) = a
 sort(a::RangeCumsum{<:Any,<:Base.OneTo}) = a
+function sort(a::RangeCumsum{<:Any,<:AbstractUnitRange{<:Integer}})
+    r = parent(a)
+    #= A RangeCumsum with a range (-a:b) for positive (a,b) may be viewed as a concatenation
+    of two components: (-a:a) and (a+1:b). The second is strictly increasing.
+    We only sort the first component, and concatenate the second to the result.
+    =#
+    a1 = RangeCumsum(r[firstindex(r):searchsortedlast(r, -first(r))])
+    a2 = RangeCumsum(r[searchsortedfirst(r, -first(r)+1):end])
+    vcat(sort!(Vector(a1)), a2)
+end
 Base.issorted(a::RangeCumsum{<:Any,<:Base.OneTo}) = true
 function Base.issorted(a::RangeCumsum{<:Any,<:AbstractUnitRange{<:Integer}})
     r = parent(a)

test/test_cumsum.jl

@@ -44,12 +44,6 @@ cmpop(p) = isinteger(real(first(p))) && isinteger(real(step(p))) ? (==) : (≈)
 
     a,b = RangeCumsum(Base.OneTo(5)), RangeCumsum(Base.OneTo(6))
     @test union(a,b) ≡ union(b,a) ≡ b
-    @test sort!(copy(a)) == a
-    @test sort!(a) ≡ a
-    @test sort(a) ≡ a == Vector(a)
-
-    r = RangeCumsum(-4:4)
-    @test sort(r) == sort(Vector(r))
 
     a = RangeCumsum(Base.OneTo(3))
     b = RangeCumsum(1:3)
@@ -97,6 +91,20 @@ cmpop(p) = isinteger(real(first(p))) && isinteger(real(step(p))) ? (==) : (≈)
         @test issorted(a)
     end
 
+    @testset "sort" begin
+        @testset "RangeCumsum($start:$stop)" for start in -15:15, stop in start-1:20
+            a = RangeCumsum(start:stop)
+            v = Vector(a)
+            @test sort(a) == sort(v)
+        end
+
+        a,b = RangeCumsum(Base.OneTo(5)), RangeCumsum(Base.OneTo(6))
+        @test union(a,b) ≡ union(b,a) ≡ b
+        @test sort!(copy(a)) == a
+        @test sort!(a) ≡ a
+        @test sort(a) ≡ a == Vector(a)
+    end
+
     @testset "minimum/maximum" begin
         @testset "RangeCumsum($start:$stop)" for start in -15:15, stop in start:20
             a = RangeCumsum(start:stop)