Fix some more tests

NDTensors/src/lib/BlockSparseArrays/src/BlockArraysExtensions/BlockArraysExtensions.jl

@@ -58,9 +58,15 @@ end
 # ```
 # but includes `BlockIndices`, where the blocks aren't contiguous.
 const BlockSliceCollection = Union{
-  BlockSlice{<:BlockRange{1}},BlockIndices{<:Vector{<:Block{1}}}
+  Base.Slice,BlockSlice{<:BlockRange{1}},BlockIndices{<:Vector{<:Block{1}}}
+}
+const SubBlockSliceCollection = BlockIndices{
+  <:BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}}
 }
 
+# Slice `a` by `I`, returning a:
+# `BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}}`
+# with the `BlockIndex{1}` corresponding to each value of `I`.
 function to_blockindices(a::BlockedOneTo{<:Integer}, I::UnitRange{<:Integer})
   return mortar(
     map(blocks(blockedunitrange_getindices(a, I))) do r
@@ -240,6 +246,12 @@ function blockrange(axis::AbstractUnitRange, r::UnitRange)
   return findblock(axis, first(r)):findblock(axis, last(r))
 end
 
+# Occurs when slicing with `a[2:4, 2:4]`.
+function blockrange(axis::BlockedOneTo{<:Integer}, r::BlockedUnitRange{<:Integer})
+  # TODO: Check the blocks are commensurate.
+  return findblock(axis, first(r)):findblock(axis, last(r))
+end
+
 function blockrange(axis::AbstractUnitRange, r::Int)
   ## return findblock(axis, r)
   return error("Slicing with integer values isn't supported.")

NDTensors/src/lib/BlockSparseArrays/src/BlockSparseArrays.jl

@@ -5,6 +5,7 @@ include("blocksparsearrayinterface/linearalgebra.jl")
 include("blocksparsearrayinterface/blockzero.jl")
 include("blocksparsearrayinterface/broadcast.jl")
 include("blocksparsearrayinterface/arraylayouts.jl")
+include("blocksparsearrayinterface/views.jl")
 include("abstractblocksparsearray/abstractblocksparsearray.jl")
 include("abstractblocksparsearray/wrappedabstractblocksparsearray.jl")
 include("abstractblocksparsearray/abstractblocksparsematrix.jl")

NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/map.jl

@@ -36,10 +36,16 @@ function Base.getindex(a::SingleBlockView{<:Any,N}, index::Vararg{Int,N}) where
 end
 
 reblock(a) = a
+# If the blocking of the slice doesn't match the blocking of the
+# parent array, reblock according to the blocking of the parent array.
 function reblock(
-  a::SubArray{<:Any,<:Any,<:AbstractBlockSparseArray,<:Tuple{Vararg{BlockSlice}}}
+  a::SubArray{<:Any,<:Any,<:AbstractBlockSparseArray,<:Tuple{Vararg{AbstractUnitRange}}}
 )
-  return @view a.parent[map(i -> i.indices, a.indices)...]
+  # TODO: This relies on the behavior that slicing a block sparse
+  # array with a UnitRange inherits the blocking of the underlying
+  # block sparse array, we might change that default behavior
+  # so this might become something like `@blocked parent(a)[...]`.
+  return @view parent(a)[UnitRange{Int}.(parentindices(a))...]
 end
 
 function SparseArrayInterface.sparse_map!(
@@ -55,7 +61,7 @@ function SparseArrayInterface.sparse_map!(
     # TODO: Investigate why this doesn't work:
     # block_srcs = ntuple(i -> @view(a_srcs[i][_block(BI_srcs[i])]), length(a_srcs))
     block_srcs = ntuple(length(a_srcs)) do i
-      return blocks(a_srcs[i])[Int.(Tuple(_block(BI_srcs[i])))...]
+      return _blocks(a_srcs[i])[Int.(Tuple(_block(BI_srcs[i])))...]
     end
     subblock_dest = @view block_dest[BI_dest.indices...]
     subblock_srcs = ntuple(i -> @view(block_srcs[i][BI_srcs[i].indices...]), length(a_srcs))

NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/views.jl

@@ -1,7 +1,23 @@
-using BlockArrays: BlockArrays, Block, viewblock
+using BlockArrays: BlockArrays, Block, BlockIndexRange, blocklength, blocksize, viewblock
 
-function blocksparse_view(a, I...)
-  return Base.invoke(view, Tuple{AbstractArray,Vararg{Any}}, a, I...)
+# This splits `BlockIndexRange{N}` into
+# `NTuple{N,BlockIndexRange{1}}`.
+# TODO: Move to `BlockArraysExtensions`.
+to_tuple(x) = Tuple(x)
+function to_tuple(x::BlockIndexRange{N}) where {N}
+  blocks = Tuple(Block(x))
+  n = length(blocks)
+  return ntuple(dim -> blocks[dim][x.indices[dim]], n)
+end
+
+# Override the default definition of `BlockArrays.blocksize`,
+# which is incorrect for certain slices.
+function BlockArrays.blocksize(a::SubArray{<:Any,<:Any,<:BlockSparseArrayLike})
+  return blocklength.(axes(a))
+end
+function BlockArrays.blocksize(a::SubArray{<:Any,<:Any,<:BlockSparseArrayLike}, i::Int)
+  # TODO: Maybe use `blocklength(axes(a, i))` which would be a bit faster.
+  return blocksize(a)[i]
 end
 
 # These definitions circumvent some generic definitions in BlockArrays.jl:
@@ -49,23 +65,60 @@ function BlockArrays.viewblock(
   return BlockView(a, block)
 end
 
+# Specialized code for getting the view of a subblock.
 function Base.view(
-  a::SubArray{T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{BlockSliceCollection,N}}},
-  block::Block{N},
+  a::AbstractBlockSparseArray{<:Any,N}, block::BlockIndexRange{N}
+) where {N}
+  return view(a, to_tuple(block)...)
+end
+
+function Base.view(a::AbstractBlockSparseArray{<:Any,N}, I::Vararg{Block{1},N}) where {N}
+  return viewblock(a, I...)
+end
+
+to_block(I::Block{1}) = I
+to_block(I::BlockIndexRange{1}) = Block(I)
+to_blockindices(I::Block{1}) = Colon()
+to_blockindices(I::BlockIndexRange{1}) = only(I.indices)
+
+function Base.view(
+  a::AbstractBlockSparseArray{<:Any,N}, I::Vararg{Union{Block{1},BlockIndexRange{1}},N}
+) where {N}
+  return @views a[to_block.(I)...][to_blockindices.(I)...]
+end
+
+function Base.view(
+  a::SubArray{
+    T,
+    N,
+    <:AbstractBlockSparseArray{T,N},
+    <:Tuple{Vararg{Union{BlockSliceCollection,SubBlockSliceCollection},N}},
+  },
+  block::Union{Block{N},BlockIndexRange{N}},
 ) where {T,N}
   return viewblock(a, block)
 end
 function Base.view(
-  a::SubArray{T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{BlockSliceCollection,N}}},
-  block::Vararg{Block{1},N},
+  a::SubArray{
+    T,
+    N,
+    <:AbstractBlockSparseArray{T,N},
+    <:Tuple{Vararg{Union{BlockSliceCollection,SubBlockSliceCollection},N}},
+  },
+  block::Vararg{Union{Block{1},BlockIndexRange{1}},N},
 ) where {T,N}
   return viewblock(a, block...)
 end
 function BlockArrays.viewblock(
-  a::SubArray{T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{BlockSliceCollection,N}}},
-  block::Block{N},
+  a::SubArray{
+    T,
+    N,
+    <:AbstractBlockSparseArray{T,N},
+    <:Tuple{Vararg{Union{BlockSliceCollection,SubBlockSliceCollection},N}},
+  },
+  block::Union{Block{N},BlockIndexRange{N}},
 ) where {T,N}
-  return viewblock(a, Tuple(block)...)
+  return viewblock(a, to_tuple(block)...)
 end
 
 # Fixes ambiguity error with `BlockSparseArrayLike` definition.
@@ -101,5 +154,90 @@ function BlockArrays.viewblock(
   if I ∈ stored_indices(blocks(a))
     return blocks(a)[I]
   end
-  return BlockView(a, block)
+  return BlockView(parent(a), Block.(Base.reindex(parentindices(blocks(a)), Tuple(I))))
+end
+
+function to_blockindexrange(
+  a::BlockIndices{<:BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}}},
+  I::Block{1},
+)
+  # TODO: Ideally we would just use `a.blocks[I]` but that doesn't
+  # work right now.
+  return blocks(a.blocks)[Int(I)]
+end
+function to_blockindexrange(a::Base.Slice{<:BlockedOneTo{<:Integer}}, I::Block{1})
+  @assert I in only(blockaxes(a.indices))
+  return I
+end
+
+function BlockArrays.viewblock(
+  a::SubArray{
+    T,
+    N,
+    <:AbstractBlockSparseArray{T,N},
+    <:Tuple{Vararg{Union{BlockSliceCollection,SubBlockSliceCollection},N}},
+  },
+  block::Vararg{Block{1},N},
+) where {T,N}
+  brs = ntuple(dim -> to_blockindexrange(parentindices(a)[dim], block[dim]), ndims(a))
+  return @view parent(a)[brs...]
+end
+
+# TODO: Define `blocksparse_viewblock`.
+function BlockArrays.viewblock(
+  a::SubArray{
+    T,
+    N,
+    <:AbstractBlockSparseArray{T,N},
+    <:Tuple{Vararg{Union{BlockSliceCollection,SubBlockSliceCollection},N}},
+  },
+  block::Vararg{BlockIndexRange{1},N},
+) where {T,N}
+  return view(viewblock(a, Block.(block)...), map(b -> only(b.indices), block)...)
+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}}}
+}
+
+function Base.view(
+  a::SubArray{T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{BlockedSlice,N}}},
+  block::Union{Block{N},BlockIndexRange{N}},
+) where {T,N}
+  return viewblock(a, block)
+end
+function Base.view(
+  a::SubArray{T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{BlockedSlice,N}}},
+  block::Vararg{Union{Block{1},BlockIndexRange{1}},N},
+) where {T,N}
+  return viewblock(a, block...)
+end
+function BlockArrays.viewblock(
+  a::SubArray{T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{BlockedSlice,N}}},
+  block::Union{Block{N},BlockIndexRange{N}},
+) where {T,N}
+  return viewblock(a, to_tuple(block)...)
+end
+# TODO: Define `blocksparse_viewblock`.
+function BlockArrays.viewblock(
+  a::SubArray{T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{BlockedSlice,N}}},
+  I::Vararg{Block{1},N},
+) where {T,N}
+  # TODO: Use `reindex`, `to_indices`, etc.
+  brs = ntuple(ndims(a)) do dim
+    # 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])]
+  end
+  return @view parent(a)[brs...]
+end
+# TODO: Define `blocksparse_viewblock`.
+function BlockArrays.viewblock(
+  a::SubArray{T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{BlockedSlice,N}}},
+  block::Vararg{BlockIndexRange{1},N},
+) where {T,N}
+  return view(viewblock(a, Block.(block)...), map(b -> only(b.indices), block)...)
 end

NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl

@@ -33,21 +33,30 @@ function Base.to_indices(
   return blocksparse_to_indices(a, inds, I)
 end
 
-# a[[Block(1)[1:2], Block(2)[1:2]], [Block(1)[1:2], Block(2)[1:2]]]
+# a[BlockVector([Block(2), Block(1)], [2]), BlockVector([Block(2), Block(1)], [2])]
+# a[BlockedVector([Block(2), Block(1)], [2]), BlockedVector([Block(2), Block(1)], [2])]
 function Base.to_indices(
-  a::BlockSparseArrayLike, inds, I::Tuple{Vector{<:BlockIndexRange{1}},Vararg{Any}}
+  a::BlockSparseArrayLike, inds, I::Tuple{AbstractBlockVector{<:Block{1}},Vararg{Any}}
 )
-  return to_indices(a, inds, (mortar(I[1]), Base.tail(I)...))
+  return blocksparse_to_indices(a, inds, I)
 end
 
-# a[BlockVector([Block(2), Block(1)], [2]), BlockVector([Block(2), Block(1)], [2])]
-# a[BlockedVector([Block(2), Block(1)], [2]), BlockedVector([Block(2), Block(1)], [2])]
+# a[mortar([Block(1)[1:2], Block(2)[1:3]])]
 function Base.to_indices(
-  a::BlockSparseArrayLike, inds, I::Tuple{AbstractBlockVector{<:Block{1}},Vararg{Any}}
+  a::BlockSparseArrayLike,
+  inds,
+  I::Tuple{BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}},Vararg{Any}},
 )
   return blocksparse_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::BlockSparseArrayLike, inds, I::Tuple{Vector{<:BlockIndexRange{1}},Vararg{Any}}
