Specialize triu/tril for StaticMatrix (JuliaArrays#1241)

Author: jishnub

Diff for: src/StaticArrays.jl

@@ -16,7 +16,8 @@ using LinearAlgebra
 import LinearAlgebra: transpose, adjoint, dot, eigvals, eigen, lyap, tr,
                       kron, diag, norm, dot, diagm, lu, svd, svdvals, pinv,
                       factorize, ishermitian, issymmetric, isposdef, issuccess, normalize,
-                      normalize!, Eigen, det, logdet, logabsdet, cross, diff, qr, \
+                      normalize!, Eigen, det, logdet, logabsdet, cross, diff, qr, \,
+                      triu, tril
 using LinearAlgebra: checksquare
 
 using PrecompileTools

Diff for: src/linalg.jl

@@ -522,3 +522,37 @@ end
 # Some shimming for special linear algebra matrix types
 @inline LinearAlgebra.Symmetric(A::StaticMatrix, uplo::Char='U') = (checksquare(A); Symmetric{eltype(A),typeof(A)}(A, uplo))
 @inline LinearAlgebra.Hermitian(A::StaticMatrix, uplo::Char='U') = (checksquare(A); Hermitian{eltype(A),typeof(A)}(A, uplo))
+
+# triu/tril
+function triu(S::StaticMatrix, k::Int=0)
+    if length(S) <= 32
+        C = CartesianIndices(S)
+        t = Tuple(S)
+        for (linind, CI) in enumerate(C)
+            i, j = Tuple(CI)
+            if j-i < k
+                t = Base.setindex(t, zero(t[linind]), linind)
+            end
+        end
+        similar_type(S)(t)
+    else
+        M = triu!(copyto!(similar(S), S), k)
+        similar_type(S)(M)
+    end
+end
+function tril(S::StaticMatrix, k::Int=0)
+    if length(S) <= 32
+        C = CartesianIndices(S)
+        t = Tuple(S)
+        for (linind, CI) in enumerate(C)
+            i, j = Tuple(CI)
+            if j-i > k
+                t = Base.setindex(t, zero(t[linind]), linind)
+            end
+        end
+        similar_type(S)(t)
+    else
+        M = tril!(copyto!(similar(S), S), k)
+        similar_type(S)(M)
+    end
+end

Diff for: test/linalg.jl

@@ -471,4 +471,14 @@ end
         m23 = SA[1 2 3; 4 5 6]
         @test_inlined checksquare(m23) false
     end
+
+    @testset "triu/tril" begin
+        for S in (SMatrix{7,5}(1:35), MMatrix{4,6}(1:24), SizedArray{Tuple{2,2}}([1 2; 3 4]))
+            M = Matrix(S)
+            for k in -10:10
+                @test triu(S, k) == triu(M, k)
+                @test tril(S, k) == tril(M, k)
+            end
+        end
+    end
 end