Specialize axes for Accumulate

Project.toml

@@ -1,6 +1,6 @@
 name = "LazyArrays"
 uuid = "5078a376-72f3-5289-bfd5-ec5146d43c02"
-version = "1.8.2"
+version = "1.8.3"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
@@ -23,6 +23,7 @@ ArrayLayouts = "1.4.1"
 Base64 = "1"
 FillArrays = "1.0"
 LinearAlgebra = "1"
+Infinities = "0.1"
 MacroTools = "0.5"
 MatrixFactorizations = "1.0, 2.0"
 SparseArrays = "1"
@@ -34,9 +35,10 @@ julia = "1.9"
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Base64 = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
+Infinities = "e1ba4f0e-776d-440f-acd9-e1d2e9742647"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 
 [targets]
-test = ["Aqua", "Base64", "StaticArrays", "Tracker", "Test"]
+test = ["Aqua", "Base64", "StaticArrays", "Tracker", "Test", "Infinities"]

src/lazyoperations.jl

@@ -428,6 +428,7 @@
 IndexStyle(::Type{<:Accumulate{<:Any,1}}) = IndexLinear()
 
 size(Q::Accumulate) = size(Q.v)
+axes(Q::Accumulate) = axes(Q.v)
 
 copy(Q::Accumulate) = Accumulate(Q.op, copy(Q.data), copy(Q.v), Q.dims, Q.datasize)
 

test/cachetests.jl

@@ -2,6 +2,10 @@ using LazyArrays, FillArrays, LinearAlgebra, ArrayLayouts, StaticArrays, SparseA
 import LazyArrays: CachedArray, CachedMatrix, CachedVector, PaddedLayout, CachedLayout, resizedata!, zero!,
                     CachedAbstractArray, CachedAbstractVector, CachedAbstractMatrix, AbstractCachedArray, AbstractCachedMatrix
 
+include("infinitearrays.jl")
+using .InfiniteArrays
+using Infinities
+
 @testset "Cache" begin
     @testset "basics" begin
         A = 1:10
@@ -172,6 +176,10 @@ import LazyArrays: CachedArray, CachedMatrix, CachedVector, PaddedLayout, Cached
         c = cache(Fill(1,10))
         @test c[2:3] isa LazyArrays.CachedVector{Int,Vector{Int},Fill{Int,1,Tuple{Base.OneTo{Int}}}}
         @test c[[2,4,6]] isa LazyArrays.CachedVector{Int,Vector{Int},Fill{Int,1,Tuple{Base.OneTo{Int}}}}
+
+        F = Fill(2, ℵ₀)
+        C = cumsum(cache(F))
+        @test axes(C) == (InfiniteArrays.OneToInf(),)
     end
 
     @testset "linalg" begin

test/infinitearrays.jl

@@ -0,0 +1,43 @@
+# Infinite Arrays implementation from
+# https://github.com/JuliaLang/julia/blob/master/test/testhelpers/InfiniteArrays.jl
+module InfiniteArrays
+    using Infinities
+
+    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()
+
+    Base.axes(::AbstractInfUnitRange) = (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 Base.getindex(v::InfUnitRange{T}, i::Integer) where T
+        checkbounds(v, i)
+        convert(T, first(v) + i - 1)
+    end
+end