Fix a bug in blocktype

Project.toml

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

src/blocksparsearrayinterface/blocksparsearrayinterface.jl

@@ -41,31 +41,37 @@
   return storedvalues(blocks(a))
 end
 
-# TODO: Generalize this, this catches simple cases
-# where the more general definition isn't specific enough.
-blocktype(a::Array) = typeof(a)
-# TODO: Maybe unwrap SubArrays?
+function blockstype(a::AbstractArray)
+  return typeof(blocks(a))
+end
 function blocktype(a::AbstractArray)
-  # TODO: Unfortunately, this doesn't always give
-  # a concrete type, even when it could be concrete, i.e.
-  #=
-  ```julia
-  julia> eltype(blocks(BlockArray(randn(2, 2), [1, 1], [1, 1])))
-  Matrix{Float64} (alias for Array{Float64, 2})
-
-  julia> eltype(blocks(BlockedArray(randn(2, 2), [1, 1], [1, 1])))
-  AbstractMatrix{Float64} (alias for AbstractArray{Float64, 2})
-
-  julia> eltype(blocks(randn(2, 2)))
-  AbstractMatrix{Float64} (alias for AbstractArray{Float64, 2})
-  ```
-  =#
   if isempty(blocks(a))
+    # TODO: Unfortunately, this doesn't always give
+    # a concrete type, even when it could be concrete, i.e.
+    #=
+    ```julia
+    julia> eltype(blocks(BlockArray(randn(2, 2), [1, 1], [1, 1])))
+    Matrix{Float64} (alias for Array{Float64, 2})
+
+    julia> eltype(blocks(BlockedArray(randn(2, 2), [1, 1], [1, 1])))
+    AbstractMatrix{Float64} (alias for AbstractArray{Float64, 2})
+
+    julia> eltype(blocks(randn(2, 2)))
+    AbstractMatrix{Float64} (alias for AbstractArray{Float64, 2})
+    ```
+    =#
+    # That is an issue in BlockArrays.jl that we should address.
     return eltype(blocks(a))
   end
-  return eltype(first(blocks(a)))
+  return mapreduce(typeof, promote_type, blocks(a))
 end
 
+using BlockArrays: BlockArray
+blockstype(::Type{<:BlockArray{<:Any,<:Any,B}}) where {B} = B
+blockstype(a::BlockArray) = blockstype(typeof(a))
+blocktype(arraytype::Type{<:BlockArray}) = eltype(blockstype(arraytype))
+blocktype(a::BlockArray) = eltype(blocks(a))
+
 abstract type AbstractBlockSparseArrayInterface <: AbstractSparseArrayInterface end
 
 # TODO: Also support specifying the `blocktype` along with the `eltype`.
@@ -78,8 +84,6 @@
 @interface ::AbstractBlockSparseArrayInterface BlockArrays.blocks(a::AbstractArray) =
   error("Not implemented")
 
-blockstype(a::AbstractArray) = blockstype(typeof(a))
-
 @interface ::AbstractBlockSparseArrayInterface function Base.getindex(
   a::AbstractArray{<:Any,N}, I::Vararg{Int,N}
 ) where {N}

test/test_basics.jl

@@ -1,13 +1,16 @@
 using Adapt: adapt
 using ArrayLayouts: zero!
 using BlockArrays:
+  BlockArrays,
   Block,
+  BlockArray,
   BlockIndexRange,
   BlockRange,
   BlockSlice,
   BlockVector,
   BlockedOneTo,
   BlockedUnitRange,
+  BlockedArray,
   BlockedVector,
   blockedrange,
   blocklength,
@@ -132,6 +135,35 @@ arrayts = (Array, JLArray)
       end
     end
   end
+  @testset "blockstype, blocktype" begin
+    a = arrayt(randn(elt, 2, 2))
+    @test blockstype(a) <: BlockArrays.BlocksView{elt,2}
+    # TODO: This is difficult to determine just from type information.
+    @test_broken blockstype(typeof(a)) <: BlockArrays.BlocksView{elt,2}
+    @test blocktype(a) <: SubArray{elt,2,arrayt{elt,2}}
+    # TODO: This is difficult to determine just from type information.
+    @test_broken blocktype(typeof(a)) <: SubArray{elt,2,arrayt{elt,2}}
+
+    a = BlockSparseMatrix{elt,arrayt{elt,2}}([1, 1], [1, 1])
+    @test blockstype(a) <: SparseMatrixDOK{arrayt{elt,2}}
+    @test blockstype(typeof(a)) <: SparseMatrixDOK{arrayt{elt,2}}
+    @test blocktype(a) <: arrayt{elt,2}
+    @test blocktype(typeof(a)) <: arrayt{elt,2}
+
+    a = BlockArray(arrayt(randn(elt, (2, 2))), [1, 1], [1, 1])
+    @test blockstype(a) === Matrix{arrayt{elt,2}}
+    @test blockstype(typeof(a)) === Matrix{arrayt{elt,2}}
+    @test blocktype(a) <: arrayt{elt,2}
+    @test blocktype(typeof(a)) <: arrayt{elt,2}
+
+    a = BlockedArray(arrayt(randn(elt, 2, 2)), [1, 1], [1, 1])
+    @test blockstype(a) <: BlockArrays.BlocksView{elt,2}
+    # TODO: This is difficult to determine just from type information.
+    @test_broken blockstype(typeof(a)) <: BlockArrays.BlocksView{elt,2}
+    @test blocktype(a) <: SubArray{elt,2,arrayt{elt,2}}
+    # TODO: This is difficult to determine just from type information.
+    @test_broken blocktype(typeof(a)) <: SubArray{elt,2,arrayt{elt,2}}
+  end
   @testset "Basics" begin
     a = dev(BlockSparseArray{elt}([2, 3], [2, 3]))
     @allowscalar @test a == dev(