Use IndexStyle in diagind to optionally return a range of CartesianIndexes (#52115)

Author: jishnub

stdlib/LinearAlgebra/src/dense.jl

@@ -198,9 +198,17 @@ function fillband!(A::AbstractMatrix{T}, x, l, u) where T
     return A
 end
 
-diagind(m::Integer, n::Integer, k::Integer=0) =
+diagind(m::Integer, n::Integer, k::Integer=0) = diagind(IndexLinear(), m, n, k)
+diagind(::IndexLinear, m::Integer, n::Integer, k::Integer=0) =
     k <= 0 ? range(1-k, step=m+1, length=min(m+k, n)) : range(k*m+1, step=m+1, length=min(m, n-k))
 
+function diagind(::IndexCartesian, m::Integer, n::Integer, k::Integer=0)
+    Cstart = CartesianIndex(1 + max(0,-k), 1 + max(0,k))
+    Cstep = CartesianIndex(1, 1)
+    length = max(0, k <= 0 ? min(m+k, n) : min(m, n-k))
+    StepRangeLen(Cstart, Cstep, length)
+end
+
 """
     diagind(M, k::Integer=0)
 
@@ -222,7 +230,7 @@ julia> diagind(A,-1)
 """
 function diagind(A::AbstractMatrix, k::Integer=0)
     require_one_based_indexing(A)
-    diagind(size(A,1), size(A,2), k)
+    diagind(IndexStyle(A), size(A,1), size(A,2), k)
 end
 
 """

stdlib/LinearAlgebra/test/bidiag.jl

@@ -843,4 +843,12 @@ end
     @test all(iszero, diag(B, 1))
 end
 
+@testset "diagind" begin
+    B = Bidiagonal(1:4, 1:3, :U)
+    M = Matrix(B)
+    @testset for k in -4:4
+        @test B[diagind(B,k)] == M[diagind(M,k)]
+    end
+end
+
 end # module TestBidiagonal

stdlib/LinearAlgebra/test/diagonal.jl

@@ -1227,4 +1227,12 @@ end
     @test *(Diagonal(ones(n)), Diagonal(1:n), Diagonal(ones(n)), Diagonal(1:n)) isa Diagonal
 end
 
+@testset "diagind" begin
+    D = Diagonal(1:4)
+    M = Matrix(D)
+    @testset for k in -4:4
+        @test D[diagind(D,k)] == M[diagind(M,k)]
+    end
+end
+
 end # module TestDiagonal