Sum for RangeCumsum (#216)

* Sum for RangeCumsum

* Don't coerce types in _half

Author: jishnub

src/cumsum.jl

@@ -15,7 +15,14 @@ Base.parent(r::RangeCumsum) = r.range
 ==(a::RangeCumsum, b::RangeCumsum) = a.range == b.range
 BroadcastStyle(::Type{<:RangeCumsum{<:Any,RR}}) where RR = BroadcastStyle(RR)
 
-_getindex(r::AbstractUnitRange{<:Integer}, k) = k * (2first(r) + k - 1) ÷ 2
+_half(x::Integer) = x ÷ 2
+_half(x) = x / 2
+
+function _getindex(r::AbstractRange{<:Real}, k)
+    v = first(r)
+    s = step(r)
+    _half(k * (2v - s + s*k))
+end
 Base.@propagate_inbounds _getindex(r::AbstractRange, k) = sum(r[range(firstindex(r), length=k)])
 
 Base.@propagate_inbounds function getindex(c::RangeCumsum{<:Any,<:AbstractRange}, k::Integer)
@@ -33,6 +40,13 @@ last(r::RangeCumsum) = sum(r.range)
 diff(r::RangeCumsum) = r.range[firstindex(r)+1:end]
 isempty(r::RangeCumsum) = isempty(r.range)
 
+function Base.sum(r::RangeCumsum{<:Real})
+    N = length(r)
+    v = first(r)
+    s = step(r.range)
+    _half((2v-s)*(N*(N+1)÷2) + s*(N*(N+1)*(2N+1)÷6))
+end
+
 union(a::RangeCumsum{<:Any,<:OneTo}, b::RangeCumsum{<:Any,<:OneTo}) =
     RangeCumsum(OneTo(max(last(a.range), last(b.range))))
 

test/test_cumsum.jl

@@ -4,20 +4,31 @@ using ArrayLayouts, Test
 
 include("infinitearrays.jl")
 
+cmpop(p) = isinteger(real(first(p))) && isinteger(real(step(p))) ? (==) : (≈)
+
 @testset "RangeCumsum" begin
-    @testset for p in (Base.OneTo(5), 2:5, 2:2:6, 6:-2:1, -1.0:3.0:5.0, (-1.0:3.0:5.0)*im,
-                        Base.IdentityUnitRange(4:6))
+    @testset for p in Any[Base.OneTo(5), 2:5, 2:2:6, 6:-2:1, Int8(2):Int8(5),
+                        UnitRange(2.5, 8.5),
+                        -1.0:1.0:10.0, -1.2:1.5:10.0,
+                        (2:5)*im, (-1:3:5)*im, (-1.0:3.0:5.0)*im, (-1.2:3.0:5.2)*(1+im),
+                        Base.IdentityUnitRange(4:6)]
+
         r = RangeCumsum(p)
         @test parent(r) == p
         @test r == r
+        cmp = cmpop(p)
+        if eltype(r) <: Complex
+            @test sum(r) isa Complex{promote_type(Int, real(eltype(r)))}
+        end
+        @test cmp(sum(r), sum(i for i in r))
         if axes(r,1) isa Base.OneTo
-            @test r == cumsum(p)
-            @test r .+ 1 == cumsum(p) .+ 1
+            @test cmp(r, cumsum(p))
+            @test cmp(r .+ 1, cumsum(p) .+ 1)
             @test r[Base.OneTo(3)] == r[1:3]
             @test @view(r[Base.OneTo(3)]) === r[Base.OneTo(3)] == r[1:3]
             @test @view(r[Base.OneTo(3)]) isa RangeCumsum
-            @test diff(r) == diff(Vector(r))
-            @test -r == -Vector(r)
+            @test cmp(diff(r),diff(Vector(r)))
+            @test cmp(-r, -Vector(r))
         end
         @test diff(r) == p[firstindex(p)+1:end]
         @test last(r) == r[end] == sum(p)
@@ -48,7 +59,7 @@ include("infinitearrays.jl")
             @test r * n isa RangeCumsum
             @test r * n ≈ w * n
         end
-        for p in (Base.OneTo(4), -4:4, -4:2:4, -1.0:3.0:5.0)
+        @testset for p in (Base.OneTo(4), -4:4, -4:2:4, -1.0:3.0:5.0)
             r = RangeCumsum(p)
             test_broadcast(3, r)
             test_broadcast(3.5, r)