Fixes for fast triangle recurrences

Project.toml

@@ -1,7 +1,7 @@
 name = "ArrayLayouts"
 uuid = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
 authors = ["Sheehan Olver <[email protected]>"]
-version = "0.4.12"
+version = "0.5.0"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
@@ -11,7 +11,7 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [compat]
 Compat = "3.16"
-FillArrays = "0.10.1"
+FillArrays = "0.11"
 julia = "1"
 
 [extras]

src/ArrayLayouts.jl

@@ -27,7 +27,7 @@ import Base: AbstractArray, AbstractMatrix, AbstractVector,
       similar, @_gc_preserve_end, @_gc_preserve_begin,
       @nexprs, @ncall, @ntuple, tuple_type_tail,
       all, any, isbitsunion, issubset, replace_in_print_matrix, replace_with_centered_mark,
-      strides, unsafe_convert, first_index, unalias
+      strides, unsafe_convert, first_index, unalias, union
 
 import Base.Broadcast: BroadcastStyle, AbstractArrayStyle, Broadcasted, broadcasted,
                         combine_eltypes, DefaultArrayStyle, instantiate, materialize,
@@ -54,7 +54,8 @@ export materialize, materialize!, MulAdd, muladd!, Ldiv, Rdiv, Lmul, Rmul, Dot,
         DiagonalLayout, ScalarLayout, SymTridiagonalLayout, TridiagonalLayout, BidiagonalLayout,
         HermitianLayout, SymmetricLayout, TriangularLayout,
         UnknownLayout, AbstractBandedLayout, ApplyBroadcastStyle, ConjLayout, AbstractFillLayout, DualLayout,
-        colsupport, rowsupport, layout_getindex, QLayout, LayoutArray, LayoutMatrix, LayoutVector
+        colsupport, rowsupport, layout_getindex, QLayout, LayoutArray, LayoutMatrix, LayoutVector,
+        RangeCumsum
 
 struct ApplyBroadcastStyle <: BroadcastStyle end
 @inline function copyto!(dest::AbstractArray, bc::Broadcasted{ApplyBroadcastStyle})
@@ -104,6 +105,7 @@ include("factorizations.jl")
 @inline sub_materialize(_, V, _) = Array(V)
 @inline sub_materialize(L, V) = sub_materialize(L, V, axes(V))
 @inline sub_materialize(V::SubArray) = sub_materialize(MemoryLayout(V), V)
+@inline sub_materialize(V::AbstractArray) = V # Anything not a SubArray is already materialized
 
 @inline layout_getindex(A, I...) = sub_materialize(view(A, I...))
 
@@ -178,6 +180,7 @@ copyto!(dest::AbstractArray{<:Any,N}, src::SubArray{<:Any,N,<:LayoutArray}) wher
 copyto!(dest::LayoutMatrix, src::AdjOrTrans{<:Any,<:LayoutArray}) = _copyto!(dest, src)
 copyto!(dest::LayoutMatrix, src::SubArray{<:Any,2,<:AdjOrTrans{<:Any,<:LayoutArray}}) = _copyto!(dest, src)
 copyto!(dest::AbstractMatrix, src::AdjOrTrans{<:Any,<:LayoutArray}) = _copyto!(dest, src)
+copyto!(dest::SubArray{<:Any,2,<:LayoutArray}, src::AdjOrTrans{<:Any,<:LayoutArray}) = _copyto!(dest, src)
 copyto!(dest::SubArray{<:Any,2,<:LayoutMatrix}, src::SubArray{<:Any,2,<:AdjOrTrans{<:Any,<:LayoutArray}}) = _copyto!(dest, src)
 copyto!(dest::AbstractMatrix, src::SubArray{<:Any,2,<:AdjOrTrans{<:Any,<:LayoutArray}}) = _copyto!(dest, src)
 # ambiguity from sparsematrix.jl
