From f52314467140a103a7c5fd5c4f0ba553cf16a85a Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Thu, 16 Jan 2025 13:25:41 -0500
Subject: [PATCH] Fix tests

---
 .../abstractblocksparsearray.jl               |  6 +++++
 src/abstractblocksparsearray/map.jl           |  3 ++-
 .../blocksparsearrayinterface.jl              |  5 ++---
 test/test_basics.jl                           | 22 +++++++++----------
 4 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/src/abstractblocksparsearray/abstractblocksparsearray.jl b/src/abstractblocksparsearray/abstractblocksparsearray.jl
index c3297c1..b03d262 100644
--- a/src/abstractblocksparsearray/abstractblocksparsearray.jl
+++ b/src/abstractblocksparsearray/abstractblocksparsearray.jl
@@ -57,6 +57,12 @@ function Base.setindex!(a::AbstractBlockSparseArray{<:Any,0}, value)
   return a
 end
 
+# Catch zero-dimensional case to avoid scalar indexing.
+function Base.setindex!(a::AbstractBlockSparseArray{<:Any,0}, value, ::Block{0})
+  blocks(a)[] = value
+  return a
+end
+
 function Base.setindex!(
   a::AbstractBlockSparseArray{<:Any,N}, value, I::Vararg{Block{1},N}
 ) where {N}
diff --git a/src/abstractblocksparsearray/map.jl b/src/abstractblocksparsearray/map.jl
index 52aec40..fd5cec7 100644
--- a/src/abstractblocksparsearray/map.jl
+++ b/src/abstractblocksparsearray/map.jl
@@ -1,6 +1,7 @@
 using ArrayLayouts: LayoutArray
 using BlockArrays: blockisequal
 using DerivableInterfaces: @interface, AbstractArrayInterface, interface
+using GPUArraysCore: @allowscalar
 using LinearAlgebra: Adjoint, Transpose
 using SparseArraysBase: SparseArraysBase, SparseArrayStyle
 
@@ -55,7 +56,7 @@ function map_zero_dim! end
 @interface ::AbstractArrayInterface function map_zero_dim!(
   f, a_dest::AbstractArray, a_srcs::AbstractArray...
 )
-  a_dest[] = f.(map(a_src -> a_src[], a_srcs)...)
+  @allowscalar a_dest[] = f.(map(a_src -> a_src[], a_srcs)...)
   return a_dest
 end
 
diff --git a/src/blocksparsearrayinterface/blocksparsearrayinterface.jl b/src/blocksparsearrayinterface/blocksparsearrayinterface.jl
index b93f275..e52c4ac 100644
--- a/src/blocksparsearrayinterface/blocksparsearrayinterface.jl
+++ b/src/blocksparsearrayinterface/blocksparsearrayinterface.jl
@@ -15,7 +15,6 @@ using BlockArrays:
   blocks,
   findblockindex
 using DerivableInterfaces: DerivableInterfaces, @interface, DefaultArrayInterface
-using GPUArraysCore: @allowscalar
 using LinearAlgebra: Adjoint, Transpose
 using SparseArraysBase:
   AbstractSparseArrayInterface,
@@ -66,7 +65,7 @@ end
 @interface ::AbstractBlockSparseArrayInterface function Base.getindex(
   a::AbstractArray{<:Any,0}
 )
-  return @allowscalar a[Block()[]]
+  return a[Block()[]]
 end
 
 # a[1:2, 1:2]
@@ -157,7 +156,7 @@ end
 )
   a_b = blocks(a)[]
   # `value[]` handles scalars and 0-dimensional arrays.
-  @allowscalar a_b[] = value[]
+  a_b[] = value[]
   # Set the block, required if it is structurally zero.
   blocks(a)[] = a_b
   return a
diff --git a/test/test_basics.jl b/test/test_basics.jl
index ca4ac6f..1754428 100644
--- a/test/test_basics.jl
+++ b/test/test_basics.jl
@@ -183,17 +183,17 @@ arrayts = (Array, JLArray)
       @test blocksize(a) == ()
       @test blocksizes(a) == fill(())
       @test iszero(blockstoredlength(a))
-      @test iszero(a[])
-      @test iszero([CartesianIndex()])
+      @test iszero(@allowscalar(a[]))
+      @test iszero(@allowscalar(a[CartesianIndex()]))
       @test a[Block()] == dev(fill(0))
-      @test iszero(a[Block()][])
-      @test iszero(a[Block()[]])
+      @test iszero(@allowscalar(a[Block()][]))
+      @test iszero(@allowscalar(a[Block()[]]))
       @test Array(a) isa Array{elt,0}
       @test Array(a) == fill(0)
       for b in (
-        (b = copy(a);  b[] = 2; b),
-        (b = copy(a); b[CartesianIndex()] = 2; b),
-        (b = copy(a); b[CartesianIndex()] = 2; b),
+        (b = copy(a); @allowscalar(b[] = 2); b),
+        (b = copy(a); @allowscalar(b[CartesianIndex()] = 2); b),
+        (b = copy(a); @allowscalar(b[Block()[]] = 2); b),
         # Regression test for https://github.com/ITensor/BlockSparseArrays.jl/issues/27.
         (b = copy(a); b[Block()] = dev(fill(2)); b),
       )
@@ -202,11 +202,11 @@ arrayts = (Array, JLArray)
         @test blocksize(b) == ()
         @test blocksizes(b) == fill(())
         @test isone(blockstoredlength(b))
-        @test b[] == 2
-        @test b[CartesianIndex()] == 2
+        @test @allowscalar(b[]) == 2
+        @test @allowscalar(b[CartesianIndex()]) == 2
         @test b[Block()] == dev(fill(2))
-        @test b[Block()][] == 2
-        @test b[Block()[]] == 2
+        @test @allowscalar(b[Block()][]) == 2
+        @test @allowscalar(b[Block()[]]) == 2
         @test Array(b) isa Array{elt,0}
         @test Array(b) == fill(2)
       end