Add missing BroadcastStyle and MemoryLayouts for cache styles (#387)

* Add missing BroadcastStyle and MemoryLayouts for cache styles

* Fix cacheddata/complex caches

* Review

* Loosen test

* Fix cached access

* StackOverflow issue

* Fix resizing where layout is a DualLayout{<:AbstractCachedLayout}

* The change

* No need to specialise on N anymore

* Broken test

Author: DanielVandH

Project.toml

@@ -1,6 +1,6 @@
 name = "LazyArrays"
 uuid = "5078a376-72f3-5289-bfd5-ec5146d43c02"
-version = "2.9"
+version = "2.9.1"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"

src/cache.jl

@@ -57,6 +57,8 @@ cache_layout(::AbstractStridedLayout, O::AbstractArray) = copy(O)
 
 const _cache = cache_layout # TODO: deprecate
 cacheddata(A::AbstractCachedArray) = view(A.data,OneTo.(A.datasize)...)
+cacheddata(A::Adjoint) = adjoint(cacheddata(parent(A)))
+cacheddata(A::Transpose) = transpose(cacheddata(parent(A)))
 
 maybe_cacheddata(A::AbstractCachedArray) = cacheddata(A)
 maybe_cacheddata(A::SubArray{<:Any,N,<:AbstractCachedArray}) where N = cacheddata(A)
@@ -333,6 +335,9 @@ MemoryLayout(C::Type{CachedArray{T,N,DAT,ARR}}) where {T,N,DAT,ARR} = cachedlayo
 
 MemoryLayout(::Type{<:AbstractCachedArray}) = GenericCachedLayout()
 
+transposelayout(::AbstractCachedLayout) = GenericCachedLayout()
+conjlayout(::Type{<:Complex}, ::AbstractCachedLayout) = GenericCachedLayout()
+
 #####
 # broadcasting
 #
@@ -345,10 +350,13 @@ CachedArrayStyle(::Val{N}) where N = CachedArrayStyle{N}()
 CachedArrayStyle{M}(::Val{N}) where {N,M} = CachedArrayStyle{N}()
 
 BroadcastStyle(::Type{<:AbstractCachedArray{<:Any,N}}) where N = CachedArrayStyle{N}()
+BroadcastStyle(::Type{<:AdjOrTrans{<:Any, <:AbstractCachedArray{<:Any,N}}}) where N = CachedArrayStyle{N}()
+BroadcastStyle(::Type{<:AdjOrTrans{<:Any, <:SubArray{<:Any,N,<:AbstractCachedArray{<:Any,M}}}}) where {N,M} = CachedArrayStyle{M}()
 BroadcastStyle(::Type{<:SubArray{<:Any,N,<:AbstractCachedArray{<:Any,M}}}) where {N,M} = CachedArrayStyle{M}()
+BroadcastStyle(::Type{<:SubArray{<:Any,N,<:AdjOrTrans{<:Any, <:AbstractCachedArray{<:Any,M}}}}) where {N,M} = CachedArrayStyle{M}()
+BroadcastStyle(::Type{<:AdjOrTrans{<:Any, <:SubArray{<:Any,N,<:AdjOrTrans{<:Any,<:AbstractCachedArray{<:Any,M}}}}}) where {N,M} = CachedArrayStyle{M}()
 BroadcastStyle(::CachedArrayStyle{N}, ::LazyArrayStyle{M}) where {N,M} = CachedArrayStyle{max(M, N)}()
 
-
 broadcasted(::AbstractLazyArrayStyle, op, A::CachedArray) = CachedArray(broadcast(op, cacheddata(A)), broadcast(op, A.array))
 layout_broadcasted(::CachedLayout, _, op, A::AbstractArray, c::Number) = CachedArray(broadcast(op, cacheddata(A), c), broadcast(op, A.array, c))
 layout_broadcasted(_, ::CachedLayout, op, c::Number, A::CachedArray) = CachedArray(broadcast(op, c, cacheddata(A)), broadcast(op, c, A.array))
@@ -394,6 +402,7 @@ function _bc_resizecacheddata!(::AbstractCachedLayout, a)
     resizedata!(a, size(a)...)
     view(cacheddata(a), axes(a)...)
 end
+_bc_resizecacheddata!(::DualLayout{ML}, a) where {ML<:AbstractCachedLayout} =  _bc_resizecacheddata!(ML(), a)
 _bc_resizecacheddata!(_, a) = a
 _bc_resizecacheddata!(a) = _bc_resizecacheddata!(MemoryLayout(a), a)
 resize_bcargs!(bc::Broadcasted{<:CachedArrayStyle}) = broadcasted(bc.f, map(_bc_resizecacheddata!, bc.args)...)

test/cachetests.jl

@@ -4,7 +4,8 @@ using LazyArrays, FillArrays, LinearAlgebra, ArrayLayouts, SparseArrays, Test
 using StaticArrays
 import LazyArrays: CachedArray, CachedMatrix, CachedVector, PaddedLayout, CachedLayout, resizedata!, zero!,
                     CachedAbstractArray, CachedAbstractVector, CachedAbstractMatrix, AbstractCachedArray, AbstractCachedMatrix,
-                    PaddedColumns, cacheddata, maybe_cacheddata
+                    PaddedColumns, cacheddata, LazyArrayStyle, maybe_cacheddata, Accumulate, CachedArrayStyle, GenericCachedLayout,
+                    AccumulateAbstractVector
 
 using ..InfiniteArrays
 using .InfiniteArrays: OneToInf
@@ -529,6 +530,18 @@ using Infinities
             copyto!(dest, src)
             @test dest == res
         end
+
+        @testset "Avoid StackOverflow for recursive CachedArrayStyles" begin
+            @test Matrix(view((1:5)', :, 1:1) .* view(Accumulate(*, 1:5)', :, 1:1)) == [1;;] # used to StackOverflow
+        end
+
+        @testset "DualLayout{<:AbstractCachedLayout}" begin
+            arg1 = view((1:100)', :, 1:10)
+            arg2 = view(AccumulateAbstractVector(*, 1:100)', :, 1:10)
+            bc = Base.Broadcast.Broadcasted(CachedArrayStyle{2}(), *, (arg1, arg2))
+            rsz_bc = LazyArrays.resize_bcargs!(bc);
+            @test rsz_bc.args[2] == view(arg2.parent.parent.data', :, 1:10)
+        end
     end
 
     @testset "maybe_cacheddata" begin
@@ -539,6 +552,28 @@ using Infinities
         C = [1, 2, 3]
         @test maybe_cacheddata(C) === C
     end
+
+    @testset "Missing BroadcastStyles/MemoryLayouts/cacheddata with CachedArrayStyles" begin
+        A = view(Accumulate(*, [1, 2, 3])', 1:1, 1:2)
+        B = view(transpose(Accumulate(*, [1, 2im, 3])), 1:1, 1:2)
+        C = Accumulate(*, [1, 2im, 3])'
+        D = transpose(Accumulate(*, [1, 2im, 3]))
+        E = view(Accumulate(*, [1, 2im, 3])', 1:1, 1:2)
+        F = view(Accumulate(*, [1, 2, 3]), 1:2)'
+        G = view(Accumulate(*, [1, 2im, 3])', 1:1, 1:2)'
+        @test all(==(CachedArrayStyle{1}()), Base.BroadcastStyle.(typeof.((A, B, C, D, E, F, G))))
+        @test all(==(GenericCachedLayout()), MemoryLayout.(typeof.((A, B, E, G))))
+        @test all(==(DualLayout{GenericCachedLayout}()), MemoryLayout.(typeof.((C, D, F))))
+        @test MemoryLayout(typeof(C)) == DualLayout{GenericCachedLayout}()
+        @test MemoryLayout(typeof(D)) == DualLayout{GenericCachedLayout}()
+        @test cacheddata(A) === view(cacheddata(parent(parent(A)))', 1:1, 1:1)
+        @test cacheddata(B) === view(transpose(cacheddata(parent(parent(B)))), 1:1, 1:1)
+        @test cacheddata(C) === cacheddata(parent(C))'
+        @test cacheddata(D) === transpose(cacheddata(parent(D)))
+        @test cacheddata(E) === view(cacheddata(parent(parent(E)))', 1:1, 1:1)
+        @test cacheddata(F) === view(cacheddata(parent(parent(F))), 1:1)'
+        @test cacheddata(G) === adjoint(view(cacheddata(parent(G)), 1:1, 1:1))
+    end
 end
 
 end # module