Make StructArrayStyle track inputs dimension

fix #185

Author: N5N3

src/structarray.jl

@@ -445,7 +445,15 @@ end
 # broadcast
 import Base.Broadcast: BroadcastStyle, ArrayStyle, AbstractArrayStyle, Broadcasted, DefaultArrayStyle
 
-struct StructArrayStyle{Style} <: AbstractArrayStyle{Any} end
+struct StructArrayStyle{S,N} <: AbstractArrayStyle{N} end
+
+# Here we define the dimension tracking behaviour of StructArrayStyle
+function StructArrayStyle{S,M}(::Val{N}) where {S,M,N}
+    if S <: AbstractArrayStyle{M}
+        return StructArrayStyle{typeof(S(Val(N))),N}()
+    end
+    return StructArrayStyle{S,N}()
+end
 
 @inline combine_style_types(::Type{A}, args...) where A<:AbstractArray =
     combine_style_types(BroadcastStyle(A), args...)
@@ -455,9 +463,9 @@ combine_style_types(s::BroadcastStyle) = s
 
 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}}, ::Type{ElType}) where {S<:DefaultArrayStyle,ElType} =
     isstructtype(ElType) ? similar(StructArray{ElType}, axes(bc)) : similar(Array{ElType}, axes(bc))
 
 # for aliasing analysis during broadcast

test/runtests.jl

@@ -926,8 +926,24 @@ 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, [1;;2]) isa Broadcast.Broadcasted{<:Broadcast.AbstractArrayStyle{2}}
+    @test Base.broadcasted(+, [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)
+
+    # 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