More generalizations

Author: mtfishman

src/BlockArraysExtensions/blockedunitrange.jl

@@ -176,6 +176,19 @@ function blockedunitrange_getindices(
   return mortar(map(b -> a[b], blocks(indices)))
 end
 
+function blockedunitrange_getindices(
+  a::AbstractBlockedUnitRange,
+  indices::BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexVector{1}}},
+)
+  return mortar(map(b -> a[b], blocks(indices)))
+end
+function blockedunitrange_getindices(
+  a::AbstractBlockedUnitRange,
+  indices::BlockVector{<:GenericBlockIndex{1},<:Vector{<:BlockIndexVector{1}}},
+)
+  return mortar(map(b -> a[b], blocks(indices)))
+end
+
 function blockedunitrange_getindices(
   a::AbstractBlockedUnitRange, indices::AbstractVector{Bool}
 )

src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl

@@ -78,13 +78,36 @@ function Base.to_indices(
   return @interface interface(a) to_indices(a, inds, I)
 end
 
+# a[mortar([Block(1)[[1, 2]], Block(2)[[1, 3]]])]
+function Base.to_indices(
+  a::AnyAbstractBlockSparseArray,
+  inds,
+  I::Tuple{BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexVector{1}}},Vararg{Any}},
+)
+  return @interface interface(a) to_indices(a, inds, I)
+end
+function Base.to_indices(
+  a::AnyAbstractBlockSparseArray,
+  inds,
+  I::Tuple{BlockVector{<:GenericBlockIndex{1},<:Vector{<:BlockIndexVector{1}}},Vararg{Any}},
+)
+  return @interface interface(a) to_indices(a, inds, I)
+end
+
 # a[[Block(1)[1:2], Block(2)[1:2]], [Block(1)[1:2], Block(2)[1:2]]]
 function Base.to_indices(
   a::AnyAbstractBlockSparseArray, inds, I::Tuple{Vector{<:BlockIndexRange{1}},Vararg{Any}}
 )
   return to_indices(a, inds, (mortar(I[1]), Base.tail(I)...))
 end
 
+# a[[Block(1)[[1, 2]], Block(2)[[1, 2]]], [Block(1)[[1, 2]], Block(2)[[1, 2]]]]
+function Base.to_indices(
+  a::AnyAbstractBlockSparseArray, inds, I::Tuple{Vector{<:BlockIndexVector{1}},Vararg{Any}}
+)
+  return to_indices(a, inds, (mortar(I[1]), Base.tail(I)...))
+end
+
 # BlockArrays `AbstractBlockArray` interface
 function BlockArrays.blocks(a::AnyAbstractBlockSparseArray)
   @interface interface(a) blocks(a)

src/blocksparsearrayinterface/blocksparsearrayinterface.jl

@@ -229,6 +229,23 @@ end
   return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...)
 end
 
+@interface ::AbstractBlockSparseArrayInterface function Base.to_indices(
+  a,
+  inds,
+  I::Tuple{BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexVector{1}}},Vararg{Any}},
+)
+  I1 = BlockIndices(I[1], blockedunitrange_getindices(inds[1], I[1]))
+  return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...)
+end
+@interface ::AbstractBlockSparseArrayInterface function Base.to_indices(
+  a,
+  inds,
+  I::Tuple{BlockVector{<:GenericBlockIndex{1},<:Vector{<:BlockIndexVector{1}}},Vararg{Any}},
+)
+  I1 = BlockIndices(I[1], blockedunitrange_getindices(inds[1], I[1]))
+  return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...)
+end
+
 # a[BlockVector([Block(2), Block(1)], [2]), BlockVector([Block(2), Block(1)], [2])]
 # Permute and merge blocks.
 # TODO: This isn't merging blocks yet, that needs to be implemented that.