optimize reductions for OffsetRanges (#202)

* optimize reductions for OffsetUnitRanges

* propagate sum to parent

Author: jishnub

src/OffsetArrays.jl

@@ -323,6 +323,8 @@ end
     A
 end
 
+Base.in(x, A::OffsetArray) = in(x, parent(A))
+
 Base.strides(A::OffsetArray) = strides(parent(A))
 Base.elsize(::Type{OffsetArray{T,N,A}}) where {T,N,A} = Base.elsize(A)
 @inline Base.unsafe_convert(::Type{Ptr{T}}, A::OffsetArray{T}) where {T} = Base.unsafe_convert(Ptr{T}, parent(A))
@@ -334,6 +336,7 @@ Broadcast.broadcast_unalias(dest::OffsetArray, src::OffsetArray) = parent(dest)
 ### Special handling for AbstractRange
 
 const OffsetRange{T} = OffsetArray{T,1,<:AbstractRange{T}}
+const OffsetUnitRange{T} = OffsetArray{T,1,<:AbstractUnitRange{T}}
 const IIUR = IdentityUnitRange{S} where S<:AbstractUnitRange{T} where T<:Integer
 
 Base.step(a::OffsetRange) = step(parent(a))
@@ -368,6 +371,19 @@ end
 # This is technically breaking, so it might be incorporated in the next major release
 # Base.getindex(a::OffsetRange, ::Colon) = OffsetArray(a.parent[:], a.offsets)
 
+# mapreduce is faster with an IdOffsetRange than with an OffsetUnitRange
+# We therefore convert OffsetUnitRanges to IdOffsetRanges with the same values and axes
+function Base.mapreduce(f, op, As::OffsetUnitRange...; kw...)
+    ofs = map(A -> first(axes(A,1)) - 1, As)
+    AIds = map((A, of) -> IdOffsetRange(UnitRange(parent(A) .- of), of), As, ofs)
+    mapreduce(f, op, AIds...; kw...)
+end
+
+# Optimize certain reductions that treat an OffsetVector as a list
+for f in [:minimum, :maximum, :extrema, :sum]
+    @eval Base.$f(r::OffsetRange) = $f(parent(r))
+end
+
 function Base.show(io::IO, r::OffsetRange)
     show(io, r.parent)
     print(io, " with indices ", UnitRange(axes(r, 1)))

test/runtests.jl

@@ -374,6 +374,13 @@ Base.convert(::Type{Int}, a::WeirdInteger) = a
         @test_throws OverflowError OffsetArray{Float64, 1, typeof(ao)}(ao, (-2, )) # inner Constructor
         @test_throws OverflowError OffsetArray(ao, (-2, )) # convinient constructor accumulate offsets
 
+        @testset "OffsetRange" begin
+            local r = 1:100
+            local a = OffsetVector(r, 4)
+            @test first(r) in a
+            @test !(last(r) + 1 in a)
+        end
+
         # disallow OffsetVector(::Array{<:Any, N}, offsets) where N != 1
         @test_throws ArgumentError OffsetVector(zeros(2,2), (2, 2))
         @test_throws ArgumentError OffsetVector(zeros(2,2), 2, 2)
@@ -1512,6 +1519,44 @@ end
 
     amin, amax = extrema(parent(A))
     @test clamp.(A, (amax+amin)/2, amax) == OffsetArray(clamp.(parent(A), (amax+amin)/2, amax), axes(A))
+
+    @testset "mapreduce for OffsetRange" begin
+        for r in Any[5:100, IdOffsetRange(1:100, 4), IdOffsetRange(4:5), # AbstractUnitRanges
+            2:4:14, 1.5:1.0:10.5, # AbstractRanges
+            ]
+
+            a = OffsetVector(r, 2);
+            @test mapreduce(identity, +, a) == mapreduce(identity, +, r)
+            @test mapreduce(x -> x^2, (x,y) -> x, a) == mapreduce(x -> x^2, (x,y) -> x, r)
+
+            b = mapreduce(identity, +, a, dims = 1)
+            br = mapreduce(identity, +, r, dims = 1)
+            @test no_offset_view(b) == no_offset_view(br)
+            @test axes(b, 1) == first(axes(a,1)):first(axes(a,1))
+
+            @test mapreduce(identity, +, a, init = 3) == mapreduce(identity, +, r, init = 3)
+            if VERSION >= v"1.2"
+                @test mapreduce((x,y) -> x*y, +, a, a) == mapreduce((x,y) -> x*y, +, r, r)
+                @test mapreduce((x,y) -> x*y, +, a, a, init = 10) == mapreduce((x,y) -> x*y, +, r, r, init = 10)
+            end
+
+            for f in [sum, minimum, maximum]
+                @test f(a) == f(r)
+
+                b = f(a, dims = 1);
+                br = f(r, dims = 1)
+                @test no_offset_view(b) == no_offset_view(br)
+                @test axes(b, 1) == first(axes(a,1)):first(axes(a,1))
+
+                b = f(a, dims = 2);
+                br = f(r, dims = 2)
+                @test no_offset_view(b) == no_offset_view(br)
+                @test axes(b, 1) == axes(a,1)
+            end
+
+            @test extrema(a) == extrema(r)
+        end
+    end
 end
 
 # v  = OffsetArray([1,1e100,1,-1e100], (-3,))*1000