Idempotent indexing in range[IdOffsetRange] (#161)

* Idempotent indexing range[IdOffsetRange]

* Add methods to index a UnitRange{<:Integer}

Author: jishnub

src/OffsetArrays.jl

@@ -308,21 +308,20 @@ end
 @propagate_inbounds Base.getindex(a::OffsetRange, r::AbstractRange) = a.parent[r .- a.offsets[1]]
 @propagate_inbounds Base.getindex(a::AbstractRange, r::OffsetRange) = OffsetArray(a[parent(r)], r.offsets)
 
-@propagate_inbounds Base.getindex(r::UnitRange, s::IIUR) =
-    OffsetArray(r[s.indices], s)
-
-@propagate_inbounds Base.getindex(r::StepRange, s::IIUR) =
-    OffsetArray(r[s.indices], s)
-
-# this method is needed for ambiguity resolution
-@propagate_inbounds Base.getindex(r::StepRangeLen{T,<:Base.TwicePrecision,<:Base.TwicePrecision}, s::IIUR) where T =
-    OffsetArray(r[s.indices], s)
+for OR in [:IIUR, :IdOffsetRange]
+    for R in [:StepRange, :StepRangeLen, :LinRange, :UnitRange]
+        @eval @propagate_inbounds Base.getindex(r::$R, s::$OR) = OffsetArray(r[_unwrap(s)], axes(s))
+    end
 
-@propagate_inbounds Base.getindex(r::StepRangeLen{T}, s::IIUR) where {T} =
-    OffsetArray(r[s.indices], s)
+    # 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[_unwrap(s)], axes(s))
+end
 
-@propagate_inbounds Base.getindex(r::LinRange, s::IIUR) =
-    OffsetArray(r[s.indices], s)
+#= Integer UnitRanges may return an appropriate AbstractUnitRange{<:Integer}, as the result may be used in indexing, and
+indexing is faster with ranges =#
+@propagate_inbounds Base.getindex(r::UnitRange{<:Integer}, s::IdOffsetRange) = IdOffsetRange(r[_unwrap(s)] .- s.offset, s.offset)
+@propagate_inbounds Base.getindex(r::UnitRange{<:Integer}, s::IIUR) = IdentityUnitRange(r[_unwrap(s)])
 
 function Base.show(io::IO, r::OffsetRange)
     show(io, r.parent)

src/utils.jl

@@ -70,4 +70,7 @@ _checkindices(A::AbstractArray, indices, label) = _checkindices(ndims(A), indice
 function _checkindices(N::Integer, indices, label)
     throw_argumenterror(N, indices, label) = throw(ArgumentError(label*" $indices are not compatible with a $(N)D array"))
     N == length(indices) || throw_argumenterror(N, indices, label)
-end
+end
+
+_unwrap(r::IdOffsetRange) = r.parent .+ r.offset
+_unwrap(r::IdentityUnitRange) = r.indices

test/runtests.jl

@@ -110,27 +110,44 @@ end
     check_indexed_by(r3, 0:2)
 
     @testset "Idempotent indexing" begin
-        r = OffsetArrays.IdOffsetRange(3:5, -1)
-        
-        # Indexing with IdentityUnitRange
-        s = IdentityUnitRange(0:2)
-        @test axes(r[s]) == axes(s)
-        for i in eachindex(s)
-            @test r[s[i]] == r[s][i]
-        end
+        @testset "Indexing into an IdOffsetRange" begin
+            r = OffsetArrays.IdOffsetRange(3:5, -1)
+            # Indexing with IdentityUnitRange
+            s = IdentityUnitRange(0:2)
+            @test axes(r[s]) == axes(s)
+            for i in eachindex(s)
+                @test r[s[i]] == r[s][i]
+            end
 
-        # Indexing with IdOffsetRange
-        s = OffsetArrays.IdOffsetRange(-4:-2, 4)
-        @test axes(r[s]) == axes(s)
-        for i in eachindex(s)
-            @test r[s[i]] == r[s][i]
+            # Indexing with IdOffsetRange
+            s = OffsetArrays.IdOffsetRange(-4:-2, 4)
+            @test axes(r[s]) == axes(s)
+            for i in eachindex(s)
+                @test r[s[i]] == r[s][i]
+            end
+
+            # Indexing with UnitRange
+            s = 0:2
+            @test axes(r[s]) == axes(s)
+            for i in eachindex(s)
+                @test r[s[i]] == r[s][i]
+            end
         end
+        @testset "Indexing using an IdOffsetRange" begin
+            r = OffsetArrays.IdOffsetRange(3:5, -1)
+            # Indexing into an IdentityUnitRange
+            s = IdentityUnitRange(-1:5)
+            @test axes(s[r]) == axes(r)
+            for i in eachindex(r)
+                @test s[r[i]] == s[r][i]
+            end
 
-        # Indexing with UnitRange
-        s = 0:2
-        @test axes(r[s]) == axes(s)
-        for i in eachindex(s)
-            @test r[s[i]] == r[s][i]
+            # Indexing into an UnitRange
+            s = -3:6
+            @test axes(s[r]) == axes(r)
+            for i in eachindex(r)
+                @test s[r[i]] == s[r][i]
+            end
         end
     end