From cbd854b0f4af800fac543676c4b3f938ebc48e62 Mon Sep 17 00:00:00 2001
From: Sakse <17059936+dalum@users.noreply.github.com>
Date: Sat, 6 Jun 2020 15:20:38 +0200
Subject: [PATCH] Add range indexing to UniformScaling (#24359)

* Add range indexing to UniformScaling

* Add tests

* Return sparse matrix when indexing UniformScaling

* start -> first

* Correctly handle non-integer steps

* Make getindex(::UniformScaling, ranges...) return a dense array

* Remove unneccessary (and slower) branch

* Add (range, integer) indexing to UniformScaling

* Update tests

* Fix unbound param

* Add NEWS item

* Update docstring and add compat annotation

Co-authored-by: Evey Dee <eveydee@users.noreply.github.com>
---
 NEWS.md                                     |  1 +
 stdlib/LinearAlgebra/src/uniformscaling.jl  | 36 +++++++++++++++++++--
 stdlib/LinearAlgebra/test/uniformscaling.jl | 26 +++++++++++++--
 3 files changed, 59 insertions(+), 4 deletions(-)

diff --git a/NEWS.md b/NEWS.md
index 9966811f8e1c7..71a2d66d7946c 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -50,6 +50,7 @@ Standard library changes
 
 #### LinearAlgebra
 * New method `LinearAlgebra.issuccess(::CholeskyPivoted)` for checking whether pivoted Cholesky factorization was successful ([#36002]).
+* `UniformScaling` can now be indexed into using ranges to return dense matrices and vectors ([#24359]).
 
 #### Markdown
 
diff --git a/stdlib/LinearAlgebra/src/uniformscaling.jl b/stdlib/LinearAlgebra/src/uniformscaling.jl
index ed007fc48fa77..214f06c6af2e1 100644
--- a/stdlib/LinearAlgebra/src/uniformscaling.jl
+++ b/stdlib/LinearAlgebra/src/uniformscaling.jl
@@ -6,8 +6,13 @@ import Base: copy, adjoint, getindex, show, transpose, one, zero, inv,
 """
     UniformScaling{T<:Number}
 
-Generically sized uniform scaling operator defined as a scalar times the
-identity operator, `λ*I`. See also [`I`](@ref).
+Generically sized uniform scaling operator defined as a scalar times
+the identity operator, `λ*I`. Although without an explicit `size`, it
+acts similarly to a matrix in many cases and includes support for some
+indexing. See also [`I`](@ref).
+
+!!! compat "Julia 1.6"
+     Indexing using ranges is available as of Julia 1.6.
 
 # Examples
 ```jldoctest
@@ -24,6 +29,11 @@ julia> J*A
 2×2 Array{Float64,2}:
  2.0  4.0
  6.0  8.0
+
+julia> J[1:2, 1:2]
+2×2 Array{Float64,2}:
+ 2.0  0.0
+ 0.0  2.0
 ```
 """
 struct UniformScaling{T<:Number}
@@ -78,6 +88,28 @@ ndims(J::UniformScaling) = 2
 Base.has_offset_axes(::UniformScaling) = false
 getindex(J::UniformScaling, i::Integer,j::Integer) = ifelse(i==j,J.λ,zero(J.λ))
 
+getindex(x::UniformScaling, n::Integer, m::AbstractRange{<:Integer}) = getindex(x, m, n)
+function getindex(x::UniformScaling{T}, n::AbstractRange{<:Integer}, m::Integer) where T
+    v = zeros(T, length(n))
+    @inbounds for (i,ii) in enumerate(n)
+        if ii == m
+            v[i] = x.λ
+        end
+    end
+    return v
+end
+
+
+function getindex(x::UniformScaling{T}, n::AbstractRange{<:Integer}, m::AbstractRange{<:Integer}) where T
+    A = zeros(T, length(n), length(m))
+    @inbounds for (j,jj) in enumerate(m), (i,ii) in enumerate(n)
+        if ii == jj
+            A[i,j] = x.λ
+        end
+    end
+    return A
+end
+
 function show(io::IO, ::MIME"text/plain", J::UniformScaling)
     s = "$(J.λ)"
     if occursin(r"\w+\s*[\+\-]\s*\w+", s)
diff --git a/stdlib/LinearAlgebra/test/uniformscaling.jl b/stdlib/LinearAlgebra/test/uniformscaling.jl
index 1564198443f24..320ca993f438c 100644
--- a/stdlib/LinearAlgebra/test/uniformscaling.jl
+++ b/stdlib/LinearAlgebra/test/uniformscaling.jl
@@ -11,8 +11,6 @@ using .Main.Quaternions
 Random.seed!(123)
 
 @testset "basic functions" begin
-    @test I[1,1] == 1 # getindex
-    @test I[1,2] == 0 # getindex
     @test I === I' # transpose
     @test ndims(I) == 2
     @test one(UniformScaling{Float32}) == UniformScaling(one(Float32))
@@ -27,6 +25,30 @@ Random.seed!(123)
     @test opnorm(UniformScaling(1+im)) ≈ sqrt(2)
 end
 
+@testset "getindex" begin
+    @test I[1,1] == 1
+    @test I[1,2] == 0
+
+    J = I(15)
+    for (a, b) in [
+        # indexing that returns a Vector
+        (1:10, 1),
+        (4, 1:10),
+        (11, 1:10),
+        # indexing that returns a Matrix
+        (1:2, 1:2),
+        (1:2:3, 1:2:3),
+        (1:2:8, 2:2:9),
+        (1:2:8, 9:-4:1),
+        (9:-4:1, 1:2:8),
+        (2:3, 1:2),
+        (2:-1:1, 1:2),
+        (1:2:9, 5:2:13),
+    ]
+        @test I[a,b] == J[a,b]
+    end
+end
+
 @testset "sqrt, exp, log, and trigonometric functions" begin
     # convert to a dense matrix with random size
     M(J) = (N = rand(1:10); Matrix(J, N, N))