Implement diag() with Val as second argument (#226)

Author: wsshin

src/StaticArrays.jl

@@ -8,7 +8,7 @@ import Base: getindex, setindex!, size, similar, vec, show,
              length, convert, promote_op, promote_rule, map, map!, reduce, reducedim, mapreducedim,
              mapreduce, broadcast, broadcast!, conj, transpose, ctranspose,
              hcat, vcat, ones, zeros, eye, one, cross, vecdot, reshape, fill,
-             fill!, det, inv, eig, eigvals, expm, sqrtm, trace, vecnorm, norm, dot, diagm,
+             fill!, det, inv, eig, eigvals, expm, sqrtm, trace, vecnorm, norm, dot, diagm, diag,
              sum, diff, prod, count, any, all, minimum,
              maximum, extrema, mean, copy, rand, randn, randexp, rand!, randn!,
              randexp!, normalize, normalize!, read, read!, write

src/linalg.jl

@@ -183,6 +183,19 @@ end
     end
 end
 
+@inline diag(m::StaticMatrix, k::Type{Val{D}}=Val{0}) where {D} = _diag(Size(m), m, k)
+@generated function _diag(::Size{S}, m::StaticMatrix, ::Type{Val{D}}) where {S,D}
+    S1, S2 = S
+    rng = D ≤ 0 ? range(1-D, S1+1, min(S1+D, S2)) : range(D*S1+1, S1+1, min(S1, S2-D))
+    Snew = length(rng)
+    T = eltype(m)
+    exprs = [:(m[$i]) for i = rng]
+    return quote
+        $(Expr(:meta, :inline))
+        @inbounds return similar_type($m, Size($Snew))(tuple($(exprs...)))
+    end
+end
+
 @inline cross(a::StaticVector, b::StaticVector) = _cross(same_size(a, b), a, b)
 _cross(::Size{S}, a::StaticVector, b::StaticVector) where {S} = error("Cross product not defined for $(S[1])-vectors")
 @inline function _cross(::Size{(2,)}, a::StaticVector, b::StaticVector)

test/linalg.jl

@@ -40,6 +40,12 @@
         @test @inferred(diagm(SVector(1,2))) === @SMatrix [1 0; 0 2]
     end
 
+    @testset "diag()" begin
+        @test @inferred(diag(@SMatrix([0 1; 2 3]))) === SVector(0, 3)
+        @test @inferred(diag(@SMatrix([0 1 2; 3 4 5]), Val{1})) === SVector(1, 5)
+        @test @inferred(diag(@SMatrix([0 1; 2 3; 4 5]), Val{-1})) === SVector(2, 5)
+    end
+
     @testset "one() and zero()" begin
         @test @inferred(one(SMatrix{2,2,Int})) === @SMatrix [1 0; 0 1]
         @test @inferred(one(SMatrix{2,2})) === @SMatrix [1.0 0.0; 0.0 1.0]