Make broadcast extendable

Author: N5N3

src/broadcast.jl

@@ -97,9 +97,22 @@ end
 scalar_getindex(x) = x
 scalar_getindex(x::Ref) = x[]
 
-@generated function _broadcast(f, ::Size{newsize}, s::Tuple{Vararg{Size}}, a...) where newsize
-    first_staticarray = a[findfirst(ai -> ai <: Union{StaticArray, Transpose{<:Any, <:StaticArray}, Adjoint{<:Any, <:StaticArray}, Diagonal{<:Any, <:StaticArray}}, a)]
+isstatic(::StaticArray) = true
+isstatic(::Transpose{<:Any, <:StaticArray}) = true
+isstatic(::Adjoint{<:Any, <:StaticArray}) = true
+isstatic(::Diagonal{<:Any, <:StaticArray}) = true
+isstatic(_) = false
+
+@inline first_statictype(x, y...) = isstatic(x) ? typeof(x) : first_statictype(y...)
+first_statictype() = error("unresolved dest type")
+
+@inline function _broadcast(f, sz::Size{newsize}, s::Tuple{Vararg{Size}}, a...) where newsize
+    first_staticarray = first_statictype(a...)
+    elements = __broadcast(f, sz, s, a...)
+    @inbounds return similar_type(first_staticarray, eltype(elements), Size(newsize))(elements)
+end
 
+@generated function __broadcast(f, ::Size{newsize}, s::Tuple{Vararg{Size}}, a...) where newsize
     if prod(newsize) == 0
         # Use inference to get eltype in empty case (see also comments in _map)
         eltys = [:(eltype(a[$i])) for i ∈ 1:length(a)]
@@ -123,8 +136,7 @@ scalar_getindex(x::Ref) = x[]
 
     return quote
         @_inline_meta
-        @inbounds elements = tuple($(exprs...))
-        @inbounds return similar_type($first_staticarray, eltype(elements), Size(newsize))(elements)
+        @inbounds return elements = tuple($(exprs...))
     end
 end
 

src/precompile.jl

@@ -22,7 +22,7 @@ function _precompile_()
     end
 
     # Some expensive generators
-    @assert precompile(Tuple{typeof(which(_broadcast,(Any,Size,Tuple{Vararg{Size}},Vararg{Any},)).generator.gen),Any,Any,Any,Any,Any,Any})
+    @assert precompile(Tuple{typeof(which(__broadcast,(Any,Size,Tuple{Vararg{Size}},Vararg{Any},)).generator.gen),Any,Any,Any,Any,Any,Any})
     @assert precompile(Tuple{typeof(which(_zeros,(Size,Type{<:StaticArray},)).generator.gen),Any,Any,Any,Type,Any})
     @assert precompile(Tuple{typeof(which(combine_sizes,(Tuple{Vararg{Size}},)).generator.gen),Any,Any})
     @assert precompile(Tuple{typeof(which(_mapfoldl,(Any,Any,Colon,Any,Size,Vararg{StaticArray},)).generator.gen),Any,Any,Any,Any,Any,Any,Any,Any})

test/broadcast.jl

@@ -282,3 +282,30 @@ end
     end
 
 end
+
+# A help struct to test style-based broadcast dispatch with unknown array wrapper.
+# `WrapArray(A)` behaves like `A` during broadcast. But its not a `StaticArray`.
+struct WrapArray{T,N,P<:AbstractArray{T,N}} <: AbstractArray{T,N}
+    data::P
+end
+Base.@propagate_inbounds Base.getindex(A::WrapArray, i::Integer...) = A.data[i...]
+Base.@propagate_inbounds Base.setindex!(A::WrapArray, v::Any, i::Integer...) = setindex!(A.data, v, i...)
+Base.size(A::WrapArray) = size(A.data)
+Broadcast.BroadcastStyle(::Type{WrapArray{T,N,P}}) where {T,N,P} = Broadcast.BroadcastStyle(P)
+StaticArrays.isstatic(A::WrapArray) = StaticArrays.isstatic(A.data)
+StaticArrays.Size(::Type{WrapArray{T,N,P}}) where {T,N,P} = StaticArrays.Size(P)
+function StaticArrays.similar_type(::Type{WrapArray{T,N,P}}, ::Type{t}, s::Size{S}) where {T,N,P,t,S}
+    return StaticArrays.similar_type(P, t, s)
+end
+
+@testset "Broadcast with unknown wrapper" begin
+    data = (1, 2)
+    for T in (SVector{2}, MVector{2})
+        a = T(data)
+        b = WrapArray(a)
+        @test @inferred(b .+ a) isa T
+        @test @inferred(b .+ b) isa T
+        @test @inferred(b .+ (1, 2)) isa T
+        @test b .+ a == b .+ b == a .+ a 
+    end
+end