Unblocked slicing

Author: mtfishman

src/BlockArraysExtensions/BlockArraysExtensions.jl

@@ -30,13 +30,14 @@ using SparseArraysBase:
 
 # A return type for `blocks(array)` when `array` isn't blocked.
 # Represents a vector with just that single block.
-struct SingleBlockView{T,N,Array<:AbstractArray{T,N}} <: AbstractArray{T,N}
+struct SingleBlockView{N,Array<:AbstractArray{<:Any,N}} <: AbstractArray{Array,N}
   array::Array
 end
 Base.parent(a::SingleBlockView) = a.array
+Base.size(a::SingleBlockView) = ntuple(Returns(1), ndims(a))
 blocks_maybe_single(a) = blocks(a)
 blocks_maybe_single(a::Array) = SingleBlockView(a)
-function Base.getindex(a::SingleBlockView{<:Any,N}, index::Vararg{Int,N}) where {N}
+function Base.getindex(a::SingleBlockView{N}, index::Vararg{Int,N}) where {N}
   @assert all(isone, index)
   return parent(a)
 end
@@ -357,7 +358,11 @@ function blockrange(axis::AbstractUnitRange, r::Base.Slice)
 end
 
 function blockrange(axis::AbstractUnitRange, r::NonBlockedVector)
-  return Block(1):Block(1)
+  return Block.(Base.OneTo(1))
+end
+
+function blockrange(axis::AbstractUnitRange, r::AbstractVector{<:Integer})
+  return Block.(Base.OneTo(1))
 end
 
 function blockrange(axis::AbstractUnitRange, r)

src/BlockSparseArrays.jl

@@ -27,6 +27,7 @@ include("abstractblocksparsearray/abstractblocksparsearray.jl")
 include("abstractblocksparsearray/abstractblocksparsematrix.jl")
 include("abstractblocksparsearray/abstractblocksparsevector.jl")
 include("abstractblocksparsearray/wrappedabstractblocksparsearray.jl")
+include("abstractblocksparsearray/unblockedsubarray.jl")
 include("abstractblocksparsearray/views.jl")
 include("abstractblocksparsearray/arraylayouts.jl")
 include("abstractblocksparsearray/sparsearrayinterface.jl")

src/abstractblocksparsearray/arraylayouts.jl

@@ -43,11 +43,20 @@ function ArrayLayouts.sub_materialize(layout::BlockLayout{<:SparseLayout}, a, ax
   return a_dest
 end
 
+function _similar(arraytype::Type{<:AbstractArray}, size::Tuple)
+  return similar(arraytype, size)
+end
+function _similar(
+  ::Type{<:SubArray{<:Any,<:Any,<:ArrayType}}, size::Tuple
+) where {ArrayType}
+  return similar(ArrayType, size)
+end
+
 # Materialize a SubArray view.
 function ArrayLayouts.sub_materialize(
   layout::BlockLayout{<:SparseLayout}, a, axes::Tuple{Vararg{Base.OneTo}}
 )
-  a_dest = blocktype(a)(undef, length.(axes))
+  a_dest = _similar(blocktype(a), length.(axes))
   a_dest .= a
   return a_dest
 end

src/abstractblocksparsearray/unblockedsubarray.jl

@@ -0,0 +1,44 @@
+using ArrayLayouts: ArrayLayouts, MemoryLayout
+using Base.Broadcast: Broadcast, BroadcastStyle
+using BlockArrays: BlockArrays
+using TypeParameterAccessors: TypeParameterAccessors, parenttype, similartype
+
+const UnblockedSubArray{T,N} = SubArray{
+  T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{Vector{<:Integer}}}
+}
+
+function BlockArrays.blocks(a::UnblockedSubArray)
+  return SingleBlockView(a)
+end
+
+function DerivableInterfaces.interface(arraytype::Type{<:UnblockedSubArray})
+  return interface(blocktype(parenttype(arraytype)))
+end
+
+function ArrayLayouts.MemoryLayout(arraytype::Type{<:UnblockedSubArray})
+  return MemoryLayout(blocktype(parenttype(arraytype)))
+end
+
+function Broadcast.BroadcastStyle(arraytype::Type{<:UnblockedSubArray})
+  return BroadcastStyle(blocktype(parenttype(arraytype)))
+end
+
+function TypeParameterAccessors.similartype(arraytype::Type{<:UnblockedSubArray}, elt::Type)
+  return similartype(blocktype(parenttype(arraytype)), elt)
+end
+
+function Base.map!(
+  f, a_dest::AbstractArray, a_src1::UnblockedSubArray, a_src_rest::UnblockedSubArray...
+)
+  return invoke(
+    map!,
+    Tuple{Any,AbstractArray,AbstractArray,Vararg{AbstractArray}},
+    f,
+    a_dest,
+    a_src1,
+    a_src_rest...,
+  )
+end
+function Base.iszero(a::UnblockedSubArray)
+  return invoke(iszero, Tuple{AbstractArray}, a)
+end

src/blocksparsearrayinterface/blocksparsearrayinterface.jl

@@ -400,7 +400,7 @@ function Base.isassigned(a::SparseSubArrayBlocks{<:Any,N}, I::Vararg{Int,N}) whe
   # TODO: Implement this properly.
   return true
 end
-function SparseArraysBase.eachstoredindex(a::SparseSubArrayBlocks)
+function SparseArraysBase.eachstoredindex(::IndexCartesian, a::SparseSubArrayBlocks)
   return eachstoredindex(view(blocks(parent(a.array)), blockrange(a)...))
 end
 # TODO: Either make this the generic interface or define