[WIP] Generalize blockwise slicing

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "BlockSparseArrays"
 uuid = "2c9a651f-6452-4ace-a6ac-809f4280fbb4"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
-version = "0.7.20"
+version = "0.7.21"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/BlockArraysExtensions/BlockArraysExtensions.jl

@@ -392,6 +392,13 @@ function blockrange(
   return map(Block, blocks(r))
 end
 
+function blockrange(
+  axis::AbstractUnitRange,
+  r::BlockVector{<:GenericBlockIndex{1},<:AbstractVector{<:BlockIndexVector}},
+)
+  return map(Block, blocks(r))
+end
+
 function blockrange(axis::AbstractUnitRange, r)
   return error("Slicing not implemented for range of type `$(typeof(r))`.")
 end

src/BlockArraysExtensions/blockedunitrange.jl

@@ -244,6 +244,15 @@ BlockArrays.blockindex(b::GenericBlockIndex{1}) = b.α[1]
 function GenericBlockIndex(indcs::Tuple{Vararg{GenericBlockIndex{1},N}}) where {N}
   GenericBlockIndex(block.(indcs), blockindex.(indcs))
 end
+
+function Base.checkindex(
+  ::Type{Bool}, axis::AbstractBlockedUnitRange, ind::GenericBlockIndex{1}
+)
+  return checkindex(Bool, axis, block(ind)) &&
+         checkbounds(Bool, axis[block(ind)], blockindex(ind))
+end
+Base.to_index(i::GenericBlockIndex) = i
+
 function print_tuple_elements(io::IO, @nospecialize(t))
   if !isempty(t)
     print(io, t[1])
@@ -261,6 +270,13 @@ function Base.show(io::IO, B::GenericBlockIndex)
   return nothing
 end
 
+# https://github.com/JuliaArrays/BlockArrays.jl/blob/v1.6.3/src/views.jl#L31-L32
+_maybetail(::Tuple{}) = ()
+_maybetail(t::Tuple) = Base.tail(t)
+@inline function Base.to_indices(A, inds, I::Tuple{GenericBlockIndex{1},Vararg{Any}})
+  return (inds[1][I[1]], to_indices(A, _maybetail(inds), Base.tail(I))...)
+end
+
 using Base: @propagate_inbounds
 @propagate_inbounds function Base.getindex(b::AbstractVector, K::GenericBlockIndex{1})
   return b[Block(K.I[1])][K.α[1]]