[Containers] fix multi-arg eachindex (#3587)

src/Containers/DenseAxisArray.jl

@@ -343,7 +343,28 @@ if VERSION >= v"1.9.0-DEV"
     end
 end
 
-Base.eachindex(A::DenseAxisArray) = CartesianIndices(size(A.data))
+Base.eachindex(A::DenseAxisArray) = eachindex(IndexStyle(A), A)
+
+function Base.eachindex(::IndexCartesian, A::DenseAxisArray)
+    return CartesianIndices(size(A.data))
+end
+
+function Base.eachindex(
+    ::IndexCartesian,
+    A::DenseAxisArray,
+    B::DenseAxisArray...,
+)
+    ret = eachindex(A)
+    for b in B
+        if eachindex(b) != ret
+            err = DimensionMismatch(
+                "incompatible dimensions in eachindex. Got $(eachindex.((A, B...)))",
+            )
+            throw(err)
+        end
+    end
+    return ret
+end
 
 # Use recursion over tuples to ensure the return-type of functions like
 # `Base.to_index` are type-stable.

src/Containers/SparseAxisArray.jl

@@ -219,6 +219,23 @@ end
 
 Base.eachindex(d::SparseAxisArray) = keys(d.data)
 
+function Base.eachindex(
+    ::IndexCartesian,
+    A::SparseAxisArray,
+    B::SparseAxisArray...,
+)
+    ret = eachindex(A)
+    for b in B
+        if eachindex(b) != ret
+            err = DimensionMismatch(
+                "incompatible dimensions in eachindex. Got $(eachindex.((A, B...)))",
+            )
+            throw(err)
+        end
+    end
+    return ret
+end
+
 ################
 # Broadcasting #
 ################

test/Containers/test_DenseAxisArray.jl

@@ -862,4 +862,16 @@ function test_sum_dims()
     return
 end
 
+function test_multi_arg_eachindex()
+    Containers.@container(x[i = 2:3], i)
+    Containers.@container(y[i = 2:3], i)
+    Containers.@container(z[i = 2:4, j = 1:2], i + j)
+    @test eachindex(x) == CartesianIndices((2,))
+    @test eachindex(y) == CartesianIndices((2,))
+    @test eachindex(z) == CartesianIndices((3, 2))
+    @test eachindex(x, y) == CartesianIndices((2,))
+    @test_throws DimensionMismatch eachindex(x, z)
+    return
+end
+
 end  # module

test/Containers/test_SparseAxisArray.jl

@@ -343,4 +343,20 @@ function test_containers_sparseaxisarray_kwarg_setindex()
     return
 end
 
+function test_multi_arg_eachindex()
+    Containers.@container(x[i = 2:3], i, container = SparseAxisArray)
+    Containers.@container(y[i = 2:3], i, container = SparseAxisArray)
+    Containers.@container(
+        z[i = 2:4, j = 1:2],
+        i + j,
+        container = SparseAxisArray,
+    )
+    @test eachindex(x) == keys(x.data)
+    @test eachindex(y) == keys(y.data)
+    @test eachindex(z) == keys(z.data)
+    @test eachindex(x, y) == eachindex(x)
+    @test_throws DimensionMismatch eachindex(x, z)
+    return
+end
+
 end  # module