diff --git a/Project.toml b/Project.toml index b2d575d6..4f5f2102 100644 --- a/Project.toml +++ b/Project.toml @@ -28,6 +28,7 @@ Adapt = "3.4" ConstructionBase = "1" DataAPI = "1" GPUArraysCore = "0.1.2" +InfiniteArrays = "0.13" StaticArrays = "1.5.6" Tables = "1" julia = "1.6" @@ -36,6 +37,7 @@ julia = "1.6" Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" +InfiniteArrays = "4858937d-0d70-526a-a4dd-2d5cb5dd786c" JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881" @@ -47,4 +49,4 @@ TypedTables = "9d95f2ec-7b3d-5a63-8d20-e2491e220bb9" WeakRefStrings = "ea10d353-3f73-51f8-a26c-33c1cb351aa5" [targets] -test = ["Adapt", "Documenter", "GPUArraysCore", "JLArrays", "LinearAlgebra", "OffsetArrays", "PooledArrays", "SparseArrays", "StaticArrays", "Test", "TypedTables", "WeakRefStrings"] \ No newline at end of file +test = ["Adapt", "Documenter", "GPUArraysCore", "InfiniteArrays", "JLArrays", "LinearAlgebra", "OffsetArrays", "PooledArrays", "SparseArrays", "StaticArrays", "Test", "TypedTables", "WeakRefStrings"] diff --git a/src/structarray.jl b/src/structarray.jl index c8231bc3..64ff17a9 100644 --- a/src/structarray.jl +++ b/src/structarray.jl @@ -628,3 +628,7 @@ function Broadcast.broadcast_unalias(dest::StructArray, src::AbstractArray) end Base.dataids(u::StructArray) = mapreduce(Base.dataids, (a, b) -> (a..., b...), values(components(u)), init=()) + +# Since all the components have the same axes, we choose the type of the first one to +# define IteratorSize for a StructArray +Base.IteratorSize(::Type{<:StructArray{<:Any,<:Any,C}}) where {C} = Base.IteratorSize(fieldtype(C, 1)) diff --git a/test/runtests.jl b/test/runtests.jl index 78318c3b..ed6f573a 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -11,6 +11,7 @@ using GPUArraysCore: backend using LinearAlgebra using Test using SparseArrays +using InfiniteArrays using Documenter: doctest if Base.VERSION >= v"1.6" && Int === Int64 @@ -1500,3 +1501,12 @@ end ss = map(x -> PS(x.a, x.b), xs) @test ss == [PS(1, 2), PS(3, nothing)] end + +@testset "IteratorSize" begin + S = StructArray{ComplexF64}((zeros(1), zeros(1))) + @test Base.IteratorSize(S) == Base.HasShape{1}() + S = StructArray{ComplexF64}((zeros(1,2), zeros(1,2))) + @test Base.IteratorSize(S) == Base.HasShape{2}() + S = StructArray{Complex{Int}}((1:∞, 1:∞)) + @test Base.IteratorSize(S) == Base.IsInfinite() +end