Band indexing for adj/trans (#1299)

The advantage of this is that we may forward the indexing to the parent,
and if the parent is a structured matrix and the band is a compile-time
constant, the value may be evaluated as a constant.

E.g.:
```julia
julia> A = UnitUpperTriangular(reshape(1:9, 3, 3))
3×3 UnitUpperTriangular{Int64, Base.ReshapedArray{Int64, 2, UnitRange{Int64}, Tuple{}}}:
 1  4  7
 ⋅  1  8
 ⋅  ⋅  1

julia> function f(A, i, ::Val{band}) where {band}
           x = Adjoint(A)[LinearAlgebra.BandIndex(band,i)]
           Val(x)
       end
f (generic function with 1 method)

julia> @inferred f(A, 1, Val(0))
Val{1}()

julia> @inferred f(A, 1, Val(1))
Val{0}()
```
In this case, the band index `0` is forwarded to the parent
(`UnitUpperTriangular`), and the `getindex` is evaluated at compile-time
to return `1` (if within bounds). Similarly, the band index `1`
corresponds to the zeros of the triangular matrix, so the value is
evaluated at compile-time. This may help with removing branches in
indexing.

Author: jishnub

src/adjtrans.jl

@@ -352,6 +352,17 @@ IndexStyle(::Type{<:AdjOrTransAbsVec}) = IndexLinear()
 @propagate_inbounds getindex(v::AdjOrTransAbsVec, ::Colon, is::AbstractArray{Int}) = wrapperop(v)(v.parent[is])
 @propagate_inbounds getindex(v::AdjOrTransAbsVec, ::Colon, ::Colon) = wrapperop(v)(v.parent[:])
 
+# band indexing
+@propagate_inbounds function getindex(A::AdjOrTransAbsMat{T}, b::BandIndex) where {T}
+    require_one_based_indexing(A)
+    wrapperop(A)(A.parent[BandIndex(-b.band, b.index)])::T
+end
+@propagate_inbounds function setindex!(A::AdjOrTransAbsMat, x, b::BandIndex)
+    require_one_based_indexing(A)
+    setindex!(A.parent, _wrapperop(A)(x), BandIndex(-b.band, b.index))
+    return A
+end
+
 # conversion of underlying storage
 convert(::Type{Adjoint{T,S}}, A::Adjoint) where {T,S} = Adjoint{T,S}(convert(S, A.parent))::Adjoint{T,S}
 convert(::Type{Transpose{T,S}}, A::Transpose) where {T,S} = Transpose{T,S}(convert(S, A.parent))::Transpose{T,S}

test/adjtrans.jl

@@ -749,4 +749,46 @@ end
     end
 end
 
+@testset "band indexing" begin
+    n = 3
+    A = UnitUpperTriangular(Matrix(reshape(1:n^2, n, n)))
+    @testset "every index" begin
+        Aadj = Adjoint(A)
+        for k in -(n-1):n-1
+            di = diagind(Aadj, k, IndexStyle(Aadj))
+            for (i,d) in enumerate(di)
+                @test Aadj[LinearAlgebra.BandIndex(k,i)] == Aadj[d]
+                if k < 0 # the adjoint is a unit lower triangular
+                    Aadj[LinearAlgebra.BandIndex(k,i)] = n^2 + i
+                    @test Aadj[d] == n^2 + i
+                end
+            end
+        end
+    end
+    @testset "inference for structured matrices" begin
+        function f(A, i, ::Val{band}) where {band}
+            x = Adjoint(A)[LinearAlgebra.BandIndex(band,i)]
+            Val(x)
+        end
+        v = @inferred f(A, 1, Val(0))
+        @test v == Val(1)
+        v = @inferred f(A, 1, Val(1))
+        @test v == Val(0)
+    end
+    @testset "non-square matrix" begin
+        r = reshape(1:6, 2, 3)
+        for d in (r, r*im)
+            @test d'[LinearAlgebra.BandIndex(1,1)] == adjoint(d[2,1])
+            @test d'[LinearAlgebra.BandIndex(-1,2)] == adjoint(d[2,3])
+            @test transpose(d)[LinearAlgebra.BandIndex(1,1)] == transpose(d[2,1])
+            @test transpose(d)[LinearAlgebra.BandIndex(-1,2)] == transpose(d[2,3])
+        end
+    end
+    @testset "block matrix" begin
+        B = reshape([[1 2; 3 4]*i for i in 1:4], 2, 2)
+        @test B'[LinearAlgebra.BandIndex(1,1)] == adjoint(B[2,1])
+        @test transpose(B)[LinearAlgebra.BandIndex(1,1)] == transpose(B[2,1])
+    end
+end
+
 end # module TestAdjointTranspose