IteratorSize for a StructArray (#290)

* IteratorSize for a StructArray

* Whitespace

Co-authored-by: Pietro Vertechi <[email protected]>

* Alphabetical order of test targets

---------

Co-authored-by: Pietro Vertechi <[email protected]>

Author: jishnub

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"]
+test = ["Adapt", "Documenter", "GPUArraysCore", "InfiniteArrays", "JLArrays", "LinearAlgebra", "OffsetArrays", "PooledArrays", "SparseArrays", "StaticArrays", "Test", "TypedTables", "WeakRefStrings"]

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))

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