Preserve block information in more slicing operations (#459)

This introduces a generalization of `BlockSlice` for views involving
vectors of `Block` and `BlockIndexRange`, building off of #445.

I call it `BlockedSlice` but I'm open to suggestions. At first I tried
to just generalize `BlockSlice` to be an `AbstractVector` subtype but
some operations depend on it being an `AbstractUnitRange` so I think it
makes sense to have both.

The motivation is that in
https://github.com/ITensor/BlockSparseArrays.jl it would be helpful to
preserve information about the original blocks being sliced when taking
views involving vectors of `Block` and `BlockIndexRange`. #445 drops the
blocks input in the slice operation, so it is hard to tell if a certain
slice operation was really originally blockwise.

@dlfivefifty curious to hear your thoughts on this since you wrote #445.

Author: mtfishman

src/blockindices.jl

@@ -290,6 +290,27 @@ _indices(B) = B
 
 Block(bs::BlockSlice{<:BlockIndexRange}) = Block(bs.block)
 
+"""
+    BlockedSlice(blocks, indices)
+
+Represents blocked indices attached to a collection of corresponding blocks.
+
+Upon calling `to_indices()`, a collection of blocks are converted to BlockedSlice objects to represent
+the indices over which the blocks span.
+
+This mimics the relationship between `Colon` and `Base.Slice`, `Block` and `BlockSlice`, etc.
+"""
+struct BlockedSlice{BB,T<:Integer,INDS<:AbstractVector{T}} <: AbstractVector{T}
+    blocks::BB
+    indices::INDS
+end
+
+for f in (:axes, :size)
+    @eval $f(S::BlockedSlice) = $f(S.indices)
+end
+
+@propagate_inbounds getindex(S::BlockedSlice, i::Integer) = getindex(S.indices, i)
+@propagate_inbounds getindex(S::BlockedSlice, k::Block{1}) = BlockSlice(S.blocks[Int(k)], getindex(S.indices, k))
 
 struct BlockRange{N,R<:NTuple{N,AbstractUnitRange{<:Integer}}} <: AbstractArray{Block{N,Int},N}
     indices::R

src/views.jl

@@ -11,7 +11,7 @@ function unblock(A, inds, I)
 end
 
 _blockslice(B, a::AbstractUnitRange) = BlockSlice(B, a)
-_blockslice(B, a) = a # drop block structure
+_blockslice(B, a) = BlockedSlice(B, a)
 
 # Allow `ones(2)[Block(1)[1:1], Block(1)[1:1]]` which is
 # similar to `ones(2)[1:1, 1:1]`.
@@ -150,10 +150,11 @@ block(A::Block) = A
 @inline view(block_arr::AbstractBlockArray{<:Any,N}, blocks::Vararg{BlockSlice1, N}) where N =
     view(block_arr, map(block,blocks)...)
 
-const BlockSlices = Union{Base.Slice,BlockSlice{<:BlockRange{1}}}
+const BlockSlices = Union{Base.Slice,BlockSlice{<:BlockRange{1}},BlockedSlice}
 # view(V::SubArray{<:Any,N,NTuple{N,BlockSlices}},
 
 _block_reindex(b::BlockSlice, i::Block{1}) = b.block[Int(i)]
+_block_reindex(b::BlockedSlice, i::Block{1}) = b.blocks[Int(i)]
 _block_reindex(b::Slice, i::Block{1}) = i
 
 @inline view(V::SubArray{<:Any,N,<:AbstractBlockArray,<:NTuple{N,BlockSlices}}, block::Block{N}) where N =

test/test_blockindices.jl

@@ -2,7 +2,7 @@ module TestBlockIndices
 
 using BlockArrays, FillArrays, Test, StaticArrays, ArrayLayouts
 using OffsetArrays
-import BlockArrays: BlockIndex, BlockIndexRange, BlockSlice
+import BlockArrays: BlockIndex, BlockIndexRange, BlockSlice, BlockedSlice
 
 @testset "Blocks" begin
     @test Int(Block(2)) === Integer(Block(2)) === Number(Block(2)) === 2
@@ -831,6 +831,16 @@ end
     end
 end
 
+@testset "BlockedSlice" begin
+    b = BlockedSlice([Block(2), Block(1)], mortar([3:5, 1:2]))
+    @test length(b) == 5
+    for i in eachindex(b.indices)
+        @test b[i] === b.indices[i]
+    end
+    @test b[Block(1)] === BlockSlice(Block(2), 3:5)
+    @test b[Block(2)] === BlockSlice(Block(1), 1:2)
+end
+
 #=
 [1,1  1,2] | [1,3  1,4  1,5]
 --------------------------

test/test_blockrange.jl

@@ -23,6 +23,15 @@ using BlockArrays, Test
     @test view(A, Block.(1:2)) == [1,2,3]
     @test A[Block.(1:2)] == [1,2,3]
 
+    V = view(A, [Block(3), Block(2)])
+    @test V == [4, 5, 6, 2, 3]
+    I = parentindices(V)[1]
+    @test I isa BlockArrays.BlockedSlice{<:Vector{<:Block{1}}}
+    @test V[Block(1)] == 4:6
+    @test V[Block(2)] == 2:3
+    @test view(V, Block(1)) === view(A, Block(3))
+    @test view(V, Block(2)) === view(A, Block(2))
+
     A = BlockArray(reshape(collect(1:(6*12)),6,12), 1:3, 3:5)
 
     @test view(A, Block.(1:2), Block.(1:2)) == A[1:3,1:7]