Convert to StaticArrayStyle

We first call broadcast from `StaticArrays` then split the output.
This should has no extra runtime overhead. But some type info might missing because the eltype change. I think there's no better ways as we don't want to depend on the full  `StaticArrays`.

Author: N5N3

src/staticarrays_support.jl

@@ -1,4 +1,4 @@
-import StaticArraysCore: StaticArray, FieldArray, tuple_prod
+using StaticArraysCore: StaticArray, FieldArray, tuple_prod
 
 """
     StructArrays.staticschema(::Type{<:StaticArray{S, T}}) where {S, T}
@@ -27,3 +27,42 @@ StructArrays.component(s::StaticArray, i) = getindex(s, i)
 end
 StructArrays.component(s::FieldArray, i) = invoke(StructArrays.component, Tuple{Any, Any}, s, i)
 StructArrays.createinstance(T::Type{<:FieldArray}, args...) = invoke(createinstance, Tuple{Type{<:Any}, Vararg}, T, args...)
+
+# Broadcast overload
+using StaticArraysCore: StaticArrayStyle
+import StaticArraysCore: Size, is_staticarray_like, similar_type
+StructStaticArrayStyle{N} = StructArrayStyle{StaticArrayStyle{N}, N}
+function Broadcast.instantiate(bc::Broadcasted{StructStaticArrayStyle{M}}) where {M}
+    bc′ = Broadcast.instantiate(convert(Broadcasted{StaticArrayStyle{M}}, bc))
+    return convert(Broadcasted{StructStaticArrayStyle{M}}, bc′)
+end
+function Broadcast._axes(bc::Broadcasted{StructStaticArrayStyle{M}}, ::Nothing) where {M}
+    return Broadcast._axes(convert(Broadcasted{StaticArrayStyle{M}}, bc), nothing)
+end
+
+# StaticArrayStyle has no similar defined.
+# Overload `Base.copy` instead.
+@inline function Base.copy(bc::Broadcasted{StructStaticArrayStyle{M}}) where {M}
+    sa = copy(convert(Broadcasted{StaticArrayStyle{M}}, bc))
+    ET = eltype(sa)
+    isnonemptystructtype(ET) || return sa
+    elements = Tuple(sa)
+    arrs = ntuple(Val(fieldcount(ET))) do i
+        similar_type(sa, fieldtype(ET, i), Size(sa))(_getfields(elements, i))
+    end
+    return StructArray{ET}(arrs)
+end
+
+@inline function _getfields(x::Tuple, i::Int)
+    if @generated
+        return Expr(:tuple, (:(getfield(x[$j], i)) for j in 1:fieldcount(x))...)
+    else
+        return map(Base.Fix2(getfield, i), x)
+    end
+end
+
+Size(::Type{SA}) where {SA<:StructArray} = Size(fieldtype(array_types(SA), 1))
+is_staticarray_like(x::StructArray) = any(is_staticarray_like, components(x))
+function similar_type(::Type{SA}, ::Type{T}, s::Size{S}) where {SA<:StructArray, T, S}
+    return similar_type(fieldtype(array_types(SA), 1), T, s)
+end

test/runtests.jl

@@ -1160,12 +1160,24 @@ Base.BroadcastStyle(::Broadcast.ArrayStyle{MyArray2}, S::Broadcast.DefaultArrayS
         end
         testset = Any[StructArray([1;2+im]),
                     1:2, 
-                    (1,2), 
+                    (1,2),
+                    StructArray(@SArray [1 1+2im]),
+                    (@SArray [1 2])
                     ]
         for aa in testset, bb in testset, cc in testset
             _test(aa, bb, cc)
         end
     end
+
+    @testset "StructStaticArray" begin
+        bclog(s) = log.(s)
+        test_allocated(f, s) = @test (@allocated f(s)) == 0
+        a = @SMatrix [float(i) for i in 1:10, j in 1:10]
+        b = @SMatrix [0. for i in 1:10, j in 1:10]
+        s = StructArray{ComplexF64}((a , b))
+        @test (@inferred bclog(s)) isa typeof(s)
+        test_allocated(bclog, s)
+    end
 end
 
 @testset "map" begin