Preserve blocking of indices when slicing a unit range with a blocked unit range (#465)

Fixes #367.

---------

Co-authored-by: Sheehan Olver <[email protected]>

Author: mtfishman

src/blockaxis.jl

@@ -105,6 +105,16 @@ function Base.AbstractUnitRange{T}(r::BlockedUnitRange) where {T}
     return _BlockedUnitRange(convert(T,first(r)), convert.(T,blocklasts(r)))
 end
 
+# See: https://github.com/JuliaLang/julia/blob/b06d26075bf7b3f4e7f1b64b120f5665d8ed76f9/base/range.jl#L991-L1004
+function Base.getindex(r::AbstractUnitRange, s::AbstractBlockedUnitRange{T}) where {T<:Integer}
+    @boundscheck checkbounds(r, s)
+
+    f = first(r)
+    start = oftype(f, f + first(s) - firstindex(r))
+    lens = map(Base.Fix1(oftype, f), blocklengths(s))
+    return blockedrange(start, lens)
+end
+
 """
     BlockedOneTo{T, <:Union{AbstractVector{T}, NTuple{<:Any,T}}} where {T}
 
@@ -164,6 +174,16 @@ function Base.AbstractUnitRange{T}(r::BlockedOneTo) where {T}
     return BlockedOneTo(convert.(T,blocklasts(r)))
 end
 
+# See: https://github.com/JuliaLang/julia/blob/b06d26075bf7b3f4e7f1b64b120f5665d8ed76f9/base/range.jl#L1006-L1010
+function getindex(r::Base.OneTo{T}, s::BlockedOneTo) where T
+    @inline
+    @boundscheck checkbounds(r, s)
+    return BlockedOneTo(convert(AbstractVector{T}, blocklasts(s)))
+end
+function getindex(r::BlockedOneTo{T}, s::BlockedOneTo) where T
+    return Base.oneto(r)[s]
+end
+
 """
     blockedrange(blocklengths::Union{Tuple, AbstractVector})
     blockedrange(first::Integer, blocklengths::Union{Tuple, AbstractVector})

test/test_blockindices.jl

@@ -330,6 +330,40 @@ end
         end
     end
 
+    @testset "BlockedUnitRange indexing" begin
+        a = 2:10
+        b = blockedrange(2, [1,2,3])
+        @test a[b] == blockedrange(3, [1,2,3])
+        @test a[b] isa BlockedUnitRange
+        @test first(a[b]) == 3
+        @test blocklengths(a[b]) == [1,2,3]
+        @test_throws BoundsError a[blockedrange(5, [1,2,3])]
+
+        a = blockedrange(2, [4,5])
+        b = blockedrange(2, [1,2,3])
+        @test a[b] == blockedrange(3, [1,2,3])
+        @test a[b] isa BlockedUnitRange
+        @test first(a[b]) == 3
+        @test blocklengths(a[b]) == [1,2,3]
+        @test_throws BoundsError a[blockedrange(5, [1,2,3])]
+
+        a = Base.OneTo(9)
+        b = blockedrange([1,2,3])
+        @test a[b] == blockedrange([1,2,3])
+        @test a[b] isa BlockedOneTo
+        @test first(a[b]) == 1
+        @test blocklengths(a[b]) == [1,2,3]
+        @test_throws BoundsError a[blockedrange([1,2,3,4])]
+
+        a = blockedrange([4,5])
+        b = blockedrange([1,2,3])
+        @test a[b] == blockedrange([1,2,3])
+        @test a[b] isa BlockedOneTo
+        @test first(a[b]) == 1
+        @test blocklengths(a[b]) == [1,2,3]
+        @test_throws BoundsError a[blockedrange([1,2,3,4])]
+    end
+
     @testset "misc" begin
         b = blockedrange(1, [1,2,3])
         @test axes(b) == Base.unsafe_indices(b) == (b,)