+)
+  return to_indices(a, inds, (mortar(I[1]), Base.tail(I)...))
+end
+
 # BlockArrays `AbstractBlockArray` interface
 BlockArrays.blocks(a::BlockSparseArrayLike) = blocksparse_blocks(a)
 

NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl

@@ -40,6 +40,15 @@ function blocksparse_to_indices(a, inds, I::Tuple{Vector{<:Block{1}},Vararg{Any}
   return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...)
 end
 
+# a[mortar([Block(1)[1:2], Block(2)[1:3]]), mortar([Block(1)[1:2], Block(2)[1:3]])]
+# a[[Block(1)[1:2], Block(2)[1:3]], [Block(1)[1:2], Block(2)[1:3]]]
+function blocksparse_to_indices(
+  a, inds, I::Tuple{BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{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.
@@ -281,13 +290,13 @@ function blocksparse_blocks(a::SubArray)
   return SparseSubArrayBlocks(a)
 end
 
-_blocks(I::BlockSlice) = I.block
-_blocks(I::BlockIndices) = I.blocks
+to_blocks_indices(I::BlockSlice{<:BlockRange{1}}) = Int.(I.block)
+to_blocks_indices(I::BlockIndices{<:Vector{<:Block{1}}}) = Int.(I.blocks)
 
 function blocksparse_blocks(
   a::SubArray{<:Any,<:Any,<:Any,<:Tuple{Vararg{BlockSliceCollection}}}
 )
-  return @view blocks(parent(a))[map(I -> Int.(_blocks(I)), parentindices(a))...]
+  return @view blocks(parent(a))[map(to_blocks_indices, parentindices(a))...]
 end
 
 using BlockArrays: BlocksView

NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/views.jl

@@ -0,0 +1,3 @@
+function blocksparse_view(a, I...)
+  return Base.invoke(view, Tuple{AbstractArray,Vararg{Any}}, a, I...)
+end