Generalize indexing AbstractUnitRanges with offset ranges

Project.toml

@@ -1,6 +1,6 @@
 name = "OffsetArrays"
 uuid = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
-version = "1.9.2"
+version = "1.10.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/OffsetArrays.jl

@@ -9,6 +9,8 @@ end
 
 export OffsetArray, OffsetMatrix, OffsetVector
 
+const IIUR = IdentityUnitRange{<:AbstractUnitRange{<:Integer}}
+
 include("axes.jl")
 include("utils.jl")
 include("origin.jl")
@@ -434,10 +436,8 @@ Base.dataids(A::OffsetArray) = Base.dataids(parent(A))
 Broadcast.broadcast_unalias(dest::OffsetArray, src::OffsetArray) = parent(dest) === parent(src) ? src : Broadcast.unalias(dest, src)
 
 ### Special handling for AbstractRange
-
 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))
 
@@ -501,23 +501,33 @@ end
     IdOffsetRange(_subtractoffset(parent(r), of), of)
 end
 
+@inline function _boundscheck_index_retaining_axes(r, s)
+    @boundscheck checkbounds(r, s)
+    @inbounds pr = r[UnitRange(s)]
+    _indexedby(pr, axes(s))
+end
+@inline _boundscheck_return(r, s) = (@boundscheck checkbounds(r, s); s)
+
 for OR in [:IIUR, :IdOffsetRange]
-    for R in [:StepRange, :StepRangeLen, :LinRange, :UnitRange]
-        @eval @inline function Base.getindex(r::$R, s::$OR)
-            @boundscheck checkbounds(r, s)
-            @inbounds pr = r[UnitRange(s)]
-            _indexedby(pr, axes(s))
-        end
+    for R in [:StepRange, :StepRangeLen, :LinRange, :AbstractUnitRange]
+        @eval @inline Base.getindex(r::$R, s::$OR) = _boundscheck_index_retaining_axes(r, s)
     end
 
     # this method is needed for ambiguity resolution
     @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)]
-        _indexedby(pr, axes(s))
+        _boundscheck_index_retaining_axes(r, s)
     end
 end
 
+# These methods are added to avoid ambiguities with Base.
+# The ones involving Base types should be ported to Base and version-limited here
+@inline Base.getindex(r::IdentityUnitRange, s::IIUR) = _boundscheck_return(r, s)
+@inline Base.getindex(r::IdentityUnitRange, s::IdOffsetRange) = _boundscheck_return(r, s)
+if IdentityUnitRange !== Base.Slice
+    @inline Base.getindex(r::Base.Slice, s::IIUR) = _boundscheck_return(r, s)
+    @inline Base.getindex(r::Base.Slice, s::IdOffsetRange) = _boundscheck_return(r, s)
+end
+
 # eltype conversion
 # This may use specialized map methods for the parent
 Base.map(::Type{T}, O::OffsetArray) where {T} = parent_call(x -> map(T, x), O)

src/axes.jl

@@ -227,6 +227,14 @@ for T in [:AbstractUnitRange, :StepRange]
     end
 end
 
+# These methods are necessary to avoid ambiguity
+for R in [:IIUR, :IdOffsetRange]
+    @eval @inline function Base.getindex(r::IdOffsetRange, s::$R)
+        @boundscheck checkbounds(r, s)
+        return _getindex(r, s)
+    end
+end
+
 # offset-preserve broadcasting
 Broadcast.broadcasted(::Base.Broadcast.DefaultArrayStyle{1}, ::typeof(-), r::IdOffsetRange{T}, x::Integer) where T =
     IdOffsetRange{T}(r.parent .- x, r.offset)

test/customranges.jl

