diff --git a/Project.toml b/Project.toml index f67e5041..66739491 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "StructArrays" uuid = "09ab397b-f2b6-538f-b94a-2f83cf4a842a" -version = "0.6.4" +version = "0.6.5" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" diff --git a/src/structarray.jl b/src/structarray.jl index 650a6b44..114c514e 100644 --- a/src/structarray.jl +++ b/src/structarray.jl @@ -445,19 +445,25 @@ end # broadcast import Base.Broadcast: BroadcastStyle, ArrayStyle, AbstractArrayStyle, Broadcasted, DefaultArrayStyle -struct StructArrayStyle{Style} <: AbstractArrayStyle{Any} end +struct StructArrayStyle{S, N} <: AbstractArrayStyle{N} end -@inline combine_style_types(::Type{A}, args...) where A<:AbstractArray = +# Here we define the dimension tracking behavior of StructArrayStyle +function StructArrayStyle{S, M}(::Val{N}) where {S, M, N} + T = S <: AbstractArrayStyle{M} ? typeof(S(Val(N))) : S + return StructArrayStyle{T, N}() +end + +@inline combine_style_types(::Type{A}, args...) where {A<:AbstractArray} = combine_style_types(BroadcastStyle(A), args...) -@inline combine_style_types(s::BroadcastStyle, ::Type{A}, args...) where A<:AbstractArray = +@inline combine_style_types(s::BroadcastStyle, ::Type{A}, args...) where {A<:AbstractArray} = combine_style_types(Broadcast.result_style(s, BroadcastStyle(A)), args...) combine_style_types(s::BroadcastStyle) = s -Base.@pure cst(::Type{SA}) where SA = combine_style_types(array_types(SA).parameters...) +Base.@pure cst(::Type{SA}) where {SA} = combine_style_types(array_types(SA).parameters...) -BroadcastStyle(::Type{SA}) where SA<:StructArray = StructArrayStyle{typeof(cst(SA))}() +BroadcastStyle(::Type{SA}) where {SA<:StructArray} = StructArrayStyle{typeof(cst(SA)), ndims(SA)}() -Base.similar(bc::Broadcasted{StructArrayStyle{S}}, ::Type{ElType}) where {S<:DefaultArrayStyle,N,ElType} = +Base.similar(bc::Broadcasted{StructArrayStyle{S, N}}, ::Type{ElType}) where {S<:DefaultArrayStyle, N, ElType} = isstructtype(ElType) ? similar(StructArray{ElType}, axes(bc)) : similar(Array{ElType}, axes(bc)) # for aliasing analysis during broadcast diff --git a/test/runtests.jl b/test/runtests.jl index 1fe7be57..2f7beddd 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -926,8 +926,25 @@ Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{MyArray}}, ::Type{El # used inside of broadcast but we also test it here explicitly @test isa(@inferred(Base.dataids(s)), NTuple{N, UInt} where {N}) - s = StructArray{ComplexF64}((MyArray(rand(2,2)), MyArray(rand(2,2)))) + s = StructArray{ComplexF64}((MyArray(rand(2)), MyArray(rand(2)))) @test_throws MethodError s .+ s + + # test for dimensionality track + @test Base.broadcasted(+, s, s) isa Broadcast.Broadcasted{<:Broadcast.AbstractArrayStyle{1}} + @test Base.broadcasted(+, s, 1:2) isa Broadcast.Broadcasted{<:Broadcast.AbstractArrayStyle{1}} + @test Base.broadcasted(+, s, reshape(1:2,1,2)) isa Broadcast.Broadcasted{<:Broadcast.AbstractArrayStyle{2}} + @test Base.broadcasted(+, reshape(1:2,1,1,2), s) isa Broadcast.Broadcasted{<:Broadcast.AbstractArrayStyle{3}} + + a = StructArray([1;2+im]) + b = StructArray([1;;2+im]) + @test a .+ b == a .+ collect(b) == collect(a) .+ b == collect(a) .+ collect(b) + @test a .+ Any[1] isa StructArray + + # issue #185 + A = StructArray(randn(ComplexF64, 3, 3)) + B = randn(ComplexF64, 3, 3) + c = StructArray(randn(ComplexF64, 3)) + @test (A .= B .* c) === A end @testset "staticarrays" begin