Faster iteration for an IdOffsetRange (#215)

* specialize for Base.OneTo

* rebase on master

* use helper functions instead of broadcasting

Author: jishnub

src/axes.jl

@@ -167,6 +167,7 @@ offset_coerce(::Type{I}, r::AbstractUnitRange) where I<:AbstractUnitRange =
 @inline Base.axes1(r::IdOffsetRange) = IdOffsetRange(Base.axes1(r.parent), r.offset)
 @inline Base.unsafe_indices(r::IdOffsetRange) = (Base.axes1(r),)
 @inline Base.length(r::IdOffsetRange) = length(r.parent)
+@inline Base.isempty(r::IdOffsetRange) = isempty(r.parent)
 Base.reduced_index(i::IdOffsetRange) = typeof(i)(first(i):first(i))
 # Workaround for #92 on Julia < 1.4
 Base.reduced_index(i::IdentityUnitRange{<:IdOffsetRange}) = typeof(i)(first(i):first(i))
@@ -178,16 +179,23 @@ for f in [:first, :last]
     @eval @inline Base.$f(r::IdOffsetRange) = eltype(r)($f(r.parent) + r.offset)
 end
 
-@inline function Base.iterate(r::IdOffsetRange)
-    ret = iterate(r.parent)
-    ret === nothing && return nothing
-    return (eltype(r)(ret[1] + r.offset), ret[2])
-end
-@inline function Base.iterate(r::IdOffsetRange, i)
-    ret = iterate(r.parent, i)
+# Iteration for an IdOffsetRange
+@inline Base.iterate(r::IdOffsetRange, i...) = _iterate(r, i...)
+# In general we iterate over the parent term by term and add the offset.
+# This might have some performance degradation when coupled with bounds-checking
+# See https://github.com/JuliaArrays/OffsetArrays.jl/issues/214
+@inline function _iterate(r::IdOffsetRange, i...)
+    ret = iterate(r.parent, i...)
     ret === nothing && return nothing
     return (eltype(r)(ret[1] + r.offset), ret[2])
 end
+# Base.OneTo(n) is known to be exactly equivalent to the range 1:n,
+# and has no specialized iteration defined for it,
+# so we may add the offset to the range directly and iterate over the result
+# This gets around the performance issue described in issue #214
+# We use the helper function _addoffset to evaluate the range instead of broadcasting
+# just in case this makes it easy for the compiler.
+@inline _iterate(r::IdOffsetRange{<:Integer, <:Base.OneTo}, i...) = iterate(_addoffset(r.parent, r.offset), i...)
 
 @inline function Base.getindex(r::IdOffsetRange, i::Integer)
     i isa Bool && throw(ArgumentError("invalid index: $i of type Bool"))

src/utils.jl

@@ -92,8 +92,12 @@ end
 # These functions are equivalent to the broadcasted operation r .- of
 # However these ensure that the result is an AbstractRange even if a specific
 # broadcasting behavior is not defined for a custom type
-_subtractoffset(r::AbstractUnitRange, of) = UnitRange(first(r) - of, last(r) - of)
-_subtractoffset(r::AbstractRange, of) = range(first(r) - of, stop = last(r) - of, step = step(r))
+@inline _subtractoffset(r::AbstractUnitRange, of) = UnitRange(first(r) - of, last(r) - of)
+@inline _subtractoffset(r::AbstractRange, of) = range(first(r) - of, stop = last(r) - of, step = step(r))
+
+# similar to _subtractoffset, except these evaluate r .+ of
+@inline _addoffset(r::AbstractUnitRange, of) = UnitRange(first(r) + of, last(r) + of)
+@inline _addoffset(r::AbstractRange, of) = range(first(r) + of, stop = last(r) + of, step = step(r))
 
 if VERSION <= v"1.7.0-DEV.1039"
     _contiguousindexingtype(r::AbstractUnitRange{<:Integer}) = UnitRange{Int}(r)

test/runtests.jl

@@ -345,6 +345,25 @@ end
             @test_throws BoundsError r[false:true:false]
         end
     end
+
+    @testset "iteration" begin
+        # parent has Base.OneTo axes
+        A = ones(4:10)
+        ax = axes(A, 1)
+        ind, st = iterate(ax)
+        @test A[ind] == A[4]
+        ind, st = iterate(ax, st)
+        @test A[ind] == A[5]
+
+        # parent doesn't have Base.OneTo axes
+        B = @view A[:]
+        C = OffsetArray(B, 0)
+        ax = axes(C, 1)
+        ind, st = iterate(ax)
+        @test C[ind] == C[4]
+        ind, st = iterate(ax, st)
+        @test C[ind] == C[5]
+    end
 end
 
 # used in testing the constructor
@@ -2472,3 +2491,15 @@ end
 end
 
 include("origin.jl")
+
+@testset "misc" begin
+    @test OffsetArrays._subtractoffset(Base.OneTo(2), 1) isa AbstractUnitRange{Int}
+    @test OffsetArrays._subtractoffset(Base.OneTo(2), 1) == 0:1
+    @test OffsetArrays._subtractoffset(3:2:9, 1) isa AbstractRange{Int}
+    @test OffsetArrays._subtractoffset(3:2:9, 1) == 2:2:8
+
+    @test OffsetArrays._addoffset(Base.OneTo(2), 1) isa AbstractUnitRange{Int}
+    @test OffsetArrays._addoffset(Base.OneTo(2), 1) == 2:3
+    @test OffsetArrays._addoffset(3:2:9, 1) isa AbstractRange{Int}
+    @test OffsetArrays._addoffset(3:2:9, 1) == 4:2:10
+end