@@ -0,0 +1,82 @@
+# Useful for testing indexing
+struct ZeroBasedRange{T,A<:AbstractRange{T}} <: AbstractRange{T}
+    a :: A
+    function ZeroBasedRange(a::AbstractRange{T}) where {T}
+        @assert !Base.has_offset_axes(a)
+        new{T, typeof(a)}(a)
+    end
+end
+
+struct ZeroBasedUnitRange{T,A<:AbstractUnitRange{T}} <: AbstractUnitRange{T}
+    a :: A
+    function ZeroBasedUnitRange(a::AbstractUnitRange{T}) where {T}
+        @assert !Base.has_offset_axes(a)
+        new{T, typeof(a)}(a)
+    end
+end
+
+for Z in [:ZeroBasedRange, :ZeroBasedUnitRange]
+    @eval Base.parent(A::$Z) = A.a
+    @eval Base.first(A::$Z) = first(A.a)
+    @eval Base.length(A::$Z) = length(A.a)
+    @eval Base.last(A::$Z) = last(A.a)
+    @eval Base.size(A::$Z) = size(A.a)
+    @eval Base.axes(A::$Z) = map(x -> IdentityUnitRange(0:x-1), size(A.a))
+    @eval Base.getindex(A::$Z, i::Int) = A.a[i + 1]
+    @eval Base.getindex(A::$Z, i::Integer) = A.a[i + 1]
+    @eval Base.firstindex(A::$Z) = 0
+    @eval Base.step(A::$Z) = step(A.a)
+    @eval OffsetArrays.no_offset_view(A::$Z) = A.a
+    @eval function Base.show(io::IO, A::$Z)
+        show(io, A.a)
+        print(io, " with indices $(axes(A,1))")
+    end
+end
+
+for Z in [:ZeroBasedRange, :ZeroBasedUnitRange]
+    for R in [:AbstractRange, :AbstractUnitRange, :StepRange]
+        @eval @inline function Base.getindex(A::$Z, r::$R{<:Integer})
+            @boundscheck checkbounds(A, r)
+            OffsetArrays._indexedby(A.a[r .+ 1], axes(r))
+        end
+    end
+
+    for R in [:ZeroBasedUnitRange, :ZeroBasedRange]
+        @eval @inline function Base.getindex(A::$Z, r::$R{<:Integer})
+            @boundscheck checkbounds(A, r)
+            OffsetArrays._indexedby(A.a[r.a .+ 1], axes(r))
+        end
+    end
+
+    for R in [:IIUR, :IdOffsetRange]
+        @eval @inline function Base.getindex(A::$Z, r::$R)
+            @boundscheck checkbounds(A, r)
+            OffsetArrays._indexedby(A.a[r .+ 1], axes(r))
+        end
+    end
+
+    for R in [:AbstractUnitRange, :IdOffsetRange, :IdentityUnitRange, :SliceIntUR, :StepRange, :StepRangeLen, :LinRange]
+        @eval @inline function Base.getindex(A::$R, r::$Z)
+            @boundscheck checkbounds(A, r)
+            OffsetArrays._indexedby(A[r.a], axes(r))
+        end
+    end
+    @eval @inline function Base.getindex(A::StepRangeLen{<:Any,<:Base.TwicePrecision,<:Base.TwicePrecision}, r::$Z)
+        @boundscheck checkbounds(A, r)
+        OffsetArrays._indexedby(A[r.a], axes(r))
+    end
+end
+
+# A basic range that does not have specialized vector indexing methods defined
+# In this case the best that we may do is to return an OffsetArray
+# Despite this, an indexing operation involving this type should preserve the axes of the indices
+struct CustomRange{T,A<:AbstractRange{T}} <: AbstractRange{T}
+    a :: A
+end
+Base.parent(r::CustomRange) = r.a
+Base.size(r::CustomRange) = size(parent(r))
+Base.axes(r::CustomRange) = axes(parent(r))
+Base.first(r::CustomRange) = first(parent(r))
+Base.last(r::CustomRange) = last(parent(r))
+Base.step(r::CustomRange) = step(parent(r))
+Base.getindex(r::CustomRange, i::Int) = getindex(parent(r), i)

test/runtests.jl

@@ -1,5 +1,6 @@
 using OffsetArrays
-using OffsetArrays: IdentityUnitRange, no_offset_view
+using OffsetArrays: IdentityUnitRange, no_offset_view, IIUR
+using Base: Slice
 using OffsetArrays: IdOffsetRange
 using Test, Aqua, Documenter
 using LinearAlgebra
@@ -11,6 +12,8 @@ using StaticArrays
 using FillArrays
 using DistributedArrays
 
+const SliceIntUR = Slice{<:AbstractUnitRange{<:Integer}}
+
 DocMeta.setdocmeta!(OffsetArrays, :DocTestSetup, :(using OffsetArrays); recursive=true)
 
 # https://github.com/JuliaLang/julia/pull/29440
@@ -26,58 +29,7 @@ struct TupleOfRanges{N}
     x ::NTuple{N, UnitRange{Int}}
 end
 
