Preserve axes when indexing ranges with IdentityUnitRange

Author: timholy

src/OffsetArrays.jl

@@ -160,13 +160,28 @@ end
 ### Special handling for AbstractRange
 
 const OffsetRange{T} = OffsetArray{T,1,<:AbstractRange{T}}
+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)
 
+@inline @propagate_inbounds Base.getindex(r::UnitRange, s::IIUR) =
+    OffsetArray(r[s.indices], s)
+
+@inline @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 =
+    OffsetArray(r[s.indices], s)
+@inline @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) =
+    OffsetArray(r[s.indices], s)
+
 ### Convenience functions ###
 
 Base.fill(x, inds::Tuple{UnitRange,Vararg{UnitRange}}) =

test/runtests.jl

@@ -312,6 +312,12 @@ end
     @test axes(v) === (IdentityUnitRange(-4:-3),)
     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 s in (IdentityUnitRange(2:3), OffsetArray(2:3, 2:3))
+            @test axes(r[s]) == axes(s)
+        end
+    end
 end
 
 @testset "logical indexing" begin