Vector indexing for BlockRange (#266)

JuliaArrays · Jun 21, 2023 · 881e81e · 881e81e · jishnub · Jun 21, 2023
1 parent f1826ef
commit 881e81e
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 3 deletions.
diff --git a/src/blockindices.jl b/src/blockindices.jl
@@ -335,15 +335,19 @@ broadcasted(::DefaultArrayStyle{0}, ::Type{Int}, block::Block{1}) = Int(block)
 
 # AbstractArray implementation
 axes(iter::BlockRange{N,R}) where {N,R} = map(axes1, iter.indices)
-@inline function Base.getindex(iter::BlockRange{N,<:NTuple{N,Base.OneTo}}, I::Vararg{Integer, N}) where {N}
+@inline function Base.getindex(iter::BlockRange{N,<:NTuple{N,Base.OneTo}}, I::Vararg{Int, N}) where {N}
     @boundscheck checkbounds(iter, I...)
     Block(I)
 end
-@inline function Base.getindex(iter::BlockRange{N,R}, I::Vararg{Integer, N}) where {N,R}
+@propagate_inbounds function Base.getindex(iter::BlockRange{N}, I::Vararg{Int, N}) where {N}
     @boundscheck checkbounds(iter, I...)
-    Block(I .- first.(axes1.(iter.indices)) .+ first.(iter.indices))
+    Block(getindex.(iter.indices, I))
 end
 
+@propagate_inbounds function Base.getindex(iter::BlockRange{N}, I::Vararg{AbstractUnitRange{<:Integer}, N}) where {N}
+    @boundscheck checkbounds(iter, I...)
+    BlockRange(getindex.(iter.indices, I))
+end
 
 @inline function iterate(iter::BlockRange)
     iterfirst, iterlast = first(iter), last(iter)

diff --git a/test/test_blockindices.jl b/test/test_blockindices.jl
@@ -156,6 +156,13 @@ end
         @test b[Block(100_000)] == 299_998:300_000
         @test_throws BlockBoundsError b[Block(0)]
         @test_throws BlockBoundsError b[Block(1_000_001)]
+
+        b = BlockRange((2:4, 3:4))
+        @test b[2,2] === Block(3,4)
+        @test b[axes(b)...] === b
+
+        b = BlockRange(OffsetArrays.IdOffsetRange.((2:4, 3:5), 2))
+        @test b[axes(b)...] === b
     end
 
     @testset "firsts/lasts/lengths" begin