diagonal special cases and colsupport for subarray (#54)

* special case of permutedims for Diagonal

* add permutedims test

* v0.5.3

* Add diagonal special case, colsupport for SubArray

* Update ArrayLayouts.jl

* test on 1.6-rc1

* Increase coverage

* Update test_layoutarray.jl

Author: dlfivefifty

Project.toml

@@ -1,7 +1,7 @@
 name = "ArrayLayouts"
 uuid = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
 authors = ["Sheehan Olver <[email protected]>"]
-version = "0.5.3"
+version = "0.5.4"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"

src/ArrayLayouts.jl

@@ -132,6 +132,7 @@ macro layoutgetindex(Typ)
         ArrayLayouts.@_layoutgetindex LinearAlgebra.Hermitian{<:Any,<:$Typ}
         ArrayLayouts.@_layoutgetindex LinearAlgebra.Adjoint{<:Any,<:$Typ}
         ArrayLayouts.@_layoutgetindex LinearAlgebra.Transpose{<:Any,<:$Typ}
+        ArrayLayouts.@_layoutgetindex LinearAlgebra.SubArray{<:Any,2,<:$Typ}
     end)
 end
 
@@ -162,6 +163,29 @@ getindex(A::LayoutVector, kr::Colon) = layout_getindex(A, kr)
 getindex(A::AdjOrTrans{<:Any,<:LayoutVector}, kr::Integer, jr::Colon) = layout_getindex(A, kr, jr)
 getindex(A::AdjOrTrans{<:Any,<:LayoutVector}, kr::Integer, jr::AbstractVector) = layout_getindex(A, kr, jr)
 
+*(A::Diagonal{<:Any,<:LayoutVector}, B::Diagonal{<:Any,<:LayoutVector}) = mul(A, B)
+*(A::Diagonal{<:Any,<:LayoutVector}, B::AbstractMatrix) = mul(A, B)
+*(A::AbstractMatrix, B::Diagonal{<:Any,<:LayoutVector}) = mul(A, B)
+*(A::Diagonal{<:Any,<:LayoutVector}, B::LayoutMatrix) = mul(A, B)
+*(A::LayoutMatrix, B::Diagonal{<:Any,<:LayoutVector}) = mul(A, B)
+*(A::Diagonal{<:Any,<:LayoutVector}, B::Diagonal) = mul(A, B)
+*(A::Diagonal, B::Diagonal{<:Any,<:LayoutVector}) = mul(A, B)
+
+for Mod in (:Adjoint, :Transpose, :Symmetric, :Hermitian)
+    @eval begin
+        *(A::Diagonal{<:Any,<:LayoutVector}, B::$Mod{<:Any,<:LayoutMatrix}) = mul(A,B)
+        *(A::$Mod{<:Any,<:LayoutMatrix}, B::Diagonal{<:Any,<:LayoutVector}) = mul(A,B)
+    end
+end
+\(A::Diagonal{<:Any,<:LayoutVector}, B::Diagonal{<:Any,<:LayoutVector}) = ldiv(A, B)
+\(A::Diagonal{<:Any,<:LayoutVector}, B::AbstractMatrix) = ldiv(A, B)
+\(A::AbstractMatrix, B::Diagonal{<:Any,<:LayoutVector}) = ldiv(A, B)
+\(A::Diagonal{<:Any,<:LayoutVector}, B::LayoutMatrix) = ldiv(A, B)
+\(A::LayoutMatrix, B::Diagonal{<:Any,<:LayoutVector}) = ldiv(A, B)
+\(A::Diagonal{<:Any,<:LayoutVector}, B::Diagonal) = ldiv(A, B)
+\(A::Diagonal, B::Diagonal{<:Any,<:LayoutVector}) = ldiv(A, B)
+
+
 _copyto!(_, _, dest::AbstractArray{T,N}, src::AbstractArray{V,N}) where {T,V,N} =
     Base.invoke(copyto!, Tuple{AbstractArray{T,N},AbstractArray{V,N}}, dest, src)
 

src/memorylayout.jl

@@ -632,6 +632,7 @@ rowsupport(A) = rowsupport(A, axes(A,1))
 
 colsupport(_, A, j) = axes(A,1)
 
+
 """"
     colsupport(A, j)
 
@@ -640,6 +641,12 @@ gives an iterator containing the possible non-zero entries in the j-th column of
 colsupport(A, j) = colsupport(MemoryLayout(A), A, j)
 colsupport(A) = colsupport(A, axes(A,2))
 
+# TODO: generalise to other subarrays
+function colsupport(A::SubArray{<:Any,N,<:Any,<:Tuple{Slice,Any}}, j) where N
+    _, jr = parentindices(A)
+    colsupport(parent(A), jr[j])
+end
+
 rowsupport(::ZerosLayout, A, _) = 1:0
 colsupport(::ZerosLayout, A, _) = 1:0
 

test/test_layoutarray.jl

@@ -192,6 +192,27 @@ MemoryLayout(::Type{MyVector}) = DenseColumnMajor()
         @test B - I ≈ A - I
         @test I - B ≈ I - B
     end
+
+    @testset "Diagonal" begin
+        D = Diagonal(MyVector(randn(5)))
+        D̃ = Diagonal(Vector(D.diag))
+        B = randn(5,5)
+        B̃ = MyMatrix(B)
+        @test D*D ≈ Matrix(D)^2
+        @test_broken D^2 ≈ D*D
+        @test D*B ≈ Matrix(D)*B
+        @test B*D ≈ B*Matrix(D)
+        @test D*B̃ ≈ Matrix(D)*B̃
+        @test B̃*D ≈ B̃*Matrix(D)
+        @test D*D̃ ≈ D̃*D
+
+        @test D\D ≈ I
+        @test D\B ≈ Matrix(D)\B
+        @test B\D ≈ B\Matrix(D)
+        @test D\B̃ ≈ Matrix(D)\B̃
+        @test B̃\D ≈ B̃\Matrix(D)
+        @test D\D̃ ≈ D̃\D
+    end
 end
 
 struct MyUpperTriangular{T} <: AbstractMatrix{T}