axes(::SDiagonal) is statically sized

vector indexing for SDiagonal produces SArrays

Author: jishnub

src/SDiagonal.jl

@@ -17,6 +17,11 @@ SDiagonal(a::StaticMatrix{N,N,T}) where {N,T} = Diagonal(diag(a))
 size(::Type{<:SDiagonal{N}}) where {N} = (N,N)
 size(::Type{<:SDiagonal{N}}, d::Int) where {N} = d > 2 ? 1 : N
 
+Base.axes(D::SDiagonal) = (ax = axes(diag(D), 1); (ax, ax))
+Base.axes(D::SDiagonal, d) = d <= 2 ? axes(D)[d] : SOneTo(1)
+
+Base.reshape(a::SDiagonal, s::Tuple{SOneTo,Vararg{SOneTo}}) = reshape(a, homogenize_shape(s))
+
 # define specific methods to avoid allocating mutable arrays
 \(D::SDiagonal, b::AbstractVector) = D.diag .\ b
 \(D::SDiagonal, b::StaticVector) = D.diag .\ b # catch ambiguity

src/indexing.jl

@@ -377,3 +377,29 @@ Base.unsafe_view(A::AbstractArray, i1::StaticIndexing, indices::StaticIndexing..
 # the tuple indices has to have at least one element to prevent infinite
 # recursion when viewing a zero-dimensional array (see issue #705)
 Base.SubArray(A::AbstractArray, indices::Tuple{StaticIndexing, Vararg{StaticIndexing}}) = Base.SubArray(A, map(unwrap, indices))
+
+###########################################################
+# SDiagonal
+###########################################################
+
+# SDiagonal uses Cartesian indexing, and the canonical indexing methods shadow getindex for Diagonal
+# these are needed for ambiguity resolution
+@inline function getindex(D::SDiagonal, i::Int, j::Int)
+    @boundscheck checkbounds(D, i, j)
+    if i == j
+        @inbounds r = diag(D)[i]
+    else
+        r = LinearAlgebra.diagzero(D, i, j)
+    end
+    r
+end
+@inline function getindex(D::SDiagonal, i::Int...)
+    @boundscheck checkbounds(D, i...)
+    @inbounds r = D[eachindex(D)[i...]]
+    r
+end
+# Ensure that vector indexing with static types lead to SArrays
+@propagate_inbounds function getindex(a::SDiagonal, inds::Union{Int, StaticArray{<:Tuple, Int}, SOneTo, Colon}...)
+    ar = reshape(a, Val(length(inds)))
+    _getindex(ar, index_sizes(Size(ar), inds...), inds)
+end

test/SDiagonal.jl

@@ -70,6 +70,15 @@ using StaticArrays, Test, LinearAlgebra
 
         @test length(m) === 4*4
 
+        m2 = SMatrix{4,4}(m)
+        @test axes(m) === axes(m2)
+        @test axes(m, 1) === axes(m2, 1)
+        @test axes(m, 3) == SOneTo(1)
+
+        @test m[:, 1] === SVector{4}(m[1,1], 0, 0, 0)
+        @test m[:, :] === m2
+        @test m[2, 2, 1] === m[2, 2]
+
         @test_throws Exception m[1] = 1
 
         b = @SVector [2,-1,2,1]