Get logical indexing working

Author: mtfishman

src/BlockArraysExtensions/BlockArraysExtensions.jl

@@ -376,8 +376,11 @@ function blockrange(axis::AbstractUnitRange, r::BlockIndexVector)
   return Block(r):Block(r)
 end
 
-function blockrange(axis::AbstractUnitRange, r::Vector{<:BlockIndexVector})
-  return map(Block, r)
+function blockrange(
+  axis::AbstractUnitRange,
+  r::BlockVector{<:BlockIndex{1},<:AbstractVector{<:BlockIndexVector}},
+)
+  return map(Block, blocks(r))
 end
 
 function blockrange(axis::AbstractUnitRange, r)

src/abstractblocksparsearray/views.jl

@@ -247,8 +247,9 @@ end
 
 # Block slice of the result of slicing `@view a[2:5, 2:5]`.
 # TODO: Move this to `BlockArraysExtensions`.
-const BlockedSlice = BlockSlice{
-  <:BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}}
+const BlockedSlice = Union{
+  BlockSlice{<:BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}}},
+  BlockIndices{<:BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexVector}}},
 }
 
 function Base.view(
@@ -269,6 +270,11 @@ function BlockArrays.viewblock(
 ) where {T,N}
   return viewblock(a, to_tuple(block)...)
 end
+
+_block(x) = error("Not implemented.")
+_block(x::BlockSlice) = x.block
+_block(x::BlockIndices) = x.blocks
+
 # TODO: Define `@interface BlockSparseArrayInterface() viewblock`.
 function BlockArrays.viewblock(
   a::SubArray{T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{BlockedSlice,N}}},
@@ -279,7 +285,7 @@ function BlockArrays.viewblock(
     # TODO: Ideally we would use this but it outputs a Vector,
     # not a range:
     # return parentindices(a)[dim].block[I[dim]]
-    return blocks(parentindices(a)[dim].block)[Int(I[dim])]
+    return blocks(_block(parentindices(a)[dim]))[Int(I[dim])]
   end
   return @view parent(a)[brs...]
 end