[RFC] Dispatch on traits instead of types in iterate for AbstractArrays

[jishnub]

This is related to the issue noted in JuliaArrays/OffsetArrays.jl#257, specifically OrdinalRanges implement custom iterate methods, but AbstactArray wrappers around such ranges use the fallback iterate method for AbstractArrays. Evidently this has a significant impact on performance:

On julia 1.6.2

julia> using OffsetArrays, BenchmarkTools

julia> r = 1:1:1000; ro = OffsetArray(r); s = collect(r); so = OffsetArray(s);

julia> @btime Float64[i for i in $r];
  1.621 μs (1 allocation: 7.94 KiB)

julia> @btime Float64[i for i in $ro];
  25.306 μs (1 allocation: 7.94 KiB)

julia> @which iterate(r)
iterate(r::OrdinalRange) in Base at range.jl:670

julia> @which iterate(ro)
iterate(A::AbstractArray) in Base at abstractarray.jl:1093

julia> @btime Float64[i for i in $s];
  1.738 μs (1 allocation: 7.94 KiB)

julia> @btime Float64[i for i in $so];
  1.991 μs (1 allocation: 7.94 KiB)

Ideally iteration for the OffsetRange ro should be as fast as that for the parent range r. This might be possible using traits, if iterate dispatches on a range-like trait instead of subtypes of AbstractRange.

Another solution to this would be for OffsetArrays to forward the iteration to the parent, however this was decided against in the PR above. In any case this is not a general solution, and addressing the issue in Base would be ideal.

[jishnub]

Here's a minimal change that gets around this issue:

In Base:

# range.jl
abstract type IterationStyle end
struct AbstractRangeIterationStyle <: IterationStyle end
struct OrdinalRangeIterationStyle <: IterationStyle end

IterationStyle(A) = IterationStyle(typeof(A))
IterationStyle(::Type{<:AbstractRange}) = AbstractRangeIterationStyle()
IterationStyle(::Type{<:OrdinalRange}) = OrdinalRangeIterationStyle()

function _iterate(r, ::AbstractRangeIterationStyle, i::Integer = zero(length(r)))
    @_inline_meta
    i += oneunit(i)
    length(r) < i && return nothing
    unsafe_getindex(r, i), i
end

iterate(A::OrdinalRange) = _iterate(A, IterationStyle(A))
iterate(A::OrdinalRange, i) = _iterate(A, IterationStyle(A), i)
_iterate(r, ::OrdinalRangeIterationStyle) = isempty(r) ? nothing : (x = first(r); (x, x))

function _iterate(r, ::OrdinalRangeIterationStyle, i)
    @_inline_meta
    i == last(r) && return nothing
    next = convert(eltype(r), i + step(r))
    (next, next)
end

# abstractarray.jl
struct AbstractArrayIterationStyle <: IterationStyle end
IterationStyle(::Type{<:AbstractArray}) = AbstractArrayIterationStyle()

iterate(A::AbstractArray, state...) = _iterate(A, IterationStyle(A), state...)

function _iterate(A::AbstractArray, ::AbstractArrayIterationStyle, state = (eachindex(A),))
    y = iterate(state...)
    y === nothing && return nothing
    A[y[1]], (state[1], tail(y)...)
end

# array.jl
struct ArrayIterationStyle <: IterationStyle end
IterationStyle(::Type{<:Array}) = ArrayIterationStyle()
iterate(A::Array, i=1) = _iterate(A, IterationStyle(A), i)
_iterate(A, ::ArrayIterationStyle, i=1) = (@_inline_meta; (i % UInt) - 1 < length(A) ? (@inbounds A[i - 1 + firstindex(A)], i + 1) : nothing)

In OffsetArrays.jl:

for f in [:first, :last, :step]
    @eval Base.$f(a::OffsetRange) = $f(parent(a))
end

Base.IterationStyle(A::OffsetArray) = Base.IterationStyle(parent(A))

With these, I obtain

julia> @btime Float64[i for i in $r];
  1.677 μs (1 allocation: 7.94 KiB)

julia> @btime Float64[i for i in $ro];
  1.609 μs (1 allocation: 7.94 KiB)

julia> @btime Float64[i for i in $s];
  1.738 μs (1 allocation: 7.94 KiB)

julia> @btime Float64[i for i in $so];
  1.764 μs (1 allocation: 7.94 KiB)

I've not looked at consequences of such a change, so comments are welcome. The trait names are somewhat arbitrary here, and these may be generalized to eg. AbstractRangeStyle from AbstractRangeIterationStyle if such traits are used elsewhere as well, eg. bounds-checking.

cc @timholy