Fix scaling block unit triangular matrices

(cherry picked from commit 2148fa4)

Author: jishnub

src/triangular.jl

@@ -1270,6 +1270,15 @@ end
 # Generic routines #
 ####################
 
+function _set_diag!(B::UpperOrLowerTriangular, x)
+    # get a mutable array to modify the diagonal
+    Bm = parent(B) isa StridedArray ? B : copy!(similar(B), B)
+    for i in diagind(Bm.data, IndexStyle(Bm.data))
+        Bm.data[i] = x
+    end
+    Bm
+end
+
 for (t, unitt) in ((UpperTriangular, UnitUpperTriangular),
                    (LowerTriangular, UnitLowerTriangular))
     tstrided = t{<:Any, <:StridedMaybeAdjOrTransMat}
@@ -1283,10 +1292,7 @@ for (t, unitt) in ((UpperTriangular, UnitUpperTriangular),
 
         function (*)(A::$unitt, x::Number)
             B = $t(A.data)*x
-            for i in axes(A, 1)
-                B.data[i,i] = x
-            end
-            return B
+            _set_diag!(B, oneunit(eltype(A)) * x)
         end
 
         (*)(x::Number, A::$t) = $t(x*A.data)
@@ -1298,10 +1304,7 @@ for (t, unitt) in ((UpperTriangular, UnitUpperTriangular),
 
         function (*)(x::Number, A::$unitt)
             B = x*$t(A.data)
-            for i in axes(A, 1)
-                B.data[i,i] = x
-            end
-            return B
+            _set_diag!(B, x * oneunit(eltype(A)))
         end
 
         (/)(A::$t, x::Number) = $t(A.data/x)
@@ -1313,11 +1316,7 @@ for (t, unitt) in ((UpperTriangular, UnitUpperTriangular),
 
         function (/)(A::$unitt, x::Number)
             B = $t(A.data)/x
-            invx = inv(x)
-            for i in axes(A, 1)
-                B.data[i,i] = invx
-            end
-            return B
+            _set_diag!(B, oneunit(eltype(A)) / x)
         end
 
         (\)(x::Number, A::$t) = $t(x\A.data)
@@ -1329,11 +1328,7 @@ for (t, unitt) in ((UpperTriangular, UnitUpperTriangular),
 
         function (\)(x::Number, A::$unitt)
             B = x\$t(A.data)
-            invx = inv(x)
-            for i in axes(A, 1)
-                B.data[i,i] = invx
-            end
-            return B
+            _set_diag!(B, x \ oneunit(eltype(A)))
         end
     end
 end

test/triangular.jl

@@ -923,4 +923,14 @@ end
     end
 end
 
+@testset "block unit triangular scaling" begin
+    m = SizedArrays.SizedArray{(2,2)}([1 2; 3 4])
+    U = UnitUpperTriangular(fill(m, 4, 4))
+    M = Matrix{eltype(U)}(U)
+    @test U/2 == M/2
+    @test 2\U == 2\M
+    @test U*2 == M*2
+    @test 2*U == 2*M
+end
+
 end # module TestTriangular