Support cacheddata on BroadcastArray and ApplyArray

Author: DanielVandH

src/cache.jl

@@ -62,6 +62,8 @@ cacheddata(A::Transpose) = transpose(cacheddata(parent(A)))
 
 maybe_cacheddata(A::AbstractCachedArray) = cacheddata(A)
 maybe_cacheddata(A::SubArray{<:Any,N,<:AbstractCachedArray}) where N = cacheddata(A)
+maybe_cacheddata(A::Adjoint) = adjoint(maybe_cacheddata(parent(A)))
+maybe_cacheddata(A::Transpose) = transpose(maybe_cacheddata(parent(A)))
 maybe_cacheddata(A) = A # no-op
 
 convert(::Type{AbstractArray{T}}, S::CachedArray{T}) where T = S
@@ -228,9 +230,13 @@ resizedata!(B::CachedArray, mn...) = resizedata!(MemoryLayout(B.data), MemoryLay
 resizedata!(B::AbstractCachedArray, mn...) = resizedata!(MemoryLayout(B.data), UnknownLayout(), B, mn...)
 resizedata!(A, mn...) = A # don't do anything
 function resizedata!(A::AdjOrTrans, m, n)
-    m ≤ 0 || resizedata!(parent(A), n)
+    resizedata!(parent(A), n, m)
     A
 end
+function resizedata!(A::AdjOrTrans{<:Any, <:AbstractVector}, m, n)
+    resizedata!(parent(A), n)
+    A
+end 
 
 function cache_filldata!(B, inds...)
     B.data[inds...] .= view(B.array,inds...)
@@ -611,4 +617,85 @@ CachedArray(data::AbstractMatrix, array::AbstractQ) = CachedArray(data, array, s
 CachedArray(data::AbstractMatrix{T}, array::AbstractQ{T}, datasize::NTuple{2,Int}) where T =
     CachedMatrix{T,typeof(data),typeof(array)}(data, array, datasize)
 
-length(A::CachedMatrix{<:T,<:AbstractMatrix{T},<:AbstractQ{T}}) where T = prod(size(A.array))
+length(A::CachedMatrix{<:T,<:AbstractMatrix{T},<:AbstractQ{T}}) where T = prod(size(A.array))
+
+##
+# Tuples of potential cached arrays
+# Assumes that the Tuples contain arrays of the same shape
+# Some of the complication here is in making sure that we can handle things like mixing vectors and matrices that have one column
+##
+@inline _tuple_wrap(x::Int) = (x,) # Some CachedArrays mistakenly use Int instead of Tuple for vector sizes (e.g. https://github.com/JuliaApproximation/SemiclassicalOrthogonalPolynomials.jl/blob/a29007f2815134180b8433fdab46a23acfcdcd01/src/neg1c.jl#L19 and https://github.com/DanielVandH/InfiniteRandomArrays.jl/blob/a859d565f5bd8278d3f3499cd26b62916b4867e0/src/vector.jl#L18, to name a couple)
+@inline _tuple_wrap(x::Tuple) = x
+@inline _datasize(x) = size(x)
+@inline _datasize(::Number) = (1,) # Numbers have size (), but we treat them as size (1,) for resizing purposes
+@inline _datasize(x::AbstractCachedArray) = x.datasize
+@inline _datasize(x::SubArray) = _datasize(parent(x))
+@inline _datasize(x::AdjOrTrans) = reverse(_datasize(parent(x)))
+@inline _datasize(x::AdjOrTrans{<:Any, <:AbstractVector}) = (sz = only(_datasize(parent(x))); (min(1, sz), sz))
+@inline _datasizes(::Tuple{}) = ()
+@inline _datasizes(t::Tuple) = (_datasize(first(t)), _datasizes(Base.tail(t))...) # equivalent to _datasize.(t)
+
+@inline _has_vector(::Tuple{}) = false # Numbers have size () 
+@inline _has_vector(::Tuple{Any}) = true
+@inline _has_vector(::Tuple{Any, Vararg}) = false
+@inline has_vector(::Tuple{}) = false
+@inline has_vector(t::Tuple) = _has_vector(first(t)) || has_vector(Base.tail(t))
+@inline to_vector_size(sz::Tuple{Any}) = sz  # Already 1-tuple
+@inline to_vector_size(sz::Tuple{Any, Vararg}) = (prod(sz),)  # Multi-tuple -> 1-tuple
+
+function max_datasize(sizes::Tuple)
+    if has_vector(sizes)
+        normalized = map(to_vector_size, sizes)
+        return reduce(@inline((a, b) -> max.(a, b)), normalized)
+    else
+        return reduce(@inline((a, b) -> max.(a, b)), sizes)
+    end
+end
+
+@inline _expand_size(sz::Tuple{Any}, ::Tuple{Any}) = sz
+@inline _expand_size(sz::Tuple{Any}, ::Tuple{Any, Vararg}) = (only(sz), 1)
+@inline _expand_size(sz::Tuple{Any, Vararg}, ::Tuple{Any, Vararg}) = sz
+
+@inline _effective_ndims(sz::Int) = 1 
+@inline function _effective_ndims(sz::Tuple)
+    length(sz) == 1 && return 1
+    length(sz) == 2 && minimum(sz) <= 1 && return 1
+    return length(sz)
+end
+
+@inline _all_same_ndims(::Tuple{}) = true
+@inline _all_same_ndims(t::Tuple{Any}) = true
+@inline function _all_same_ndims(t::Tuple)
+    first_sz = _arraysize(first(t))
+    first_ndim = _effective_ndims(first_sz)
+    _check_same_ndims_args(first_ndim, Base.tail(t))
+end
+
+@inline _check_same_ndims_args(::Int, ::Tuple{}) = true
+@inline function _check_same_ndims_args(expected_ndim::Int, t::Tuple)
+    sz = _arraysize(first(t))
+    _effective_ndims(sz) == expected_ndim && _check_same_ndims_args(expected_ndim, Base.tail(t))
+end
+
+@inline _arraysize(x::Number) = (1,)
+@inline _arraysize(x) = size(x)
+
+function conforming_resize!(args::Tuple)
+    isempty(args) && return args
+    if !_all_same_ndims(args)
+        throw(ArgumentError("Cannot conform arrays with incompatible dimensions: $(map(_arraysize, args))"))
+    end
+    sizes = _datasizes(args)
+    sz = max_datasize(sizes)
+    _resize_each!(args, sizes, sz)
+    return args
+end  
+
+@inline _resize_each!(::Tuple{}, ::Tuple{}, sz) = nothing
+@inline function _resize_each!(args::Tuple, sizes::Tuple, sz)
+    arg = first(args)
+    orig_sz = first(sizes)
+    target_sz = _expand_size(sz, orig_sz)
+    resizedata!(arg, target_sz...)
+    _resize_each!(Base.tail(args), Base.tail(sizes), sz)
+end

src/lazyapplying.jl

@@ -228,7 +228,11 @@ AbstractArray{T,N}(A::ApplyArray{<:Any,N}) where {T,N} = ApplyArray{T,N}(A.f, ma
 
 @inline applied_axes(f, args...) = map(oneto, applied_size(f, args...))
 
-
+function cacheddata(A::ApplyArray{T, N, F}) where {T, N, F}
+    args = arguments(A) 
+    conforming_resize!(args)
+    ApplyArray{T, N}(A.f, map(maybe_cacheddata, args)...)
+end
 
 # immutable arrays don't need to copy.
 # Some special cases like vcat overload setindex! and therefore
@@ -329,9 +333,6 @@ function show(io::IO, A::Applied)
     print(io, ')')
 end
 
-# BroadcastStyle(::Type{<:LinearAlgebra.QRCompactWYQ}) = DefaultArrayStyle{2}()
-# BroadcastStyle(::Type{<:LinearAlgebra.AdjointQ}) = DefaultArrayStyle{2}()
-
 applybroadcaststyle(::Type{<:AbstractArray{<:Any,N}}, _2) where N = DefaultArrayStyle{N}()
 applybroadcaststyle(::Type{<:AbstractArray{<:Any,N}}, ::AbstractLazyLayout) where N = LazyArrayStyle{N}()
 applybroadcaststyle(::Type{<:ApplyArray{<:Any,N,<:Any,Args}}, ::AbstractLazyLayout) where {N,Args<:Tuple} = result_style(LazyArrayStyle{N}(), tuple_type_broadcastlayout(Args))

src/lazybroadcasting.jl

@@ -122,6 +122,12 @@ sub_materialize(::BroadcastLayout, A) = converteltype(eltype(A), sub_materialize
 
 copy(bc::Broadcasted{<:AbstractLazyArrayStyle}) = BroadcastArray(bc)
 
+function cacheddata(A::BroadcastArray{T, N, F}) where {T, N, F}
+    args = arguments(A) 
+    conforming_resize!(args)
+    BroadcastArray{T, N}(A.f, map(maybe_cacheddata, args)...)
+end
+
 # BroadcastArray are immutable
 copy(bc::BroadcastArray) = bc
 map(::typeof(copy), bc::BroadcastArray) = bc

src/padded.jl

@@ -149,8 +149,6 @@ _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

@@ -552,6 +552,11 @@ using Infinities
         @test maybe_cacheddata(B) === cacheddata(B)
         C = [1, 2, 3]
         @test maybe_cacheddata(C) === C
+
+        v = cache(1:10)'
+        @test maybe_cacheddata(v) === cacheddata(parent(v))'
+        v = transpose(cache(1:10))
+        @test maybe_cacheddata(v) === transpose(cacheddata(parent(v)))
     end
 
     @testset "Missing BroadcastStyles/MemoryLayouts/cacheddata with CachedArrayStyles" begin
@@ -587,6 +592,149 @@ using Infinities
         @test BroadcastStyle(typeof(Hcat(d, (1:10)))) == CachedArrayStyle{2}()
         @test BroadcastStyle(typeof(Hcat((1:10), d))) == CachedArrayStyle{2}()
     end
-end
+
+    @testset "Enforce same-size arguments for cacheddata" begin
+        @testset "max_datasize" begin
+            @test LazyArrays._datasize(1:10) == (10,)
+            x = cache(1:10)
+            @test LazyArrays._datasize(x) == (0,)
+            resizedata!(x, 3)
+            @test LazyArrays._datasize(x) == (3,)
+            @test LazyArrays._datasize(view(x, 3:5)) == (3,)
+            @test LazyArrays._datasize(transpose(x)) == (1, 3)
+            
+            arr = (1:10, cache(1:10))
+            @test LazyArrays.max_datasize(LazyArrays._datasizes(arr)) == (10,)
+            
+            arr = (cache(1:10), cache(1:10))
+            resizedata!(arr[1], 3)
+            @test LazyArrays.max_datasize(LazyArrays._datasizes(arr)) == (3,)
+            resizedata!(arr[2], 7)
+            @test LazyArrays.max_datasize(LazyArrays._datasizes(arr)) == (7,)
+            
+            arr = (rand(10, 5), LazyArrays.CachedArray(rand(10, 5)))
+            @test LazyArrays.max_datasize(LazyArrays._datasizes(arr)) == (10, 5)
+            
+            arr = (LazyArrays.CachedArray(rand(10, 5)), LazyArrays.CachedArray(rand(10, 5)))
+            @test LazyArrays.max_datasize(LazyArrays._datasizes(arr)) == (0, 0)
+            resizedata!(arr[1], 3, 4)
+            @test LazyArrays.max_datasize(LazyArrays._datasizes(arr)) == (3, 4)
+            resizedata!(arr[2], 2, 5);
+            @test LazyArrays.max_datasize(LazyArrays._datasizes(arr)) == (3, 5) 
+
+            arr = (1, [1, 2])
+            @test LazyArrays._datasize(1) == (1, )
+            @test LazyArrays.max_datasize(LazyArrays._datasizes(arr)) == (2,)
+        end
+
+        @testset "conforming_resize!" begin
+            args = (cache(1:10), cache(1:10));
+            LazyArrays.conforming_resize!(args);
+            @test LazyArrays._datasizes(args) == ((0,), (0,));
+            LazyArrays.resizedata!(args[2], 4);
+            LazyArrays.conforming_resize!(args);
+            @test LazyArrays._datasizes(args) == ((4,), (4,));
+            
+            args = (1:10, cache(1:10));
+            LazyArrays.conforming_resize!(args);
+            @test LazyArrays._datasizes(args) == ((10,), (10,));
+            
+            args = (cache(1:10)', cache(1:10)');
+            LazyArrays.conforming_resize!(args);
+            @test LazyArrays._datasizes(args) == ((0, 0), (0, 0));
+            LazyArrays.resizedata!(args[1], 1, 3);
+            LazyArrays.conforming_resize!(args);
+            @test LazyArrays._datasizes(args) == ((1, 3), (1, 3));
+            
+            args = (cache(1:10)', LazyArrays.CachedArray(rand(1, 10)));
+            @test LazyArrays._datasizes(args) == ((0, 0), (0, 0));
+            LazyArrays.conforming_resize!(args);
+            @test LazyArrays._datasizes(args) == ((0, 0), (0, 0));
+            LazyArrays.resizedata!(args[1], 1, 4);
+            LazyArrays.conforming_resize!(args);
+            @test LazyArrays._datasizes(args) == ((1, 4), (1, 4));
+            LazyArrays.resizedata!(args[2], 1, 6)
+            LazyArrays.conforming_resize!(args);
+            @test LazyArrays._datasizes(args) == ((1, 6), (1, 6));
+            
+            args = (cache(1:10), LazyArrays.CachedArray(rand(1, 10))');
+            LazyArrays.resizedata!(args[1], 4);
+            LazyArrays.conforming_resize!(args);
+            @test LazyArrays._datasizes(args) == ((4,), (4, 1));
+            
+            args = (cache(1:10), LazyArrays.CachedArray(rand(1, 10))', 1:10);
+            LazyArrays.conforming_resize!(args);
+            @test LazyArrays._datasizes(args) == ((10,), (10, 1), (10,));
+
+            args = (1, cache(1:2));
+            LazyArrays.conforming_resize!(args);
+            @test LazyArrays._datasizes(args) == ((1,), (1,)) # because scalars are treated as size (1,)
+
+            args = (rand(10, 10), LazyArrays.CachedArray(rand(10, 10)));
+            LazyArrays.conforming_resize!(args);
+            @test LazyArrays._datasizes(args) == ((10, 10), (10, 10));
+
+            args = (LazyArrays.CachedArray(rand(10, 10)), LazyArrays.CachedArray(rand(10, 10)));
+            LazyArrays.conforming_resize!(args);
+            @test LazyArrays._datasizes(args) == ((0, 0), (0, 0));
+            LazyArrays.resizedata!(args[1], 3, 4);
+            LazyArrays.conforming_resize!(args);
+            @test LazyArrays._datasizes(args) == ((3, 4), (3, 4));
+            LazyArrays.resizedata!(args[2], 5, 6);
+            LazyArrays.conforming_resize!(args);
+            @test LazyArrays._datasizes(args) == ((5, 6), (5, 6));
+
+            @testset "conforming_resize! dimension mismatch" begin
+                args = (cache(1:10), cache(1:10))
+                @test LazyArrays.conforming_resize!(args) === args  
+                
+                args = (cache(1:10), LazyArrays.CachedArray(rand(10, 5)))
+                @test_throws ArgumentError LazyArrays.conforming_resize!(args)
+                
+                args = (LazyArrays.CachedArray(rand(3, 4)), LazyArrays.CachedArray(rand(5, 2)))
+                @test LazyArrays.conforming_resize!(args) === args  
+                
+                args = (1, cache(1:10))
+                @test LazyArrays.conforming_resize!(args) === args 
+                
+                args = (cache(1:10), LazyArrays.CachedArray(rand(2, 3)), reshape(cache(1:8), 2, 2, 2))
+                @test_throws ArgumentError LazyArrays.conforming_resize!(args)
+                
+                args = (reshape(cache(1:8), 2, 2, 2), reshape(cache(1:27), 3, 3, 3))
+                @test LazyArrays.conforming_resize!(args) === args  
+                
+                args = ()
+                @test LazyArrays.conforming_resize!(args) === args
+                
+                args = (cache(1:10),)
+                @test LazyArrays.conforming_resize!(args) === args
+            end
+        end
+    end
+
+    @testset "cacheddata for ApplyArray and BroadcastArray" begin
+        x = ApplyArray(+, 1:10, cache(11:20));
+        @test cacheddata(x) == ApplyArray(+, 1:10, 11:20) 
+        @test Base.Broadcast.BroadcastStyle(typeof(cacheddata(x))) == LazyArrays.LazyArrayStyle{1}()
+
+        x = BroadcastVector(*, 1:10, cache(1:10));
+        @test cacheddata(x) == BroadcastVector(*, 1:10, 1:10)
+        @test Base.Broadcast.BroadcastStyle(typeof(cacheddata(x))) == LazyArrays.LazyArrayStyle{1}()
+
+        x = ApplyArray(+, cache(1:10), cache(11:20));
+        @test cacheddata(x) == ApplyArray(+, 1:0, 1:0)
+        @test Base.Broadcast.BroadcastStyle(typeof(cacheddata(x))) == LazyArrays.LazyArrayStyle{1}()
+        LazyArrays.resizedata!(x.args[1], 3)
+        @test cacheddata(x) == ApplyArray(+, 1:3, 11:13)
+        @test Base.Broadcast.BroadcastStyle(typeof(cacheddata(x))) == LazyArrays.LazyArrayStyle{1}()
+
+        x = BroadcastVector(*, cache(1:10), cache(11:20));
+        @test cacheddata(x) == BroadcastVector(*, 1:0, 1:0)
+        @test Base.Broadcast.BroadcastStyle(typeof(cacheddata(x))) == LazyArrays.LazyArrayStyle{1}()
+        LazyArrays.resizedata!(x.args[1], 4)
+        @test cacheddata(x) == BroadcastVector(*, 1:4, 11:14)
+        @test Base.Broadcast.BroadcastStyle(typeof(cacheddata(x))) == LazyArrays.LazyArrayStyle{1}()
+    end
+end 
 
 end # module

test/concattests.jl

@@ -224,11 +224,14 @@ import Base.Broadcast: BroadcastStyle
             end
 
             @testset "cacheddata" begin
-                v = Vcat(1, cache(1:2))
-                @test @inferred(cacheddata(v)) == [1]
-                resizedata!(v, 2)
+                v = Vcat(1, cache(1:2));
                 @test @inferred(cacheddata(v)) == [1, 1]
+                resizedata!(v, 3)
+                @test @inferred(cacheddata(v)) == [1, 1, 2]
                 @test cacheddata(v) isa Vcat
+
+                v = Vcat((1:10)', cache(11:20)');
+                @test @inferred(cacheddata(v)) == Vcat((1:10)', (11:20)')
             end
 
             p =  Vcat([1,2], Zeros(4));