broadcasting for Diagonal StaticArrays (#914)

* broadcasting for Diagonal StaticArrays

* Add tests for SizedArrays

* add tests for inplace broadcasting

* version bump to v1.2.3

Author: jishnub

Project.toml

@@ -1,6 +1,6 @@
 name = "StaticArrays"
 uuid = "90137ffa-7385-5640-81b9-e52037218182"
-version = "1.2.2"
+version = "1.2.3"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/broadcast.jl

@@ -13,6 +13,7 @@ StaticArrayStyle{M}(::Val{N}) where {M,N} = StaticArrayStyle{N}()
 BroadcastStyle(::Type{<:StaticArray{<:Tuple, <:Any, N}}) where {N} = StaticArrayStyle{N}()
 BroadcastStyle(::Type{<:Transpose{<:Any, <:StaticArray{<:Tuple, <:Any, N}}}) where {N} = StaticArrayStyle{N}()
 BroadcastStyle(::Type{<:Adjoint{<:Any, <:StaticArray{<:Tuple, <:Any, N}}}) where {N} = StaticArrayStyle{N}()
+BroadcastStyle(::Type{<:Diagonal{<:Any, <:StaticArray{<:Tuple, <:Any, 1}}}) = StaticArrayStyle{2}()
 # Precedence rules
 BroadcastStyle(::StaticArrayStyle{M}, ::DefaultArrayStyle{N}) where {M,N} =
     DefaultArrayStyle(Val(max(M, N)))
@@ -97,7 +98,7 @@ 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}}, a)]
+    first_staticarray = a[findfirst(ai -> ai <: Union{StaticArray, Transpose{<:Any, <:StaticArray}, Adjoint{<:Any, <:StaticArray}, Diagonal{<:Any, <:StaticArray}}, a)]
 
     if prod(newsize) == 0
         # Use inference to get eltype in empty case (see also comments in _map)

test/broadcast.jl

@@ -240,4 +240,34 @@ end
         # Unfortunately this case of nested broadcasting is not inferred
         @test_broken @inferred(SA[1,2,3] .* (SA[1,0],))
     end
+
+    @testset "SDiagonal" begin
+        for DS in Any[Diagonal(SVector{2}(1:2)), Diagonal(MVector{2}(1:2)), Diagonal(SizedArray{Tuple{2}}(1:2))],
+                S in Any[SVector{2}(1:2), MVector{2}(1:2), SizedArray{Tuple{2}}(1:2)]
+            @test DS .* S isa StaticArray
+            @test DS .* S == collect(DS) .* collect(S)
+            @test DS .* collect(S) == collect(DS) .* collect(S)
+            @test DS .* S' isa StaticArray
+            @test DS .* S' == collect(DS) .* collect(S)
+            @test DS .* collect(S') == collect(DS) .* collect(S)
+            @test S .* DS .* S' isa StaticArray
+            @test S .* DS .* S' == collect(S) .* collect(DS) .* collect(S)
+            DS2 = Diagonal(S)
+            @test DS .* DS2 isa StaticArray
+            @test DS .* DS2 == collect(DS) .* collect(DS2)
+            @test DS .* collect(DS2) == collect(DS) .* collect(DS2)
+
+            # inplace broadcasting for mutable diagonal types
+            DS2 = Diagonal(MVector{2}(diag(DS)))
+            DS2 .*= S
+            @test DS2 == DS .* S
+            DS2 .= DS
+            @test DS2 == DS
+            DS2 .+= DS
+            @test DS2 == DS .+ DS
+            DS2 .= DS
+            DS2 .*= DS
+            @test DS2 == DS .* DS
+        end
+    end
 end