-# Useful for testing indexing
-struct ZeroBasedRange{T,A<:AbstractRange{T}} <: AbstractRange{T}
-    a :: A
-    function ZeroBasedRange(a::AbstractRange{T}) where {T}
-        @assert !Base.has_offset_axes(a)
-        new{T, typeof(a)}(a)
-    end
-end
-
-struct ZeroBasedUnitRange{T,A<:AbstractUnitRange{T}} <: AbstractUnitRange{T}
-    a :: A
-    function ZeroBasedUnitRange(a::AbstractUnitRange{T}) where {T}
-        @assert !Base.has_offset_axes(a)
-        new{T, typeof(a)}(a)
-    end
-end
-
-for Z in [:ZeroBasedRange, :ZeroBasedUnitRange]
-    @eval Base.parent(A::$Z) = A.a
-    @eval Base.first(A::$Z) = first(A.a)
-    @eval Base.length(A::$Z) = length(A.a)
-    @eval Base.last(A::$Z) = last(A.a)
-    @eval Base.size(A::$Z) = size(A.a)
-    @eval Base.axes(A::$Z) = map(x -> IdentityUnitRange(0:x-1), size(A.a))
-    @eval Base.getindex(A::$Z, i::Int) = A.a[i + 1]
-    @eval Base.firstindex(A::$Z) = 0
-    @eval Base.axes(A::$Z) = map(x -> IdentityUnitRange(0:x-1), size(A.a))
-    @eval Base.getindex(A::$Z, i::Integer) = A.a[i + 1]
-    @eval Base.step(A::$Z) = step(A.a)
-    @eval OffsetArrays.no_offset_view(A::$Z) = A.a
-    @eval function Base.show(io::IO, A::$Z)
-        show(io, A.a)
-        print(io, " with indices $(axes(A,1))")
-    end
-
-    for R in [:AbstractRange, :AbstractUnitRange, :StepRange]
-        @eval @inline function Base.getindex(A::$Z, r::$R{<:Integer})
-            @boundscheck checkbounds(A, r)
-            OffsetArrays._indexedby(A.a[r .+ 1], axes(r))
-        end
-    end
-    for R in [:UnitRange, :StepRange, :StepRangeLen, :LinRange]
-        @eval @inline function Base.getindex(A::$R, r::$Z)
-            @boundscheck checkbounds(A, r)
-            OffsetArrays._indexedby(A[r.a], axes(r))
-        end
-    end
-    @eval @inline function Base.getindex(A::StepRangeLen{<:Any,<:Base.TwicePrecision,<:Base.TwicePrecision}, r::$Z)
-        @boundscheck checkbounds(A, r)
-        OffsetArrays._indexedby(A[r.a], axes(r))
-    end
-end
+include("customranges.jl")
 
 function same_value(r1, r2)
     length(r1) == length(r2) || return false
@@ -1128,6 +1080,10 @@ end
         OffsetArray(IdOffsetRange(IdOffsetRange(10:1000, -1), 1), 3), # offset index
 
         # AbstractRanges
+        Base.OneTo(1000),
+        CustomRange(Base.OneTo(1000)),
+        Slice(Base.OneTo(1000)),
+        SOneTo(1000),
         1:1000,
         UnitRange(1.0, 1000.0),
         1:3:1000,
@@ -1144,6 +1100,7 @@ end
         ZeroBasedUnitRange(1:1000), # offset range
         ZeroBasedRange(1:1000), # offset range
         ZeroBasedRange(1:1:1000), # offset range
+        CustomRange(ZeroBasedRange(1:1:1000)), # offset range
         ]
 
         # AbstractArrays with 1-based indices
@@ -1158,7 +1115,7 @@ end
             test_indexing_axes_and_vals(r1, r2)
             test_indexing_axes_and_vals(r1, collect(r2))
 
-            if r1 isa AbstractRange && axes(r2, 1) isa Base.OneTo
+            if r1 isa AbstractRange && !(r1 isa CustomRange) && axes(r2, 1) isa Base.OneTo
                 @test r1[r2] isa AbstractRange
             end
         end
@@ -1269,13 +1226,18 @@ end
         OffsetArray(IdOffsetRange(IdOffsetRange(10:1000, -1), 1), 3), # offset index
 
         # AbstractRanges
+        Base.OneTo(1000),
+        Slice(Base.OneTo(1000)),
+        SOneTo(1000),
+        CustomRange(Base.OneTo(1000)),
         1:1000,
         UnitRange(1.0, 1000.0),
         1:2:2000,
         2000:-1:1,
         1.0:2.0:2000.0,
         StepRangeLen(Float64(1), Float64(1000), 1000),
         LinRange(1.0, 2000.0, 2000),
+        Base.Slice(Base.OneTo(1000)), # 1-based index
         IdOffsetRange(Base.OneTo(1000)), # 1-based index
         IdOffsetRange(1:1000, 0), # 1-based index
         IdOffsetRange(Base.OneTo(1000), 4), # offset index
@@ -1288,6 +1250,7 @@ end
         ZeroBasedRange(1:1000), # offset index
         ZeroBasedRange(1:1:1000), # offset index
         ZeroBasedUnitRange(IdentityUnitRange(1:1000)), # offset index
+        CustomRange(ZeroBasedUnitRange(IdentityUnitRange(1:1000))), # offset index
         ]
 
         # AbstractArrays with offset axes