Add missing BlockUnitRange functionality

Author: mtfishman

src/BlockArraysExtensions/blockrange.jl

@@ -8,13 +8,21 @@ end
 function blockrange(eachblockaxis)
   return BlockUnitRange(blockedrange(length.(eachblockaxis)), eachblockaxis)
 end
+function blockrange(first::Integer, eachblockaxis)
+  return BlockUnitRange(blockedrange(first, length.(eachblockaxis)), eachblockaxis)
+end
 Base.first(r::BlockUnitRange) = first(r.r)
 Base.last(r::BlockUnitRange) = last(r.r)
 BlockArrays.blocklasts(r::BlockUnitRange) = blocklasts(r.r)
 eachblockaxis(r::BlockUnitRange) = r.eachblockaxis
 function Base.getindex(r::BlockUnitRange, I::Block{1})
   return eachblockaxis(r)[Int(I)] .+ (first(r.r[I]) - 1)
 end
+function Base.getindex(r::BlockUnitRange, I::BlockRange{1})
+  return blockrange(first(r), eachblockaxis(r)[Int.(I)])
+end
+Base.axes(r::BlockUnitRange) = (blockrange(eachblockaxis(r)),)
+Base.axes1(r::BlockUnitRange) = blockrange(eachblockaxis(r))
 
 using BlockArrays: BlockedOneTo
 const BlockOneTo{T<:Integer,B,CS,R<:BlockedOneTo{T,CS}} = BlockUnitRange{T,B,CS,R}

src/abstractblocksparsearray/abstractblocksparsearray.jl

@@ -87,18 +87,18 @@ function Base.setindex!(
       ),
     )
   end
-  if !isconcretetype(blocktype(a))
+  if isstored(a, I...)
+    # This writes into existing blocks, or constructs blocks
+    # using the axes.
+    aI = @view! a[I...]
+    aI .= value
+  else
     # Custom `_convert` works around the issue that
     # `convert(::Type{<:Diagonal}, ::AbstractMatrix)` isnt' defined
     # in Julia v1.10 (https://github.com/JuliaLang/julia/pull/48895,
     # https://github.com/JuliaLang/julia/pull/52487).
-    # TODO: Delete once we drop support for Julia v1.10.
+    # TODO: Delete `_convert` once we drop support for Julia v1.10.
     blocks(a)[Int.(I)...] = _convert(blocktype(a), value)
-  else
-    # This writes into existing blocks, or constructs blocks
-    # using the axes.
-    aI = @view! a[I...]
-    aI .= value
   end
   return a
 end

src/blocksparsearray/blocksparsearray.jl

@@ -262,11 +262,14 @@ function blocksparsezeros(::BlockType{A}, axes...) where {A<:AbstractArray}
   # to make a bit more generic.
   return BlockSparseArray{eltype(A),ndims(A),A}(undef, axes...)
 end
-function blocksparse(d::Dict{<:Block,<:AbstractArray}, ax::Tuple)
-  blockaxtype = Tuple{map(eltype ∘ eachblockaxis, ax)...}
-  # TODO: Catch if inference fails and use `valtype(d)` instead.
-  blockt = Base.promote_op(similar, Type{valtype(d)}, blockaxtype)
-  a = blocksparsezeros(BlockType(blockt), ax)
+function blocksparse(d::Dict, ax::Tuple)
+  a = blocksparsezeros(BlockType(valtype(d)), ax...)
+
+  ## blockaxtype = Tuple{map(eltype ∘ eachblockaxis, ax)...}
+  ## # TODO: Catch if inference fails and use `valtype(d)` instead.
+  ## blockt = Base.promote_op(similar, Type{valtype(d)}, blockaxtype)
+  ## a = blocksparsezeros(BlockType(blockt), ax)
+
   for I in eachindex(d)
     a[I] = d[I]
   end

src/blocksparsearrayinterface/getunstoredblock.jl

@@ -18,10 +18,17 @@ function Base.AbstractArray{A}(a::ZeroBlocks{N}) where {N,A}
 end
 
 @inline function Base.getindex(a::ZeroBlocks{N,A}, I::Vararg{Int,N}) where {N,A}
+
+  #foreach(display, a.parentaxes)
+  #foreach(display ∘ eachblockaxis, a.parentaxes)
+
   # TODO: Use `BlockArrays.eachblockaxes`.
   ax = ntuple(N) do d
     return eachblockaxis(a.parentaxes[d])[I[d]]
   end
+
+  #@show ax
+
   !isconcretetype(A) && return zero!(similar(Array{eltype(A),N}, ax))
   return zero!(similar(A, ax))
 end