row/colsupport for OneElementMatrix (#200)

Author: jishnub

Project.toml

@@ -1,7 +1,7 @@
 name = "ArrayLayouts"
 uuid = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
 authors = ["Sheehan Olver <[email protected]>"]
-version = "1.5.2"
+version = "1.5.3"
 
 [deps]
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"

src/memorylayout.jl

@@ -641,9 +641,14 @@ rowsupport(_, A, k) = axes(A,2)
 """
     rowsupport(A, k)
 
-gives an iterator containing the possible non-zero entries in the k-th row of A.
+Return an iterator containing the column indices of the possible non-zero entries in the `k`-th row of `A`.
 """
 rowsupport(A, k) = rowsupport(MemoryLayout(A), A, k)
+"""
+    rowsupport(A)
+
+Return an iterator containing the column indices of the possible non-zero entries in `A`.
+"""
 rowsupport(A) = rowsupport(A, axes(A,1))
 
 colsupport(_, A, j) = axes(A,1)
@@ -652,9 +657,14 @@ colsupport(_, A, j) = axes(A,1)
 """
     colsupport(A, j)
 
-gives an iterator containing the possible non-zero entries in the j-th column of A.
+Return an iterator containing the row indices of the possible non-zero entries in the `j`-th column of `A`.
 """
 colsupport(A, j) = colsupport(MemoryLayout(A), A, j)
+"""
+    colsupport(A)
+
+Return an iterator containing the row indices of the possible non-zero entries in `A`.
+"""
 colsupport(A) = colsupport(A, axes(A,2))
 
 # TODO: generalise to other subarrays
@@ -668,6 +678,12 @@ colsupport(::ZerosLayout, A, _) = 1:0
 
 colsupport(::UnknownLayout, A::OneElement{<:Any,1}, _) =
     intersect(axes(A,1), A.ind[1]:A.ind[1])
+function colsupport(::UnknownLayout, A::OneElement{<:Any,2}, j)
+    intersect(axes(A,1), range(A.ind[1], length = Int(A.ind[2] ∈ j)))
+end
+function rowsupport(::UnknownLayout, A::OneElement{<:Any,2}, k)
+    intersect(axes(A,2), range(A.ind[2], length = Int(A.ind[1] ∈ k)))
+end
 
 rowsupport(::DiagonalLayout, _, k) = k
 colsupport(::DiagonalLayout, _, j) = j

src/muladd.jl

@@ -179,7 +179,7 @@ function default_blasmul!(α, A::AbstractMatrix, B::AbstractMatrix, β, C::Abstr
     jindsid = all(k -> rowsupport(B,rowsupport(A,k)) == r, colsupport(A))
 
     if jindsid
-        for j in rowsupport(B,rowsupport(A,1)), k in colsupport(A)
+        for j in r, k in colsupport(A)
             _default_blasmul_loop!(α, A, B, β, C, k, j)
         end
     else

test/test_layouts.jl

@@ -316,9 +316,21 @@ struct FooNumber <: Number end
         @test isempty(rowsupport(Zeros(5,10), 2))
 
         @testset "OneElement" begin
-            for ind in (4, 20)
+            @testset for ind in (4, 20)
                 o = OneElement(2, ind, 10)
                 @test sum(o) == sum(o[colsupport(o)])
+                @test sum(o) == sum(o[colsupport(o),rowsupport(o)])
+            end
+            @testset for ind in ((3,4), (15,20))
+                O = OneElement(2, ind, (10,10))
+                @test isempty(colsupport(O,1))
+                if ind[2] < size(O,2)
+                    @test colsupport(O,ind[2]) == ind[1]:ind[1]
+                end
+                if ind[1] < size(O,1)
+                    @test rowsupport(O,ind[1]) == ind[2]:ind[2]
+                end
+                @test sum(O) == sum(O[colsupport(O),rowsupport(O)])
             end
         end