Introduce StructuredMatrixStyle{Matrix} (fixes #33397) (#33506)

* Introduce `StructuredMatrixStyle{Matrix}` (fixes #33397)

* Remove `Matrix` from `const StructuredMatrix`

* Prevent `StructuredMatrixStyle{<:Matrix}` from falling back to `DefaultArrayStyle{2}`

Author: goggle

stdlib/LinearAlgebra/src/structuredbroadcast.jl

@@ -40,6 +40,17 @@ Broadcast.BroadcastStyle(::StructuredMatrixStyle{<:UnitLowerTriangular}, ::Struc
 Broadcast.BroadcastStyle(::StructuredMatrixStyle{<:UnitUpperTriangular}, ::StructuredMatrixStyle{<:Union{Diagonal,UpperTriangular,UnitUpperTriangular}}) =
     StructuredMatrixStyle{UpperTriangular}()
 
+Broadcast.BroadcastStyle(::StructuredMatrixStyle{<:Union{LowerTriangular,UnitLowerTriangular}}, ::StructuredMatrixStyle{<:Union{UpperTriangular,UnitUpperTriangular}}) =
+    StructuredMatrixStyle{Matrix}()
+Broadcast.BroadcastStyle(::StructuredMatrixStyle{<:Union{UpperTriangular,UnitUpperTriangular}}, ::StructuredMatrixStyle{<:Union{LowerTriangular,UnitLowerTriangular}}) =
+    StructuredMatrixStyle{Matrix}()
+
+# Make sure that `StructuredMatrixStyle{<:Matrix}` doesn't ever end up falling
+# through and give back `DefaultArrayStyle{2}`
+Broadcast.BroadcastStyle(T::StructuredMatrixStyle{<:Matrix}, ::StructuredMatrixStyle) = T
+Broadcast.BroadcastStyle(::StructuredMatrixStyle, T::StructuredMatrixStyle{<:Matrix}) = T
+Broadcast.BroadcastStyle(T::StructuredMatrixStyle{<:Matrix}, ::StructuredMatrixStyle{<:Matrix}) = T
+
 # All other combinations fall back to the default style
 Broadcast.BroadcastStyle(::StructuredMatrixStyle, ::StructuredMatrixStyle) = DefaultArrayStyle{2}()
 
@@ -69,6 +80,8 @@ structured_broadcast_alloc(bc, ::Type{<:UnitLowerTriangular}, ::Type{ElType}, n)
     UnitLowerTriangular(Array{ElType}(undef, n, n))
 structured_broadcast_alloc(bc, ::Type{<:UnitUpperTriangular}, ::Type{ElType}, n) where {ElType} =
     UnitUpperTriangular(Array{ElType}(undef, n, n))
+structured_broadcast_alloc(bc, ::Type{<:Matrix}, ::Type{ElType}, n) where {ElType} =
+    Matrix(Array{ElType}(undef, n, n))
 
 # A _very_ limited list of structure-preserving functions known at compile-time. This list is
 # derived from the formerly-implemented `broadcast` methods in 0.6. Note that this must

stdlib/LinearAlgebra/test/structuredbroadcast.jl

@@ -14,7 +14,8 @@ using Test, LinearAlgebra
     T = Tridiagonal(rand(N - 1), rand(N), rand(N - 1))
     U = UpperTriangular(rand(N,N))
     L = LowerTriangular(rand(N,N))
-    structuredarrays = (D, B, T, U, L)
+    M = Matrix(rand(N,N))
+    structuredarrays = (D, B, T, U, L, M)
     fstructuredarrays = map(Array, structuredarrays)
     for (X, fX) in zip(structuredarrays, fstructuredarrays)
         @test (Q = broadcast(sin, X); typeof(Q) == typeof(X) && Q == broadcast(sin, fX))
@@ -56,19 +57,22 @@ end
     T = Tridiagonal(rand(N - 1), rand(N), rand(N - 1))
     ◣ = LowerTriangular(rand(N,N))
     ◥ = UpperTriangular(rand(N,N))
+    M = Matrix(rand(N,N))
 
     @test broadcast!(sin, copy(D), D) == Diagonal(sin.(D))
     @test broadcast!(sin, copy(Bu), Bu) == Bidiagonal(sin.(Bu), :U)
     @test broadcast!(sin, copy(Bl), Bl) == Bidiagonal(sin.(Bl), :L)
     @test broadcast!(sin, copy(T), T) == Tridiagonal(sin.(T))
     @test broadcast!(sin, copy(◣), ◣) == LowerTriangular(sin.(◣))
     @test broadcast!(sin, copy(◥), ◥) == UpperTriangular(sin.(◥))
+    @test broadcast!(sin, copy(M), M) == Matrix(sin.(M))
     @test broadcast!(*, copy(D), D, A) == Diagonal(broadcast(*, D, A))
     @test broadcast!(*, copy(Bu), Bu, A) == Bidiagonal(broadcast(*, Bu, A), :U)
     @test broadcast!(*, copy(Bl), Bl, A) == Bidiagonal(broadcast(*, Bl, A), :L)
     @test broadcast!(*, copy(T), T, A) == Tridiagonal(broadcast(*, T, A))
     @test broadcast!(*, copy(◣), ◣, A) == LowerTriangular(broadcast(*, ◣, A))
     @test broadcast!(*, copy(◥), ◥, A) == UpperTriangular(broadcast(*, ◥, A))
+    @test broadcast!(*, copy(M), M, A) == Matrix(broadcast(*, M, A))
 
     @test_throws ArgumentError broadcast!(cos, copy(D), D) == Diagonal(sin.(D))
     @test_throws ArgumentError broadcast!(cos, copy(Bu), Bu) == Bidiagonal(sin.(Bu), :U)
@@ -93,7 +97,8 @@ end
     T = Tridiagonal(rand(N - 1), rand(N), rand(N - 1))
     U = UpperTriangular(rand(N,N))
     L = LowerTriangular(rand(N,N))
-    structuredarrays = (D, B, T, U, L)
+    M = Matrix(rand(N,N))
+    structuredarrays = (M, D, B, T, U, L)
     fstructuredarrays = map(Array, structuredarrays)
     for (X, fX) in zip(structuredarrays, fstructuredarrays)
         @test (Q = map(sin, X); typeof(Q) == typeof(X) && Q == map(sin, fX))
@@ -123,4 +128,22 @@ end
     end
 end
 
+@testset "Issue #33397" begin
+    N = 5
+    U = UpperTriangular(rand(N, N))
+    L = LowerTriangular(rand(N, N))
+    UnitU = UnitUpperTriangular(rand(N, N))
+    UnitL = UnitLowerTriangular(rand(N, N))
+    D = Diagonal(rand(N))
+    @test U .+ L .+ D == U + L + D
+    @test L .+ U .+ D == L + U + D
+    @test UnitU .+ UnitL .+ D == UnitU + UnitL + D
+    @test UnitL .+ UnitU .+ D == UnitL + UnitU + D
+    @test U .+ UnitL .+ D == U + UnitL + D
+    @test L .+ UnitU .+ D == L + UnitU + D
+    @test L .+ U .+ L .+ U == L + U + L + U
+    @test U .+ L .+ U .+ L == U + L + U + L
+    @test L .+ UnitL .+ UnitU .+ U .+ D == L + UnitL + UnitU + U + D
+    @test L .+ U .+ D .+ D .+ D .+ D == L + U + D + D + D + D
+end
 end