Convert to StaticArrayStyle

Author: N5N3

src/staticarrays_support.jl

@@ -1,4 +1,4 @@
-import StaticArrays: StaticArray, FieldArray, tuple_prod
+using StaticArrays: StaticArrays, StaticArray, FieldArray, tuple_prod, StaticArrayStyle
 
 """
     StructArrays.staticschema(::Type{<:StaticArray{S, T}}) where {S, T}
@@ -26,4 +26,62 @@ StructArrays.component(s::StaticArray, i) = getindex(s, i)
     invoke(StructArrays.staticschema, Tuple{Type{<:Any}}, T)
 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...)
+StructArrays.createinstance(T::Type{<:FieldArray}, args...) = invoke(createinstance, Tuple{Type{<:Any}, Vararg}, T, args...)
+
+# Broadcast overload
+import StaticArrays: Size, isstatic, similar_type
+using StaticArrays: first_statictype, broadcast_sizes, SOneTo
+import Base.Broadcast: instantiate
+StructStaticArrayStyle{N} = StructArrayStyle{StaticArrayStyle{N}, N}
+function instantiate(bc::Broadcasted{StructStaticArrayStyle{M}}) where {M}
+    bc′ = instantiate(convert(Broadcasted{StaticArrayStyle{M}}, bc))
+    return convert(Broadcasted{StructStaticArrayStyle{M}}, bc′)
+end
+function Broadcast._axes(bc::Broadcasted{<:StructStaticArrayStyle}, ::Nothing)
+    return StaticArrays.static_combine_axes(bc.args...)
+end
+
+# StaticArrayStyle has no similar defined.
+# Overload `Base.copy` instead.
+@inline function Base.copy(B::Broadcasted{<:StructStaticArrayStyle})
+    flat = Broadcast.flatten(B); as = flat.args; f = flat.f
+    argsizes = broadcast_sizes(as...)
+    ax = axes(B)
+    ax isa Tuple{Vararg{SOneTo}} || error("Dimension is not static. Please file a bug.")
+    return _broadcast(f, Size(map(length, ax)), argsizes, as...)
+end
+@inline function _broadcast(f, sz::Size{newsize}, s::Tuple{Vararg{Size}}, a...) where newsize
+    AT = first_statictype(a...)
+    if prod(newsize) == 0
+        # Use inference to get eltype in empty case (see also comments in _map)
+        eltys = Tuple{map(eltype, a)...}
+        T = Core.Compiler.return_type(f, eltys)
+        return _struct_static_similar(T, AT, sz, ())
+    end
+    elements = StaticArrays.__broadcast(f, sz, s, a...)
+    return _struct_static_similar(eltype(elements), AT, sz, elements)
+end
+function _struct_static_similar(::Type{ET}, ::Type{AT}, sz, elements::Tuple) where {ET, AT}
+    if isnonemptystructtype(ET)
+        arrs = ntuple(Val(fieldcount(ET))) do i
+            similar_type(AT, fieldtype(ET, i), sz)(_getfields(elements, i))
+        end
+        return StructArray{ET}(arrs)
+    else
+        return similar_type(AT, ET, sz)(elements)
+    end
+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))
+isstatic(x::StructArray) = isstatic(component(x, 1))
+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