From 281b496c760fa3fd6e7442f43e03e9c1491f91a6 Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Mon, 18 Aug 2025 20:08:35 -0400
Subject: [PATCH] More general ZeroBlocks definition

---
 Project.toml                                  |  2 +-
 .../KroneckerArraysBlockSparseArraysExt.jl    | 64 ++++++-------------
 2 files changed, 22 insertions(+), 44 deletions(-)

diff --git a/Project.toml b/Project.toml
index 517d91c..8a5f85c 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "KroneckerArrays"
 uuid = "05d0b138-81bc-4ff7-84be-08becefb1ccc"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
-version = "0.1.29"
+version = "0.1.30"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
diff --git a/ext/KroneckerArraysBlockSparseArraysExt/KroneckerArraysBlockSparseArraysExt.jl b/ext/KroneckerArraysBlockSparseArraysExt/KroneckerArraysBlockSparseArraysExt.jl
index 624b269..a28d98d 100644
--- a/ext/KroneckerArraysBlockSparseArraysExt/KroneckerArraysBlockSparseArraysExt.jl
+++ b/ext/KroneckerArraysBlockSparseArraysExt/KroneckerArraysBlockSparseArraysExt.jl
@@ -25,18 +25,8 @@ end
 
 using BlockArrays: AbstractBlockedUnitRange
 using BlockSparseArrays: Block, ZeroBlocks, eachblockaxis, mortar_axis
-using DerivableInterfaces: zero!
-using FillArrays: Eye
-using KroneckerArrays:
-  KroneckerArrays,
-  EyeEye,
-  EyeKronecker,
-  KroneckerEye,
-  KroneckerMatrix,
-  ⊗,
-  arg1,
-  arg2,
-  _similar
+using KroneckerArrays: KroneckerArrays, KroneckerArray, ⊗, arg1, arg2, _similar
+using BlockSparseArrays.TypeParameterAccessors: unwrap_array_type
 
 function KroneckerArrays.arg1(r::AbstractBlockedUnitRange)
   return mortar_axis(arg1.(eachblockaxis(r)))
@@ -58,41 +48,29 @@ end
 
 ## TODO: Is this needed?
 function Base.getindex(
-  a::ZeroBlocks{2,KroneckerMatrix{T,A,B}}, I::Vararg{Int,2}
-) where {T,A<:AbstractMatrix{T},B<:AbstractMatrix{T}}
+  a::ZeroBlocks{N,KroneckerArray{T,N,A,B}}, I::Vararg{Int,N}
+) where {T,N,A<:AbstractArray{T,N},B<:AbstractArray{T,N}}
   ax_a1 = map(arg1, a.parentaxes)
-  a1 = ZeroBlocks{2,A}(ax_a1)[I...]
   ax_a2 = map(arg2, a.parentaxes)
-  a2 = ZeroBlocks{2,B}(ax_a2)[I...]
-  return a1 ⊗ a2
-end
-function Base.getindex(
-  a::ZeroBlocks{2,EyeKronecker{T,A,B}}, I::Vararg{Int,2}
-) where {T,A<:Eye{T},B<:AbstractMatrix{T}}
-  block_ax_a1 = arg1.(block_axes(a.parentaxes, Block(I)))
-  a1 = _similar(A, block_ax_a1)
-
-  ax_a2 = arg2.(a.parentaxes)
-  a2 = ZeroBlocks{2,B}(ax_a2)[I...]
-
-  return a1 ⊗ a2
-end
-function Base.getindex(
-  a::ZeroBlocks{2,KroneckerEye{T,A,B}}, I::Vararg{Int,2}
-) where {T,A<:AbstractMatrix{T},B<:Eye{T}}
-  ax_a1 = arg1.(a.parentaxes)
-  a1 = ZeroBlocks{2,A}(ax_a1)[I...]
-
-  block_ax_a2 = arg2.(block_axes(a.parentaxes, Block(I)))
-  a2 = _similar(B, block_ax_a2)
-
+  # TODO: Instead of mutability, maybe have a trait like
+  # `isstructural` or `isdata`.
+  ismut1 = ismutabletype(unwrap_array_type(A))
+  ismut2 = ismutabletype(unwrap_array_type(B))
+  (ismut1 || ismut2) || error("Can't get zero block.")
+  a1 = if ismut1
+    ZeroBlocks{N,A}(ax_a1)[I...]
+  else
+    block_ax_a1 = arg1.(block_axes(a.parentaxes, Block(I)))
+    _similar(A, block_ax_a1)
+  end
+  a2 = if ismut2
+    ZeroBlocks{N,B}(ax_a2)[I...]
+  else
+    block_ax_a2 = arg2.(block_axes(a.parentaxes, Block(I)))
+    a2 = _similar(B, block_ax_a2)
+  end
   return a1 ⊗ a2
 end
-function Base.getindex(
-  a::ZeroBlocks{2,EyeEye{T,A,B}}, I::Vararg{Int,2}
-) where {T,A<:Eye{T},B<:Eye{T}}
-  return error("Not implemented.")
-end
 
 using BlockSparseArrays: BlockSparseArrays
 using KroneckerArrays: KroneckerArrays, KroneckerVector