Idempotent indexing using offset ranges (#156)

* preserve axes while indexing with IdOffsetRange

* Idempotent indexing for IdentityUnitRange

Author: jishnub

src/OffsetArrays.jl

@@ -305,22 +305,29 @@ const IIUR = IdentityUnitRange{S} where S<:AbstractUnitRange{T} where T<:Integer
 
 Base.step(a::OffsetRange) = step(parent(a))
 
-Base.getindex(a::OffsetRange, r::OffsetRange) = OffsetArray(a[parent(r)], r.offsets)
-Base.getindex(a::OffsetRange, r::AbstractRange) = a.parent[r .- a.offsets[1]]
-Base.getindex(a::AbstractRange, r::OffsetRange) = OffsetArray(a[parent(r)], r.offsets)
+@propagate_inbounds Base.getindex(a::OffsetRange, r::OffsetRange) = OffsetArray(a[parent(r)], r.offsets)
+@propagate_inbounds function Base.getindex(a::OffsetRange, r::IdOffsetRange)
+    OffsetArray(a.parent[r.parent .+ (r.offset - a.offsets[1])], r.offset)
+end
+@propagate_inbounds Base.getindex(r::OffsetRange, s::IIUR) =
+    OffsetArray(r[s.indices], s)
+@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)
 
-@inline @propagate_inbounds Base.getindex(r::StepRangeLen{T,<:Base.TwicePrecision,<:Base.TwicePrecision}, s::IIUR) where T =
+# 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)
-@inline @propagate_inbounds Base.getindex(r::StepRangeLen{T}, s::IIUR) where {T} =
+
+@propagate_inbounds Base.getindex(r::StepRangeLen{T}, s::IIUR) where {T} =
     OffsetArray(r[s.indices], s)
 
-@inline @propagate_inbounds Base.getindex(r::LinRange, s::IIUR) =
+@propagate_inbounds Base.getindex(r::LinRange, s::IIUR) =
     OffsetArray(r[s.indices], s)
 
 function Base.show(io::IO, r::OffsetRange)

test/runtests.jl

@@ -231,8 +231,15 @@ end
     r = -2:5
     y = OffsetArray(r, r)
     @test axes(y) == (r,)
+    @test step(y) == step(r)
     y = OffsetArray(r, (r,))
     @test axes(y) == (r,)
+
+    s = -2:2:4
+    r = 5:8
+    y = OffsetArray(s, r)
+    @test axes(y) == (r,)
+    @test step(y) == step(s)
 end
 
 @testset "OffsetArray of OffsetArray construction" begin
@@ -351,6 +358,19 @@ end
     r1 = OffsetArray(8:10, -1:1)[OffsetArray(0:1, -5:-4)]
     @test axes(r1) == (IdentityUnitRange(-5:-4),)
     @test parent(r1) === 9:10
+
+    a = OffsetVector(3:4, 10:11)
+    ax = OffsetArrays.IdOffsetRange(5:6, 5)
+    @test axes(a[ax]) == axes(ax)
+    for i in axes(ax,1)
+        @test a[ax[i]] == a[ax][i]
+    end
+
+    ax = IdentityUnitRange(10:11)
+    @test axes(a[ax]) == axes(ax)
+    for i in axes(ax,1)
+        @test a[ax[i]] == a[ax][i]
+    end
 end
 
 @testset "CartesianIndexing" begin
@@ -651,7 +671,7 @@ end
     v = view(A0, 1:1, i1)
     @test axes(v) == (Base.OneTo(1), IdentityUnitRange(-4:-3))
 
-    for r in (1:10, 1:1:10, StepRangeLen(1, 1, 10), LinRange(1, 10, 10))
+    for r in (1:10, 1:1:10, StepRangeLen(1, 1, 10), LinRange(1, 10, 10), 0.1:0.2:0.9)
         for s in (IdentityUnitRange(2:3), OffsetArray(2:3, 2:3))
             @test axes(r[s]) == axes(s)
         end