Index into a CartesianIndices with a Block (#354)

* View into ranges/CartesianIndices with BlockSlices

* Index into a CartesianIndices with a Block

* Preserve indexing of Blocks in BlockSlice

* Block{1} indexing

* Add BlockSlice type test

Author: jishnub

src/blockindices.jl

@@ -113,6 +113,16 @@ function Base.show(io::IO, B::Block{N,Int}) where N
     print(io, ")")
 end
 
+# Some views may be computed eagerly without the SubArray wrapper
+Base.@propagate_inbounds Base.view(r::AbstractRange, B::Block{1}) = r[to_indices(r, (B,))...]
+Base.@propagate_inbounds function Base.view(C::CartesianIndices{N}, b1::Block{1}, B::Block{1}...) where {N}
+    blk = Block((b1, B...))
+    view(C, to_indices(C, (blk,))...)
+end
+Base.@propagate_inbounds function Base.view(C::CartesianIndices{N}, B::Block{N}) where {N}
+    view(C, to_indices(C, (B,))...)
+end
+
 """
     BlockIndex{N}
 
@@ -289,11 +299,19 @@ for f in (:axes, :unsafe_indices, :axes1, :first, :last, :size, :length,
     @eval $f(S::BlockSlice) = $f(S.indices)
 end
 
+_indices(B::BlockSlice) = B.indices
+_indices(B) = B
+
 @propagate_inbounds getindex(S::BlockSlice, i::Integer) = getindex(S.indices, i)
-@propagate_inbounds getindex(S::BlockSlice{<:Block}, k::AbstractUnitRange{Int}) = BlockSlice(S.block[k],S.indices[k])
-@propagate_inbounds getindex(S::BlockSlice{<:BlockIndexRange}, k::AbstractUnitRange{Int}) = BlockSlice(S.block[k],S.indices[k])
+@propagate_inbounds getindex(S::BlockSlice{<:Block{1}}, k::AbstractUnitRange{Int}) =
+    BlockSlice(S.block[_indices(k)], S.indices[_indices(k)])
+@propagate_inbounds getindex(S::BlockSlice{<:BlockIndexRange{1}}, k::AbstractUnitRange{Int}) =
+    BlockSlice(S.block[_indices(k)], S.indices[_indices(k)])
 show(io::IO, r::BlockSlice) = print(io, "BlockSlice(", r.block, ",", r.indices, ")")
 
+# Avoid creating a SubArray wrapper in certain non-allocating cases
+Base.@propagate_inbounds Base.view(C::CartesianIndices{N}, bs::Vararg{BlockSlice,N}) where {N} = view(C, map(x->x.indices, bs)...)
+
 Block(bs::BlockSlice{<:BlockIndexRange}) = Block(bs.block)
 
 

test/test_blockindices.jl

@@ -103,6 +103,37 @@ import BlockArrays: BlockIndex, BlockIndexRange, BlockSlice
         @test intersect(Block.(2:5), Block.(3:6)) ≡ Block.(3:5)
     end
 
+    @testset "view into CartesianIndices/ranges" begin
+        r = 2:10
+        @test view(r, Block(1)) === r
+        @test_throws BlockBoundsError view(r, Block(2))
+        C = CartesianIndices((2:10, 2:10))
+        ==ᵥ = VERSION >= v"1.10" ? (===) : (==)
+
+        @test view(C, Block(1,1)) ==ᵥ C
+        @test view(C, Block(1), Block(1)) == C
+        @test_throws BlockBoundsError view(C, Block(1), Block(2))
+        @test_throws BlockBoundsError view(C, Block(1,2))
+
+        B = BlockArray([1:3;], [2,1])
+        Cb = CartesianIndices(B)
+        @test view(Cb, Block(1)) ==ᵥ Cb[Block(1)] ==ᵥ CartesianIndices((1:2,))
+        @test view(Cb, Block(2)) ==ᵥ Cb[Block(2)] ==ᵥ CartesianIndices((3:3,))
+
+        B = BlockArray(reshape([1:9;],3,3), [2,1], [2,1])
+        Cb = CartesianIndices(B)
+        @test view(Cb, Block(1,1)) ==ᵥ Cb[Block(1,1)] ==ᵥ CartesianIndices((1:2,1:2))
+        @test view(Cb, Block(1,2)) ==ᵥ Cb[Block(1,2)] ==ᵥ CartesianIndices((1:2, 3:3))
+        @test view(Cb, Block(2,1)) ==ᵥ Cb[Block(2,1)] ==ᵥ CartesianIndices((3:3,1:2))
+        @test view(Cb, Block(2,2)) ==ᵥ Cb[Block(2,2)] ==ᵥ CartesianIndices((3:3, 3:3))
+        for i in 1:2, j in 1:2
+            @test view(Cb, Block(j), Block(i)) ==ᵥ view(Cb, Block(j, i))
+        end
+        # ensure that calls with mismatched ndims don't error
+        @test view(Cb, Block(1)) == view(Cb, to_indices(Cb, (Block(1),))...)
+        @test reshape(view(Cb, Block(1), Block(1), Block(1)), 2, 2) == view(Cb, Block(1), Block(1))
+    end
+
     @testset "print" begin
         @test sprint(show, "text/plain", Block()) == "Block()"
         @test sprint(show, "text/plain", Block(1)) == "Block(1)"
@@ -405,6 +436,8 @@ end
 
 @testset "BlockSlice" begin
     b = BlockSlice(Block(5),1:3)
+    @test b[b] == b
+    @test b[b] isa BlockSlice{<:BlockIndexRange}
     @test b[Base.Slice(1:3)] ≡ b
     @test b[1:2] ≡ b[1:2][1:2] ≡ BlockSlice(Block(5)[1:2],1:2)
     @test Block(b) ≡ Block(5)
@@ -414,6 +447,16 @@ end
             @test convert(typeof(b), Base.OneTo(1)) ≡ b
         end
     end
+
+    @testset "view into CartesianIndices/ranges" begin
+        C = CartesianIndices((1:3,))
+        r = 1:2
+        b = BlockSlice(Block(1), r)
+        @test view(C, b) === view(C, r)
+        @test view(1:10, b) === view(1:10, r)
+        C = CartesianIndices((1:3, 1:3))
+        @test view(C, b, b) === view(C, r, r)
+    end
 end
 
 #=