Support resizedata! and cacheddata for Vcat (#386)

* Support resizedata! and cacheddata for Vcat

* Functionality + more tests

* EOL

* Revert paddedlayout vcat

---------

Co-authored-by: Daniel VandenHeuvel <[email protected]>

Author: dlfivefifty

src/cache.jl

@@ -58,6 +58,10 @@ cache_layout(::AbstractStridedLayout, O::AbstractArray) = copy(O)
 const _cache = cache_layout # TODO: deprecate
 cacheddata(A::AbstractCachedArray) = view(A.data,OneTo.(A.datasize)...)
 
+maybe_cacheddata(A::AbstractCachedArray) = cacheddata(A)
+maybe_cacheddata(A::SubArray{<:Any,N,<:AbstractCachedArray}) where N = cacheddata(A)
+maybe_cacheddata(A) = A # no-op
+
 convert(::Type{AbstractArray{T}}, S::CachedArray{T}) where T = S
 convert(::Type{AbstractArray{T,N}}, S::CachedArray{T,N}) where {T,N} = S
 convert(::Type{AbstractArray{T}}, S::CachedArray{<:Any,N}) where {T,N} = convert(AbstractArray{T,N}, S)

src/padded.jl

@@ -128,11 +128,29 @@ function _vcat_resizedata!(::Union{AbstractPaddedLayout, DualLayout{<:PaddedRows
     any(iszero,m) || Base.checkbounds(paddeddata(B), m...)
     B
 end
+function _vcat_resizedata!(::Union{DualLayout{<:PaddedRows}, AbstractPaddedLayout}, B::Vcat{<:Any, 1}, m) # ambiguity
+    iszero(m) || Base.checkbounds(paddeddata(B), m)
+    B
+end
+
+function _vcat_resizedata!(_, B::Vcat{<:Any, 1}, n) 
+    m = n 
+    for arg in arguments(B)
+        m ≤ 0 && break
+        len = length(arg)
+        s = min(len, max(m, 0))
+        resizedata!(arg, s)
+        m -= s 
+    end
+    B
+end
 
 _vcat_resizedata!(_, B, m...) = B # by default we can't resize
 
 resizedata!(B::Vcat, m...) = _vcat_resizedata!(MemoryLayout(B), B, m...)
 
+cacheddata(B::Vcat) = Vcat(map(maybe_cacheddata, arguments(B))...)
+
 function ==(A::CachedVector{<:Any,<:Any,<:Zeros}, B::CachedVector{<:Any,<:Any,<:Zeros})
     length(A) == length(B) || return false
     n = max(A.datasize[1], B.datasize[1])

test/cachetests.jl

@@ -4,7 +4,7 @@ using LazyArrays, FillArrays, LinearAlgebra, ArrayLayouts, SparseArrays, Test
 using StaticArrays
 import LazyArrays: CachedArray, CachedMatrix, CachedVector, PaddedLayout, CachedLayout, resizedata!, zero!,
                     CachedAbstractArray, CachedAbstractVector, CachedAbstractMatrix, AbstractCachedArray, AbstractCachedMatrix,
-                    PaddedColumns
+                    PaddedColumns, cacheddata, maybe_cacheddata
 
 using ..InfiniteArrays
 using .InfiniteArrays: OneToInf
@@ -487,6 +487,15 @@ using Infinities
         B[5, 7] = 3.4
         @test A == B
     end
+
+    @testset "maybe_cacheddata" begin
+        A = cache(1:10)
+        @test maybe_cacheddata(A) === cacheddata(A)
+        B = view(A, 1:5)
+        @test maybe_cacheddata(B) === cacheddata(B)
+        C = [1, 2, 3]
+        @test maybe_cacheddata(C) === C
+    end
 end
 
 end # module

test/concattests.jl

@@ -4,7 +4,7 @@ using LazyArrays, FillArrays, LinearAlgebra, ArrayLayouts, Test, Base64
 using StaticArrays
 import LazyArrays: MemoryLayout, DenseColumnMajor, materialize!, call, paddeddata,
                     MulAdd, Applied, ApplyLayout, DefaultApplyStyle, sub_materialize, resizedata!,
-                    CachedVector, ApplyLayout, arguments, BroadcastVector, LazyLayout
+                    CachedVector, ApplyLayout, arguments, BroadcastVector, LazyLayout, cacheddata
 
 @testset "concat" begin
     @testset "Vcat" begin
@@ -199,11 +199,38 @@ import LazyArrays: MemoryLayout, DenseColumnMajor, materialize!, call, paddeddat
             @test A ≠ [1,2,4]
         end
 
-        @testset "resizedata!" begin
-            # allow emulating a cached Vector
-            a = Vcat([1,2], Zeros(8))
-            @test resizedata!(a, 2) ≡ a
-            @test_throws BoundsError resizedata!(a,3)
+        @testset "cached vcat" begin
+            @testset "resizedata!" begin
+                v = Vcat(1, cache(1:10));
+                resizedata!(v, 3);
+                @test v.args[2].datasize == (2,)
+                v = Vcat(1, [2, 4, 6], cache(7:10), cache(1:5), 50);
+                resizedata!(v, 1)
+                @test v.args[3].datasize == v.args[4].datasize == (0,)
+                resizedata!(v, 5)
+                @test v.args[3].datasize == (1,) 
+                @test v.args[4].datasize == (0,)
+                resizedata!(v, 10)
+                @test v.args[3].datasize == (4,)
+                @test v.args[4].datasize == (2,)
+                resizedata!(v, 14)
+                @test v.args[3].datasize == (4,)
+                @test v.args[4].datasize == (5,)
+                resizedata!(v, 50) # test it doesn't break for excess resize 
+                @test v.args[3].datasize == (4,)
+                @test v.args[4].datasize == (5,)
+            end
+
+            @testset "cacheddata" begin
+                v = Vcat(1, cache(1:2))
+                @test @inferred(cacheddata(v)) == [1]
+                resizedata!(v, 2)
+                @test @inferred(cacheddata(v)) == [1, 1]
+                @test cacheddata(v) isa Vcat
+            end
+
+            p =  Vcat([1,2], Zeros(4));
+            # TODO: special behaviour?
         end
 
         @testset "Axpy" begin