Specialize axes for rangecumsum (#191)

Author: jishnub

Project.toml

@@ -1,7 +1,7 @@
 name = "ArrayLayouts"
 uuid = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
 authors = ["Sheehan Olver <[email protected]>"]
-version = "1.4.4"
+version = "1.4.5"
 
 [deps]
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
@@ -17,6 +17,7 @@ ArrayLayoutsSparseArraysExt = "SparseArrays"
 [compat]
 Aqua = "0.8"
 FillArrays = "1.2.1"
+Infinities = "0.1"
 LinearAlgebra = "1.6"
 Quaternions = "0.7"
 Random = "1.6"
@@ -27,11 +28,12 @@ julia = "1.6"
 
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+Infinities = "e1ba4f0e-776d-440f-acd9-e1d2e9742647"
 Quaternions = "94ee1d12-ae83-5a48-8b1c-48b8ff168ae0"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "Random", "StableRNGs", "SparseArrays", "Test", "Quaternions"]
+test = ["Aqua", "Infinities", "Quaternions", "Random", "StableRNGs", "SparseArrays", "Test"]

src/cumsum.jl

@@ -8,6 +8,7 @@ struct RangeCumsum{T, RR<:AbstractRange{T}} <: LayoutVector{T}
 end
 
 size(c::RangeCumsum) = size(c.range)
+axes(c::RangeCumsum) = axes(c.range)
 
 ==(a::RangeCumsum, b::RangeCumsum) = a.range == b.range
 BroadcastStyle(::Type{<:RangeCumsum{<:Any,RR}}) where RR = BroadcastStyle(RR)

test/infinitearrays.jl

@@ -0,0 +1,41 @@
+# Infinite Arrays implementation from
+# https://github.com/JuliaLang/julia/blob/master/test/testhelpers/InfiniteArrays.jl
+module InfiniteArrays
+    using Infinities
+    export OneToInf
+
+    abstract type AbstractInfUnitRange{T<:Real} <: AbstractUnitRange{T} end
+    Base.length(r::AbstractInfUnitRange) = ℵ₀
+    Base.size(r::AbstractInfUnitRange) = (ℵ₀,)
+    Base.last(r::AbstractInfUnitRange) = ℵ₀
+    Base.axes(r::AbstractInfUnitRange) = (OneToInf(),)
+
+    Base.IteratorSize(::Type{<:AbstractInfUnitRange}) = Base.IsInfinite()
+
+    """
+        OneToInf(n)
+    Define an `AbstractInfUnitRange` that behaves like `1:∞`, with the added
+    distinction that the limits are guaranteed (by the type system) to
+    be 1 and ∞.
+    """
+    struct OneToInf{T<:Integer} <: AbstractInfUnitRange{T} end
+
+    OneToInf() = OneToInf{Int}()
+
+    Base.axes(r::OneToInf) = (r,)
+    Base.first(r::OneToInf{T}) where {T} = oneunit(T)
+    Base.oneto(::InfiniteCardinal{0}) = OneToInf()
+
+    struct InfUnitRange{T<:Real} <: AbstractInfUnitRange{T}
+        start::T
+    end
+    Base.first(r::InfUnitRange) = r.start
+    InfUnitRange(a::InfUnitRange) = a
+    InfUnitRange{T}(a::AbstractInfUnitRange) where T<:Real = InfUnitRange{T}(first(a))
+    InfUnitRange(a::AbstractInfUnitRange{T}) where T<:Real = InfUnitRange{T}(first(a))
+    Base.:(:)(start::T, stop::InfiniteCardinal{0}) where {T<:Integer} = InfUnitRange{T}(start)
+    function getindex(v::InfUnitRange{T}, i::Integer) where T
+        @boundscheck i > 0 || Base.throw_boundserror(v, i)
+        convert(T, first(v) + i - 1)
+    end
+end

test/test_cumsum.jl

@@ -1,5 +1,7 @@
 using ArrayLayouts, Test
 
+include("infinitearrays.jl")
+
 @testset "RangeCumsum" begin
     for r in (RangeCumsum(Base.OneTo(5)), RangeCumsum(2:5), RangeCumsum(2:2:6), RangeCumsum(6:-2:1))
         @test r == cumsum(r.range)
@@ -20,4 +22,7 @@ using ArrayLayouts, Test
     a = RangeCumsum(Base.OneTo(3))
     b = RangeCumsum(1:3)
     @test oftype(a, b) === a
+
+    r = RangeCumsum(InfiniteArrays.OneToInf())
+    @test axes(r, 1) == InfiniteArrays.OneToInf()
 end