diff --git a/src/linalg.jl b/src/linalg.jl index 4d36c472..49da6b92 100644 --- a/src/linalg.jl +++ b/src/linalg.jl @@ -172,11 +172,18 @@ end end end -@inline diagm(v::StaticVector) = _diagm(Size(v), v) -@generated function _diagm(::Size{S}, v::StaticVector) where {S} - Snew = (S[1], S[1]) +@inline diagm(v::StaticVector, k::Type{Val{D}}=Val{0}) where {D} = _diagm(Size(v), v, k) +@generated function _diagm(::Size{S}, v::StaticVector, ::Type{Val{D}}) where {S,D} + S1 = S[1] + Snew1 = S1+abs(D) + Snew = (Snew1, Snew1) + Lnew = Snew1 * Snew1 T = eltype(v) - exprs = [i == j ? :(v[$i]) : zero(T) for i = 1:S[1], j = 1:S[1]] + ind = diagind(Snew1, Snew1, D) + exprs = fill(:(zero($T)), Lnew) + for n = 1:S[1] + exprs[ind[n]] = :(v[$n]) + end return quote $(Expr(:meta, :inline)) @inbounds return similar_type($v, Size($Snew))(tuple($(exprs...))) @@ -323,4 +330,3 @@ end @inline Base.LinAlg.Symmetric(A::StaticMatrix, uplo::Char='U') = (Base.LinAlg.checksquare(A);Symmetric{eltype(A),typeof(A)}(A, uplo)) @inline Base.LinAlg.Hermitian(A::StaticMatrix, uplo::Char='U') = (Base.LinAlg.checksquare(A);Hermitian{eltype(A),typeof(A)}(A, uplo)) - diff --git a/test/linalg.jl b/test/linalg.jl index ab7a7486..22ebf253 100644 --- a/test/linalg.jl +++ b/test/linalg.jl @@ -40,6 +40,8 @@ using StaticArrays, Base.Test @testset "diagm()" begin @test @inferred(diagm(SVector(1,2))) === @SMatrix [1 0; 0 2] + @test @inferred(diagm(SVector(1,2,3), Val{2})) == diagm([1,2,3], 2) + @test @inferred(diagm(SVector(1,2,3), Val{-2})) == diagm([1,2,3], -2) end @testset "diag()" begin