@@ -281,6 +284,6 @@ Base.print_matrix_row(io::IO,
         axes_print_matrix_row(axes(X), io, X, A, i, cols, sep)
 
 
-
+include("cumsum.jl")
 
 end

src/cumsum.jl

@@ -0,0 +1,35 @@
+"""
+    RangeCumsum(range)
+
+represents the cumsum of a `AbstractRange`. 
+"""
+struct RangeCumsum{T, RR<:AbstractRange{T}} <: LayoutVector{T}
+    range::RR
+end
+
+size(c::RangeCumsum) = size(c.range)
+
+==(a::RangeCumsum, b::RangeCumsum) = a.range == b.range
+BroadcastStyle(::Type{<:RangeCumsum{<:Any,RR}}) where RR = BroadcastStyle(RR)
+
+
+Base.@propagate_inbounds function getindex(c::RangeCumsum{<:Any,<:AbstractRange}, k::Integer)
+    @boundscheck checkbounds(c, k)
+    r = c.range
+    k * (first(r) + r[k]) ÷ 2
+end
+Base.@propagate_inbounds function getindex(c::RangeCumsum{<:Any,<:AbstractUnitRange}, k::Integer)
+    @boundscheck checkbounds(c, k)
+    r = c.range
+    k * (2first(r) + k - 1) ÷ 2
+end
+
+Base.@propagate_inbounds getindex(c::RangeCumsum, kr::OneTo) = RangeCumsum(c.range[kr])
+
+last(r::RangeCumsum) = sum(r.range)
+diff(r::RangeCumsum) = r.range[2:end]
+
+union(a::RangeCumsum{<:Any,<:Base.OneTo}, b::RangeCumsum{<:Any,<:Base.OneTo}) = 
+    RangeCumsum(Base.OneTo(max(last(a.range),last(b.range))))
+
+sort!(a::RangeCumsum{<:Any,<:AbstractUnitRange}) = a

src/memorylayout.jl

@@ -530,10 +530,11 @@ sublayout(L::AbstractFillLayout, inds::Type) = L
 reshapedlayout(L::AbstractFillLayout, _) = L
 adjointlayout(::Type, L::AbstractFillLayout) = L
 transposelayout(L::AbstractFillLayout) = L
-# TODO: Move to FillArrays.jl
-_getindex_value(V::SubArray) = getindex_value(parent(V))
 
-sub_materialize(::AbstractFillLayout, V, ax) = Fill(_getindex_value(V), ax)
+_copyto!(_, ::AbstractFillLayout, dest::AbstractArray{<:Any,N}, src::AbstractArray{<:Any,N}) where N =
+    fill!(dest, getindex_value(src))
+
+sub_materialize(::AbstractFillLayout, V, ax) = Fill(getindex_value(V), ax)
 sub_materialize(::ZerosLayout, V, ax) = Zeros{eltype(V)}(ax)
 sub_materialize(::OnesLayout, V, ax) = Ones{eltype(V)}(ax)
 

src/mul.jl

@@ -142,6 +142,7 @@ macro veclayoutmul(Typ)
             LinearAlgebra.mul!(dest::AbstractVector, A::$Mod{<:Any,<:$Typ}, b::AbstractVector) =
                 ArrayLayouts.mul!(dest,A,b)
 
+            Base.:*(A::$Mod{<:Any,<:$Typ}, B::ArrayLayouts.LayoutMatrix) = ArrayLayouts.mul(A,B)
             Base.:*(A::$Mod{<:Any,<:$Typ}, B::AbstractMatrix) = ArrayLayouts.mul(A,B)
             Base.:*(A::$Mod{<:Any,<:$Typ}, B::AbstractVector) = ArrayLayouts.mul(A,B)
             Base.:*(A::$Mod{<:Number,<:$Typ}, B::AbstractVector{<:Number}) = ArrayLayouts.mul(A,B)

test/runtests.jl

@@ -79,6 +79,7 @@ MemoryLayout(::Type{MyVector}) = DenseColumnMajor()
         @test copyto!(MyMatrix(Array{Float64}(undef,5,5)), view(A',:,:)) == A'
         @test copyto!(Array{Float64}(undef,5,5), A') == A'
         @test copyto!(Array{Float64}(undef,5,5), view(A',:,:)) == A'
+        @test copyto!(view(MyMatrix(Array{Float64}(undef,5,5)),:,:), A') == A'
 
         @test copyto!(view(MyMatrix(Array{Float64}(undef,5,5)),:,:), view(A',:,:)) == A'
 
@@ -219,3 +220,4 @@ triangulardata(A::MyUpperTriangular) = triangulardata(A.A)
     @test_skip lmul!(U,view(copy(B),collect(1:5),1:5)) ≈ U * B
 end
 
+include("test_cumsum.jl")

test/test_cumsum.jl

@@ -0,0 +1,16 @@
+using ArrayLayouts, Test
+
+@testset "RangeCumsum" begin
+    for r in (RangeCumsum(Base.OneTo(5)), RangeCumsum(2:5), RangeCumsum(2:2:6), RangeCumsum(6:-2:1))
+        @test r == cumsum(r.range)
+        @test r == r
+        @test r .+ 1 == cumsum(r.range) .+ 1
+        @test r[Base.OneTo(3)] == r[1:3]
+        @test last(r) == r[end]
+        @test diff(r) == diff(Vector(r))
+    end
+
+    a,b = RangeCumsum(Base.OneTo(5)), RangeCumsum(Base.OneTo(6))
+    @test union(a,b) ≡ union(b,a) ≡ b
+    @test sort!(a) ≡ a
+end

test/test_layouts.jl

@@ -287,9 +287,9 @@ struct FooNumber <: Number end
         @test MemoryLayout(Fill(1,10)) == FillLayout()
         @test MemoryLayout(Ones(10)) == OnesLayout()
         @test MemoryLayout(Zeros(10)) == ZerosLayout()
-        @test MemoryLayout(view(Fill(1,10),1:3)) == FillLayout()
+        @test MemoryLayout(SubArray(Fill(1,10),(1:3,))) == FillLayout()
         @test MemoryLayout(view(Fill(1,10),1:3,1)) == FillLayout()
-        @test MemoryLayout(view(Fill(1,10),[1,3,2])) == FillLayout()
+        @test MemoryLayout(SubArray(Fill(1,10),([1,3,2],))) == FillLayout()
         @test MemoryLayout(reshape(Fill(1,10),2,5)) == FillLayout()
         @test MemoryLayout(Fill(1+0im,10)') == DualLayout{FillLayout}()
         @test MemoryLayout(Adjoint(Fill(1+0im,10,2))) == FillLayout()
@@ -298,6 +298,12 @@ struct FooNumber <: Number end
         @test layout_getindex(Fill(1,10), 1:3) ≡ Fill(1,3)
         @test layout_getindex(Ones{Int}(1,10), 1, 1:3) ≡ Ones{Int}(3)
         @test layout_getindex(Zeros{Int}(5,10,12), 1, 1:3,4:6) ≡ Zeros{Int}(3,3)
+
+        # views of Fill no longer create Sub Arrays, but are supported
+        # as there was no strong need to delete their support
+        v = SubArray(Fill(1,10),(1:3,))
+        @test ArrayLayouts.sub_materialize(v) ≡ Fill(1,3)
+        @test ArrayLayouts._copyto!(Vector{Float64}(undef,3), v) == ones(3)
     end
 
     @testset "Triangular col/rowsupport" begin