Preserve ranges in indexing with IdentityUnitRange(::Base.OneTo) (#211)

* Preserve ranges in indexing with IIUR(::OneTo)

* remove _maybewrapaxes

* return AbstractUnitRanges in indexing OffsetUnitRanges

* restrict _maybewrapoffset to integer ranges

* to_indices for OffsetUnitRange{<:Integer}

* nightly docstring hack

* checkindex for OffsetRange

* trailing newline

* Stronger test for indexing with UnitRange

* improve comments in tests

* inline to_indices

* getindex(::Array, ::OffsetUnitRange)

* Add tests

* merge getindex method definitions

* Test AbstractUnitRange constructor

* Add tests

* remove specific reductions for Int32 and Int64

* fix test on nightly

* fix conversion to AbstractUnitRange

* avoid broadcasting in indexing with custom ranges

* fix range on v1.0

Author: jishnub

benchmark/benchmarks.jl

@@ -3,10 +3,14 @@ using OffsetArrays
 
 const dim = 1000
 
-x = Array{Float64}(undef, 2*dim)
-y = OffsetArray{Float64}(undef, -dim + 1 : dim)
-x2d = Array{Float64}(undef, 2*dim, 2*dim)
-y2d = OffsetArray{Float64}(undef, -dim + 1 : dim, -dim + 1 : dim)
+x = Array{Float64}(undef, 2*dim);
+y = OffsetArray{Float64}(undef, -dim + 1 : dim);
+x2d = Array{Float64}(undef, 2*dim, 2*dim);
+y2d = OffsetArray{Float64}(undef, -dim + 1 : dim, -dim + 1 : dim);
+
+s = OffsetVector(1:dim, 0);
+sur = 1:dim;
+sior = OffsetArrays.IdOffsetRange(parent(s));
 
 fill1d(x) = for i in axes(x,1); x[i] = i; end
 fill2d(x) = for j in axes(x,2); for i in axes(x,1); x[i,j] = i + j; end; end
@@ -20,24 +24,63 @@ unsafe_update(x) = @inbounds(for i in axes(x,1); x[i] = x[i] + i; end)
 unsafe_update2d(x) = @inbounds(for j in axes(x,2); for i in axes(x,1); x[i,j] = x[i,j] + i + j; end; end)
 unsafe_update_eachindex(x) = @inbounds(for i in eachindex(x); x[i] = x[i] + i; end)
 
-@show @benchmark fill1d(x)
-@show @benchmark fill1d(y)
-@show @benchmark fill2d(x2d)
-@show @benchmark fill2d(y2d)
-@show @benchmark update(x)
-@show @benchmark update(y)
-@show @benchmark update2d(x2d)
-@show @benchmark update2d(y2d)
-@show @benchmark update_eachindex(x)
-@show @benchmark update_eachindex(y)
-
-@show @benchmark unsafe_fill(x)
-@show @benchmark unsafe_fill(y)
-@show @benchmark unsafe_fill2d(x2d)
-@show @benchmark unsafe_fill2d(y2d)
-@show @benchmark unsafe_update(x)
-@show @benchmark unsafe_update(y)
-@show @benchmark unsafe_update2d(x2d)
-@show @benchmark unsafe_update2d(y2d)
-@show @benchmark unsafe_update_eachindex(x)
-@show @benchmark unsafe_update_eachindex(y)
+vectorlinearindexing(a, ax) = a[ax]
+vectorCartesianindexing(a, ax1, ax2) = a[ax1, ax2]
+nestedvectorlinearindexing(a, ax1, ax2) = a[ax1[ax2]]
+
+macro showbenchmark(ex)
+	print(ex, " : ")
+	quote
+		println(@benchmark esc($ex))
+	end
+end
+
+@showbenchmark fill1d(x)
+@showbenchmark fill1d(y)
+@showbenchmark fill2d(x2d)
+@showbenchmark fill2d(y2d)
+@showbenchmark update(x)
+@showbenchmark update(y)
+@showbenchmark update2d(x2d)
+@showbenchmark update2d(y2d)
+@showbenchmark update_eachindex(x)
+@showbenchmark update_eachindex(y)
+
+@showbenchmark unsafe_fill(x)
+@showbenchmark unsafe_fill(y)
+@showbenchmark unsafe_fill2d(x2d)
+@showbenchmark unsafe_fill2d(y2d)
+@showbenchmark unsafe_update(x)
+@showbenchmark unsafe_update(y)
+@showbenchmark unsafe_update2d(x2d)
+@showbenchmark unsafe_update2d(y2d)
+@showbenchmark unsafe_update_eachindex(x)
+@showbenchmark unsafe_update_eachindex(y)
+
+# Benchmarks of vector indexing using OffsetRanges as axes
+@showbenchmark vectorlinearindexing(x, s)
+@showbenchmark vectorlinearindexing(x, sur)
+@showbenchmark vectorlinearindexing(x, sior)
+@showbenchmark vectorlinearindexing(y, s)
+@showbenchmark vectorlinearindexing(y, sur)
+@showbenchmark vectorlinearindexing(y, sior)
+
+@showbenchmark vectorlinearindexing(sur, s)
+@showbenchmark vectorlinearindexing(sur, sur)
+@showbenchmark vectorlinearindexing(sur, sior)
+
+@showbenchmark vectorCartesianindexing(x2d, s, s)
+@showbenchmark vectorCartesianindexing(x2d, sur, sur)
+@showbenchmark vectorCartesianindexing(x2d, sior, sior)
+
+@showbenchmark nestedvectorlinearindexing(x, s, s)
+@showbenchmark nestedvectorlinearindexing(x, sur, sur)
+@showbenchmark nestedvectorlinearindexing(x, s, sior)
+@showbenchmark nestedvectorlinearindexing(x, sur, sior)
+@showbenchmark nestedvectorlinearindexing(x, sior, sior)
+@showbenchmark vectorlinearindexing(x, sior[sior])
+@showbenchmark nestedvectorlinearindexing(x, sur, sior)
+@showbenchmark vectorlinearindexing(x, sur[sior])
+@showbenchmark nestedvectorlinearindexing(x, sior, sur)
+@showbenchmark vectorlinearindexing(x, sior[sur])
+@showbenchmark nestedvectorlinearindexing(x, sur, sur)

src/OffsetArrays.jl

@@ -392,44 +392,92 @@ 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 OffsetRange{T} = OffsetVector{T,<:AbstractRange{T}}
+const OffsetUnitRange{T} = OffsetVector{T,<:AbstractUnitRange{T}}
 const IIUR = IdentityUnitRange{S} where S<:AbstractUnitRange{T} where T<:Integer
 
 Base.step(a::OffsetRange) = step(parent(a))
 
-@propagate_inbounds function Base.getindex(a::OffsetRange, r::OffsetRange)
-    OffsetArray(a.parent[r.parent .- a.offsets[1]], axes(r))
+Base.checkindex(::Type{Bool}, inds::AbstractUnitRange, or::OffsetRange) = Base.checkindex(Bool, inds, parent(or))
+
+# Certain special methods for linear indexing with integer ranges (or OffsetRanges)
+# These may bypass the default getindex(A, I...) pathway if the parent types permit this
+# For example AbstractUnitRanges and Arrays have special linear indexing behavior defined
+
+# If both the arguments are offset, we may unwrap the indices to call (::OffsetArray)[::AbstractRange{Int}]
+@propagate_inbounds function Base.getindex(A::OffsetArray, r::OffsetRange{Int})
+    _maybewrapoffset(A[parent(r)], axes(r))
+end
+# If the indices are offset, we may unwrap them and pass the parent to getindex
+@propagate_inbounds function Base.getindex(A::AbstractRange, r::OffsetRange{Int})
+    _maybewrapoffset(A[parent(r)], axes(r))
+end
+
+# An OffsetUnitRange might use the rapid getindex(::Array, ::AbstractUnitRange{Int}) for contiguous indexing
+@propagate_inbounds function Base.getindex(A::Array, r::OffsetUnitRange{Int})
+    B = A[_contiguousindexingtype(parent(r))]
+    OffsetArray(B, axes(r))
+end
+
+# avoid hitting the slow method getindex(::Array, ::AbstractRange{Int})
+# instead use the faster getindex(::Array, ::UnitRange{Int})
+@propagate_inbounds function Base.getindex(A::Array, r::Union{IdOffsetRange, IIUR})
+    B = A[_contiguousindexingtype(r)]
+    _maybewrapoffset(B, axes(r))
 end
-@propagate_inbounds function Base.getindex(a::OffsetRange, r::IdOffsetRange)
-    OffsetArray(a.parent[r.parent .+ (r.offset - a.offsets[1])], axes(r))
+
+# Linear Indexing of OffsetArrays with AbstractUnitRanges may use the faster contiguous indexing methods
+@inline function Base.getindex(A::OffsetArray, r::AbstractUnitRange{Int})
+    @boundscheck checkbounds(A, r)
+    # nD OffsetArrays do not have their linear indices shifted, so we may forward the indices provided to the parent
+    @inbounds B = parent(A)[_contiguousindexingtype(r)]
+    _maybewrapoffset(B, axes(r))
+end
+@inline function Base.getindex(A::OffsetVector, r::AbstractUnitRange{Int})
+    @boundscheck checkbounds(A, r)
+    # OffsetVectors may have their linear indices shifted, so we subtract the offset from the indices provided
+    @inbounds B = parent(A)[_subtractoffset(r, A.offsets[1])]
+    _maybewrapoffset(B, axes(r))
+end
+
+# This method added mainly to index an OffsetRange with another range
+@inline function Base.getindex(A::OffsetVector, r::AbstractRange{Int})
+    @boundscheck checkbounds(A, r)
+    @inbounds B = parent(A)[_subtractoffset(r, A.offsets[1])]
+    _maybewrapoffset(B, axes(r))
+end
+
+# In general we would pass through getindex(A, I...) which calls to_indices(A, I) and finally to_index(I)
+# An OffsetUnitRange{Int} has an equivalent IdOffsetRange with the same values and axes,
+# something similar also holds for OffsetUnitRange{BigInt}
+# We may replace the former with the latter in an indexing operation to obtain a performance boost
+@inline function Base.to_index(r::OffsetUnitRange{<:Union{Int,BigInt}})
+    of = first(axes(r,1)) - 1
+    IdOffsetRange(_subtractoffset(parent(r), of), of)
 end
-@propagate_inbounds Base.getindex(a::OffsetRange, r::AbstractRange) = _maybewrapaxes(a.parent[r .- a.offsets[1]], axes(r,1))
-@propagate_inbounds Base.getindex(a::AbstractRange, r::OffsetRange) = OffsetArray(a[parent(r)], axes(r))
 
 for OR in [:IIUR, :IdOffsetRange]
     for R in [:StepRange, :StepRangeLen, :LinRange, :UnitRange]
-        @eval @propagate_inbounds Base.getindex(r::$R, s::$OR) = OffsetArray(r[UnitRange(s)], axes(s))
+        @eval @inline function Base.getindex(r::$R, s::$OR)
+            @boundscheck checkbounds(r, s)
+            @inbounds pr = r[UnitRange(s)]
+            _maybewrapoffset(pr, axes(s,1))
+        end
     end
 
     # this method is needed for ambiguity resolution
-    @eval @propagate_inbounds Base.getindex(r::StepRangeLen{T,<:Base.TwicePrecision,<:Base.TwicePrecision}, s::$OR) where T =
-    OffsetArray(r[UnitRange(s)], axes(s))
-
-    #= Integer UnitRanges may return an appropriate AbstractUnitRange{<:Integer}, as the result may be used in indexing, and
-    indexing is faster with ranges =#
-    @eval @propagate_inbounds function Base.getindex(r::UnitRange{<:Integer}, s::$OR)
-        rs = r[UnitRange(s)]
-        offset_s = first(axes(s,1)) - 1
-        IdOffsetRange(UnitRange(rs .- offset_s), offset_s)
+    @eval @inline function Base.getindex(r::StepRangeLen{T,<:Base.TwicePrecision,<:Base.TwicePrecision}, s::$OR) where T
+        @boundscheck checkbounds(r, s)
+        @inbounds pr = r[UnitRange(s)]
+        _maybewrapoffset(pr, axes(s,1))
     end
 end
 
 # 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{<:Integer}...; kw...)
     ofs = map(A -> first(axes(A,1)) - 1, As)
-    AIds = map((A, of) -> IdOffsetRange(UnitRange(parent(A)) .- of, of), As, ofs)
+    AIds = map((A, of) -> IdOffsetRange(_subtractoffset(parent(A), of), of), As, ofs)
     mapreduce(f, op, AIds...; kw...)
 end
 

src/axes.jl

@@ -20,7 +20,7 @@ julia> ro[-1]
 -1
 
 julia> ro[3]
-ERROR: BoundsError: attempt to access 3-element UnitRange{$Int} at index [5]
+ERROR: BoundsError: attempt to access 3-element OffsetArrays.$(IdOffsetRange{Int,UnitRange{Int}}) with indices -1:1 at index [3]
 ```
 
 If the range doesn't start at 1, the values may be different from the indices:
@@ -35,7 +35,7 @@ julia> ro[-1]
 9
 
 julia> ro[3]
-ERROR: BoundsError: attempt to access 3-element UnitRange{$Int} at index [5]
+ERROR: BoundsError: attempt to access 3-element OffsetArrays.$(IdOffsetRange{Int,UnitRange{Int}}) with indices -1:1 at index [3]
 ```
 
 # Extended help
@@ -121,11 +121,11 @@ end
 
 # Conversions to an AbstractUnitRange{Int} (and to an OrdinalRange{Int,Int} on Julia v"1.6") are necessary
 # to evaluate CartesianIndices for BigInt ranges, as their axes are also BigInt ranges
-AbstractUnitRange{T}(r::IdOffsetRange) where {T<:Integer} = IdOffsetRange{T}(r)
+Base.AbstractUnitRange{T}(r::IdOffsetRange) where {T<:Integer} = IdOffsetRange{T}(r)
 
 # A version upper bound on this may be set after https://github.com/JuliaLang/julia/pull/40038 is merged
 if v"1.6" <= VERSION
-    OrdinalRange{T,T}(r::IdOffsetRange) where {T<:Integer} = IdOffsetRange{T}(r)
+    Base.OrdinalRange{T,T}(r::IdOffsetRange) where {T<:Integer} = IdOffsetRange{T}(r)
 end
 
 # TODO: uncomment these when Julia is ready
@@ -172,15 +172,19 @@ end
     return (ret[1] + r.offset, ret[2])
 end
 
-@propagate_inbounds Base.getindex(r::IdOffsetRange, i::Integer) = r.parent[i - r.offset] + r.offset
-@propagate_inbounds function Base.getindex(r::IdOffsetRange, s::AbstractUnitRange{<:Integer})
-    offset_s = first(axes(s,1)) - 1
-    pr = r.parent[s .- r.offset] .+ (r.offset - offset_s)
-    _maybewrapoffset(pr, offset_s, axes(s,1))
+@inline function Base.getindex(r::IdOffsetRange, i::Integer)
+    @boundscheck checkbounds(r, i)
+    @inbounds r.parent[i - r.offset] + r.offset
+end
+@inline function Base.getindex(r::IdOffsetRange, s::AbstractUnitRange{<:Integer})
+    @boundscheck checkbounds(r, s)
+    @inbounds pr = r.parent[_subtractoffset(s, r.offset)] .+ r.offset
+    _maybewrapoffset(pr, axes(s,1))
 end
 # The following method is required to avoid falling back to getindex(::AbstractUnitRange, ::StepRange{<:Integer})
-@propagate_inbounds function Base.getindex(r::IdOffsetRange, s::StepRange{<:Integer})
-    rs = r.parent[s .- r.offset] .+ r.offset
+@inline function Base.getindex(r::IdOffsetRange, s::StepRange{<:Integer})
+    @boundscheck checkbounds(r, s)
+    @inbounds rs = r.parent[s .- r.offset] .+ r.offset
     return no_offset_view(rs)
 end
 

src/utils.jl

@@ -16,17 +16,17 @@ _offset(axparent::AbstractUnitRange, ::Union{Integer, Colon}) = 1 - first(axpare
 """
     OffsetArrays.AxisConversionStyle(typeof(indices))
 
-`AxisConversionStyle` declares if `indices` should be converted to a single `AbstractUnitRange{Int}` 
-or to a `Tuple{Vararg{AbstractUnitRange{Int}}}` while flattening custom types into indices. 
-This method is called after `to_indices(A::Array, axes(A), indices)` to provide 
+`AxisConversionStyle` declares if `indices` should be converted to a single `AbstractUnitRange{Int}`
+or to a `Tuple{Vararg{AbstractUnitRange{Int}}}` while flattening custom types into indices.
+This method is called after `to_indices(A::Array, axes(A), indices)` to provide
 further information in case `to_indices` does not return a `Tuple` of `AbstractUnitRange{Int}`.
 
-Custom index types should extend `AxisConversionStyle` and return either `OffsetArray.SingleRange()`, 
-which is the default, or `OffsetArray.TupleOfRanges()`. In the former case, the type `T` should 
-define `Base.convert(::Type{AbstractUnitRange{Int}}, ::T)`, whereas in the latter it should define 
-`Base.convert(::Type{Tuple{Vararg{AbstractUnitRange{Int}}}}, ::T)`. 
+Custom index types should extend `AxisConversionStyle` and return either `OffsetArray.SingleRange()`,
+which is the default, or `OffsetArray.TupleOfRanges()`. In the former case, the type `T` should
+define `Base.convert(::Type{AbstractUnitRange{Int}}, ::T)`, whereas in the latter it should define
+`Base.convert(::Type{Tuple{Vararg{AbstractUnitRange{Int}}}}, ::T)`.
 
-An example of the latter is `CartesianIndices`, which is converted to a `Tuple` of 
+An example of the latter is `CartesianIndices`, which is converted to a `Tuple` of
 `AbstractUnitRange{Int}` while flattening the indices.
 
 # Example
@@ -75,10 +75,23 @@ function _checkindices(N::Integer, indices, label)
     N == length(indices) || throw_argumenterror(N, indices, label)
 end
 
-_maybewrapaxes(A::AbstractVector, ::Base.OneTo) = no_offset_view(A)
-_maybewrapaxes(A::AbstractVector, ax) = OffsetArray(A, ax)
-
-_maybewrapoffset(r::AbstractUnitRange, of, ::Base.OneTo) = no_offset_view(r)
-_maybewrapoffset(r::AbstractVector, of, ::Base.OneTo) = no_offset_view(r)
-_maybewrapoffset(r::AbstractUnitRange, of, ::Any) = IdOffsetRange(UnitRange(r), of)
-_maybewrapoffset(r::AbstractVector, of, axs) = OffsetArray(r .+ of, axs)
+@inline _maybewrapoffset(r::AbstractVector, ax::Tuple{Any}) = _maybewrapoffset(r, ax[1])
+@inline _maybewrapoffset(r::AbstractUnitRange{<:Integer}, ::Base.OneTo) = no_offset_view(r)
+@inline _maybewrapoffset(r::AbstractVector, ::Base.OneTo) = no_offset_view(r)
+@inline function _maybewrapoffset(r::AbstractUnitRange{<:Integer}, ax::AbstractUnitRange)
+	of = first(ax) - 1
+	IdOffsetRange(_subtractoffset(r, of), of)
+end
+@inline _maybewrapoffset(r::AbstractVector, ax::AbstractUnitRange) = OffsetArray(r, ax)
+
+# 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))
+
+if VERSION <= v"1.7.0-DEV.1039"
+    _contiguousindexingtype(r::AbstractUnitRange{<:Integer}) = UnitRange{Int}(r)
+else
+    _contiguousindexingtype(r::AbstractUnitRange{<:Integer}) = r
+end