Fix indexing for non-Real and offset ranges (#212)

* Fix indexing for non-Real and offset ranges

* Extra test for diff with offset ranges

Author: jishnub

src/cumsum.jl

@@ -13,16 +13,13 @@ axes(c::RangeCumsum) = axes(c.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
+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)
     @boundscheck checkbounds(c, k)
     r = c.range
-    k * (first(r) + r[k]) ÷ 2
-end
-Base.@propagate_inbounds function getindex(c::RangeCumsum{<:Any,<:AbstractUnitRange}, k::Integer)
-    @boundscheck checkbounds(c, k)
-    r = c.range
-    k * (2first(r) + k - 1) ÷ 2
+    _getindex(r, k-firstindex(r)+1)
 end
 
 Base.@propagate_inbounds getindex(c::RangeCumsum, kr::OneTo) = RangeCumsum(c.range[kr])
@@ -31,7 +28,7 @@ Base.@propagate_inbounds view(c::RangeCumsum, kr::OneTo) = c[kr]
 
 first(r::RangeCumsum) = first(r.range)
 last(r::RangeCumsum) = sum(r.range)
-diff(r::RangeCumsum) = r.range[2:end]
+diff(r::RangeCumsum) = r.range[firstindex(r)+1:end]
 isempty(r::RangeCumsum) = isempty(r.range)
 
 union(a::RangeCumsum{<:Any,<:OneTo}, b::RangeCumsum{<:Any,<:OneTo}) =

test/test_cumsum.jl

@@ -5,16 +5,21 @@ 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)
+    @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))
+        r = RangeCumsum(p)
         @test r == r
-        @test r .+ 1 == cumsum(r.range) .+ 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 last(r) == r[end]
-        @test diff(r) == diff(Vector(r))
-        @test first(r) == r[1]
+        if axes(r) isa Base.OneTo
+            @test r == cumsum(p)
+            @test 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))
+        end
+        @test diff(r) == p[firstindex(p)+1:end]
+        @test last(r) == r[end] == sum(p)
+        @test first(r) == r[firstindex(r)] == first(p)
     end
 
     a,b = RangeCumsum(Base.OneTo(5)), RangeCumsum(Base.OneTo(6))