diff --git a/src/blockaxis.jl b/src/blockaxis.jl index b6c151ae..9bc0b3f7 100644 --- a/src/blockaxis.jl +++ b/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}) diff --git a/test/test_blockindices.jl b/test/test_blockindices.jl index d9e622af..0485a239 100644 --- a/test/test_blockindices.jl +++ b/test/test_blockindices.jl @@ -329,6 +329,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,)