diff --git a/src/Containers/DenseAxisArray.jl b/src/Containers/DenseAxisArray.jl index 7124835bfd7..1b23ee4f057 100644 --- a/src/Containers/DenseAxisArray.jl +++ b/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. diff --git a/src/Containers/SparseAxisArray.jl b/src/Containers/SparseAxisArray.jl index 0bfb020cc85..64aebd67fdf 100644 --- a/src/Containers/SparseAxisArray.jl +++ b/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 # ################ diff --git a/test/Containers/test_DenseAxisArray.jl b/test/Containers/test_DenseAxisArray.jl index 546a16169f1..fd780720a35 100644 --- a/test/Containers/test_DenseAxisArray.jl +++ b/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 diff --git a/test/Containers/test_SparseAxisArray.jl b/test/Containers/test_SparseAxisArray.jl index 914f71b0e58..94a30762b33 100644 --- a/test/Containers/test_SparseAxisArray.jl +++ b/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