Use lazy Cumsum for InfRanges (#35)

* Use lazy Cumsum for InfRanges

* v0.7.2

Author: dlfivefifty

Project.toml

@@ -1,6 +1,6 @@
 name = "InfiniteArrays"
 uuid = "4858937d-0d70-526a-a4dd-2d5cb5dd786c"
-version = "0.7.1"
+version = "0.7.2"
 
 [deps]
 DSP = "717857b8-e6f2-59f4-9121-6e50c889abd2"
@@ -13,12 +13,12 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 [compat]
 DSP = "0.6"
 FillArrays = "0.8.10"
-LazyArrays = "0.16.11"
+LazyArrays = "0.16.12"
 julia = "1.2"
 
 [extras]
-Base64 = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
+Base64 = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
 LazyBandedMatrices = "d7e5e226-e90b-4449-9968-0f923699bf6f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 

src/infrange.jl

@@ -28,7 +28,7 @@ _range(a::AbstractFloat, ::Nothing,         ::Nothing, len::Infinity) = _range(a
 _rangestyle(::Ordered, ::ArithmeticWraps, a::T, step::S, len::Infinity) where {T,S} =
     InfStepRange{T,S}(a, step)
 _rangestyle(::Ordered, ::ArithmeticUnknown, a::T, step::S, len::Infinity) where {T,S} =
-    InfStepRange{T,S}(a, step)    
+    InfStepRange{T,S}(a, step)
 _range(a::T, st::T, ::Nothing, ::Infinity) where T<:Union{Float16,Float32,Float64} =
     InfStepRange{T,T}(a, st)
 _range(a::T, st::T, ::Nothing, ::Infinity) where T<:AbstractFloat =
@@ -395,7 +395,19 @@ broadcast(f, r::Adjoint{<:Any,<:InfRanges}, a::Number) = broadcast(f,parent(r),a
 broadcast(f, r::Transpose{<:Any,<:InfRanges}, a::Number) = transpose(broadcast(f,parent(r),a))
 
 
-cumsum(r::InfRanges) = OneToInf() .* (first(r) .+ r) .÷ 2
+# cumsum(r::InfRanges) = OneToInf() .* (first(r) .+ r) .÷ 2
+
+cumsum(r::InfRanges) = Cumsum(r)
+function getindex(c::Cumsum{<:Integer,1,<:InfStepRange}, k::Integer)
+    r = c.v
+    k * (first(r) + r[k]) ÷ 2
+end
+function getindex(c::Cumsum{<:Integer,1,<:AbstractInfUnitRange}, k::Integer)
+    r = c.v
+    k * (2first(r) + k - 1) ÷ 2
+end
+getindex(c::Cumsum{<:Integer,1,<:OneToInf}, k::Integer) = k * (k+1) ÷ 2
+
 
 diff(r::InfRanges) = Fill(step(r),∞)
 diff(r::OneToInf{T}) where T = Ones{T}(∞)
@@ -418,22 +430,22 @@ conv(::Ones{Bool,1,<:Tuple{<:OneToInf}}, ::Ones{Bool,1,<:Tuple{<:OneToInf}}) =
 conv(::Ones{T,1,<:Tuple{<:OneToInf}}, ::Ones{V,1,<:Tuple{<:OneToInf}}) where {T,V} =
     one(promote_type(T,V)):∞
 
-function conv(::Ones{T,1,<:Tuple{<:OneToInf}}, a::AbstractVector{V}) where {T,V} 
+function conv(::Ones{T,1,<:Tuple{<:OneToInf}}, a::AbstractVector{V}) where {T,V}
     cs = cumsum(convert(AbstractVector{promote_type(T,V)}, a))
     Vcat(cs, Fill(last(cs), ∞))
 end
 
-function conv(::Ones{T,1,<:Tuple{<:OneToInf}}, a::Vector{V}) where {T,V} 
+function conv(::Ones{T,1,<:Tuple{<:OneToInf}}, a::Vector{V}) where {T,V}
     cs = cumsum(convert(AbstractVector{promote_type(T,V)}, a))
     Vcat(cs, Fill(last(cs), ∞))
 end
 
-function conv(a::AbstractVector{V}, ::Ones{T,1,<:Tuple{<:OneToInf}}) where {T,V} 
+function conv(a::AbstractVector{V}, ::Ones{T,1,<:Tuple{<:OneToInf}}) where {T,V}
     cs = cumsum(convert(AbstractVector{promote_type(T,V)}, a))
     Vcat(cs, Fill(last(cs), ∞))
 end
 
-function conv(a::Vector{V}, ::Ones{T,1,<:Tuple{<:OneToInf}}) where {T,V} 
+function conv(a::Vector{V}, ::Ones{T,1,<:Tuple{<:OneToInf}}) where {T,V}
     cs = cumsum(convert(AbstractVector{promote_type(T,V)}, a))
     Vcat(cs, Fill(last(cs), ∞))
 end
@@ -493,16 +505,16 @@ function _step_findfirst(p, r::InfStepRange{T,S}) where {T,S}
 end
 
 for op in (:isequal, :(==))
-    @eval begin 
+    @eval begin
         findfirst(p::Fix2{typeof($op),T}, r::InfStepRange{T,S}) where {T,S} =
             _step_findfirst(p, r)
 
-        findfirst(p::Fix2{typeof($op),T}, r::InfStepRange{T,<:Integer}) where {T<:Integer} = 
+        findfirst(p::Fix2{typeof($op),T}, r::InfStepRange{T,<:Integer}) where {T<:Integer} =
+            _step_findfirst(p, r)
+
+        findfirst(p::Fix2{typeof($op),<:Integer}, r::InfStepRange{<:Integer,<:Integer}) =
             _step_findfirst(p, r)
 
-        findfirst(p::Fix2{typeof($op),<:Integer}, r::InfStepRange{<:Integer,<:Integer}) = 
-            _step_findfirst(p, r)            
-        
         function findfirst(p::Fix2{typeof($op),<:Integer}, r::AbstractInfUnitRange{<:Integer})
             first(r) <= p.x || return nothing
             p.x - first(r) + 1
@@ -511,6 +523,6 @@ for op in (:isequal, :(==))
         findfirst(p::Fix2{typeof($op),<:Number}, r::AbstractInfUnitRange{<:Integer}) =
             isinteger(p.x) ? findfirst($op(convert(Integer, p.x)), r) : nothing
         findfirst(p::Fix2{typeof($op),<:Number}, r::InfStepRange{<:Integer,<:Integer}) =
-            isinteger(p.x) ? findfirst($op(convert(V, p.x)), r) : nothing            
+            isinteger(p.x) ? findfirst($op(convert(V, p.x)), r) : nothing
     end
 end

test/runtests.jl

@@ -1,5 +1,5 @@
 using LinearAlgebra, SparseArrays, InfiniteArrays, FillArrays, LazyArrays, Statistics, DSP, BandedMatrices, LazyBandedMatrices, Test, Base64
-import InfiniteArrays: OrientedInfinity, SignedInfinity, OneToInf, InfUnitRange, InfStepRange, OneToInf
+import InfiniteArrays: OrientedInfinity, SignedInfinity, InfUnitRange, InfStepRange, OneToInf
 import LazyArrays: CachedArray, MemoryLayout, LazyLayout, DiagonalLayout, LazyArrayStyle, colsupport
 import BandedMatrices: _BandedMatrix, BandedColumns
 import Base.Broadcast: broadcasted, Broadcasted, instantiate
@@ -663,8 +663,13 @@ end
 
     @test cumsum(x).args[2] ≡ 8:12
     @test last(y.args) == sum(x[1:9]):2:∞
-    r = (3:4:∞)
-    @test cumsum(r)[1:20] == cumsum(r[1:20])
+
+    for r in (3:4:∞, 2:∞, Base.OneTo(∞))
+        c = cumsum(r)
+        @test c isa Cumsum
+        @test c[1:20] == [c[k] for k=1:20] == cumsum(r[1:20])
+        @test c == c
+    end
 end
 
 @testset "Sub-array" begin