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,41 @@
   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
+
+#=
+Ideally this would just be defined as `eltype(blockstype(a))`.
+However, BlockArrays.jl doesn't make `eltype(blocks(a))` concrete
+even when it could be
+(https://github.com/JuliaArrays/BlockArrays.jl/blob/v1.4.0/src/blocks.jl#L71-L74):
+```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})
+```
+Also note the current definition doesn't handle the limit
+when `blocks(a)` is empty
+=#
 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))
-    return eltype(blocks(a))
+    error("`blocktype` can't be determined if `isempty(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 +88,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/Project.toml

@@ -18,3 +18,4 @@ Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
 SymmetrySectors = "f8a8ad64-adbc-4fce-92f7-ffe2bb36a86e"
 TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a"

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,
@@ -34,6 +37,7 @@ using LinearAlgebra: Adjoint, Transpose, dot, mul!, norm
 using SparseArraysBase: SparseArrayDOK, SparseMatrixDOK, SparseVectorDOK, storedlength
 using TensorAlgebra: contract
 using Test: @test, @test_broken, @test_throws, @testset, @inferred
+using TestExtras: @constinferred
 include("TestBlockSparseArraysUtils.jl")
 
 arrayts = (Array, JLArray)
@@ -132,6 +136,35 @@ arrayts = (Array, JLArray)
       end
     end
   end
+  @testset "blockstype, blocktype" begin
+    a = arrayt(randn(elt, 2, 2))
+    @test (@constinferred 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 (@constinferred 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 (@constinferred blockstype(a)) <: SparseMatrixDOK{arrayt{elt,2}}
+    @test (@constinferred blockstype(typeof(a))) <: SparseMatrixDOK{arrayt{elt,2}}
+    @test (@constinferred blocktype(a)) <: arrayt{elt,2}
+    @test (@constinferred blocktype(typeof(a))) <: arrayt{elt,2}
+
+    a = BlockArray(arrayt(randn(elt, (2, 2))), [1, 1], [1, 1])
+    @test (@constinferred blockstype(a)) === Matrix{arrayt{elt,2}}
+    @test (@constinferred blockstype(typeof(a))) === Matrix{arrayt{elt,2}}
+    @test (@constinferred blocktype(a)) <: arrayt{elt,2}
+    @test (@constinferred blocktype(typeof(a))) <: arrayt{elt,2}
+
+    a = BlockedArray(arrayt(randn(elt, 2, 2)), [1, 1], [1, 1])
+    @test (@constinferred 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 (@constinferred 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(