Support views of ApplyArray(*, ...) (#73)

* col/rowsupport can take ranges

* views of ApplyArray(*, ...) are ApplyLayout{typeof(*)}

* add call o apply array interface

* only test allocated on 1.2+

* add vcat/hcat_copyto!

* allocate for fast Matrix * Vcat

* double cache size to avoid O(n^2) complexity

* Use datasize instead of size of data, fix flatten

* tests pass on 1.0

* use copyto! instead of broadcast

* add tests

Author: dlfivefifty

Project.toml

@@ -1,6 +1,6 @@
 name = "LazyArrays"
 uuid = "5078a376-72f3-5289-bfd5-ec5146d43c02"
-version = "0.12.2"
+version = "0.12.3"
 
 [deps]
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"

src/cache.jl

@@ -92,6 +92,26 @@ end
 
 ## Array caching
 
+function resizedata!(B::CachedVector{T,Vector{T}}, n::Integer) where T<:Number
+    @boundscheck checkbounds(Bool, B, n) || throw(ArgumentError("Cannot resize beyound size of operator"))
+
+    # increase size of array if necessary
+    olddata = B.data
+    ν, = B.datasize
+    n = max(ν,n)
+    if n > length(B.data) # double memory to avoid O(n^2) growing
+        B.data = Array{T}(undef, min(2n,length(B.array)))
+        B.data[axes(olddata,1)] = olddata
+    end
+
+    inds = ν+1:n
+    B.data[inds] .= view(B.array,inds)
+
+    B.datasize = (n,)
+
+    B
+end
+
 function resizedata!(B::CachedArray{T,N,Array{T,N}},nm::Vararg{Integer,N}) where {T<:Number,N}
     @boundscheck checkbounds(Bool, B, nm...) || throw(ArgumentError("Cannot resize beyound size of operator"))
 
@@ -124,9 +144,12 @@ function convexunion(a::AbstractVector, b::AbstractVector)
     min(minimum(a),minimum(b)):max(maximum(a),maximum(b))
 end
 
-colsupport(A::CachedMatrix, i) = i ≤ size(A.data,2) ? convexunion(colsupport(A.array, i),colsupport(A.data,i)) : colsupport(A.array, i)
-colsupport(A::CachedVector, i) = convexunion(colsupport(A.array, i),colsupport(A.data,i))
-rowsupport(A::CachedMatrix, i) = i ≤ size(A.data,1) ? convexunion(rowsupport(A.array, i),rowsupport(A.data,i)) : rowsupport(A.array, i)
+colsupport(A::CachedMatrix, i) = 
+    minimum(i) ≤ A.datasize[2] ? convexunion(colsupport(A.array, i),colsupport(A.data,i) ∩ Base.OneTo(A.datasize[1])) : colsupport(A.array, i)
+colsupport(A::CachedVector, i) = 
+    convexunion(colsupport(A.array, i),colsupport(A.data,i) ∩ Base.OneTo(A.datasize[1]))
+rowsupport(A::CachedMatrix, i) = 
+    minimum(i) ≤ A.datasize[1] ? convexunion(rowsupport(A.array, i),rowsupport(A.data,i) ∩ Base.OneTo(A.datasize[2])) : rowsupport(A.array, i)
 
 Base.replace_in_print_matrix(A::CachedMatrix, i::Integer, j::Integer, s::AbstractString) =
     i in colsupport(A,j) ? s : Base.replace_with_centered_mark(s)

src/lazyapplying.jl

@@ -16,6 +16,10 @@ end
 @inline Applied{Style}(f::F, args::Args) where {Style,F,Args<:Tuple} = Applied{Style,F,Args}(f, args)
 @inline Applied{Style}(A::Applied) where Style = Applied{Style}(A.f, A.args)
 
+
+call(a) = a.f
+call(_, a) = a.f
+call(a::AbstractArray) = call(MemoryLayout(typeof(a)), a)
 arguments(a) = a.args
 arguments(_, a) = a.args
 arguments(a::AbstractArray) = arguments(MemoryLayout(typeof(a)), a)
@@ -182,8 +186,6 @@ AbstractArray{T}(A::ApplyArray{<:Any,N}) where {T,N} = ApplyArray{T,N}(A.f, map(
 AbstractArray{T,N}(A::ApplyArray{T,N}) where {T,N} = copy(A)
 AbstractArray{T,N}(A::ApplyArray{<:Any,N}) where {T,N} = ApplyArray{T,N}(A.f, map(copy,A.args)...)
 
-
-@inline Applied(A::ApplyArray) = applied(A.f, A.args...)
 @inline axes(A::ApplyArray) = axes(Applied(A))
 @inline size(A::ApplyArray) = map(length, axes(A))
 @inline copy(A::ApplyArray) = ApplyArray(A.f, map(copy,A.args)...)
@@ -230,6 +232,9 @@ MemoryLayout(::Type{Applied{Style,F,Args}}) where {Style,F,Args} =
 MemoryLayout(::Type{ApplyArray{T,N,F,Args}}) where {T,N,F,Args} = 
     applylayout(F, tuple_type_memorylayouts(Args)...)
 
+@inline Applied(A::AbstractArray) = Applied(call(A), arguments(A)...)
+@inline ApplyArray(A::AbstractArray) = ApplyArray(call(A), arguments(A)...)
+
 function show(io::IO, A::Applied) 
     print(io, "Applied(", A.f)
     for a in A.args
@@ -273,7 +278,10 @@ end
 @inline getindex(A::LazyMatrix, kr::AbstractUnitRange, jr::Colon) = lazy_getindex(A, kr, jr)
 @inline getindex(A::LazyMatrix, kr::AbstractUnitRange, jr::AbstractUnitRange) = lazy_getindex(A, kr, jr)
 
+@inline copyto!(dest::AbstractArray{T,N}, src::ApplyArray{T,N}) where {T,N} = copyto!(dest, Applied(src))
+@inline copyto!(dest::AbstractArray, src::ApplyArray) = copyto!(dest, Applied(src))    
 
-diagonallayout(::LazyLayout) = DiagonalLayout{LazyLayout}()
-diagonallayout(::ApplyLayout) = DiagonalLayout{LazyLayout}()
-
+# avoid infinite-loop
+_base_copyto!(dest::AbstractArray{T,N}, src::AbstractArray{T,N}) where {T,N} = Base.invoke(copyto!, NTuple{2,AbstractArray{T,N}}, dest, src)
+_base_copyto!(dest::AbstractArray, src::AbstractArray) = Base.invoke(copyto!, NTuple{2,AbstractArray}, dest, src)
+@inline copyto!(dest::AbstractArray, M::Applied{LazyArrayApplyStyle}) = _base_copyto!(dest, materialize(M))

src/lazybroadcasting.jl

@@ -152,8 +152,6 @@ broadcasted(::LazyArrayStyle{N}, ::typeof(*), a::AbstractArray{T,N}, b::Zeros{V,
 broadcasted(::LazyArrayStyle{N}, ::typeof(*), a::Zeros{T,N}, b::AbstractArray{V,N}) where {T,V,N} =
     broadcast(DefaultArrayStyle{N}(), *, a, b)
 
-diagonallayout(::BroadcastLayout) = DiagonalLayout{LazyLayout}()    
-
 
 ###
 # support

src/lazyconcat.jl

@@ -116,9 +116,8 @@ end
 
 ## copyto!
 # based on Base/array.jl, Base/abstractarray.jl
-
-function copyto!(dest::AbstractMatrix, V::Vcat{<:Any,2})
-    arrays = V.args
+copyto!(dest::AbstractArray, V::Vcat) = vcat_copyto!(dest, arguments(V)...)
+function vcat_copyto!(dest::AbstractMatrix, arrays...)
     nargs = length(arrays)
     nrows = size(dest,1)
     nrows == sum(a->size(a, 1), arrays) || throw(DimensionMismatch("sum of rows each matrix must equal $nrows"))
@@ -131,35 +130,13 @@ function copyto!(dest::AbstractMatrix, V::Vcat{<:Any,2})
     pos = 1
     for a in arrays
         p1 = pos+size(a,1)-1
-        dest[pos:p1, :] .= a
+        copyto!(view(dest,pos:p1, :), a)
         pos = p1+1
     end
     return dest
 end
 
-# this is repeated to avoid allocation in .=
-function copyto!(dest::AbstractMatrix, V::Vcat{<:Any,2,<:Tuple{Vararg{<:AbstractMatrix}}})
-    arrays = V.args
-    nargs = length(arrays)
-    nrows = size(dest,1)
-    nrows == sum(a->size(a, 1), arrays) || throw(DimensionMismatch("sum of rows each matrix must equal $nrows"))
-    ncols = size(dest, 2)
-    for a in arrays
-        if size(a, 2) != ncols
-            throw(DimensionMismatch("number of columns of each array must match (got $(map(x->size(x,2), A)))"))
-        end
-    end
-    pos = 1
-    for a in arrays
-        p1 = pos+size(a,1)-1
-        dest[pos:p1, :] = a
-        pos = p1+1
-    end
-    return dest
-end
-
-function copyto!(arr::AbstractVector, A::Vcat{<:Any,1,<:Tuple{Vararg{<:AbstractVector}}})
-    arrays = A.args
+function vcat_copyto!(arr::AbstractVector, arrays...)
     n = 0
     for a in arrays
         n += length(a)
@@ -168,16 +145,14 @@ function copyto!(arr::AbstractVector, A::Vcat{<:Any,1,<:Tuple{Vararg{<:AbstractV
 
     i = 0
     @inbounds for a in arrays
-        for ai in a
-            i += 1
-            arr[i] = ai
-        end
+        m = length(a)
+        copyto!(view(arr,i+1:i+m), a)
+        i += m
     end
     arr
 end
 
-function copyto!(arr::Vector{T}, A::Vcat{T,1,<:Tuple{Vararg{<:Vector{T}}}}) where T
-    arrays = A.args
+function vcat_copyto!(arr::Vector{T}, arrays::Vector{T}...) where T
     n = 0
     for a in arrays
         n += length(a)
@@ -217,8 +192,8 @@ function copyto!(arr::Vector{T}, A::Vcat{T,1,<:Tuple{Vararg{<:Vector{T}}}}) wher
     return arr
 end
 
-function copyto!(dest::AbstractMatrix, H::Hcat)
-    arrays = H.args
+copyto!(dest::AbstractMatrix, H::Hcat) = hcat_copyto!(dest, arguments(H)...)
+function hcat_copyto!(dest::AbstractMatrix, arrays...)
     nargs = length(arrays)
     nrows = size(dest, 1)
     ncols = 0
@@ -229,7 +204,7 @@ function copyto!(dest::AbstractMatrix, H::Hcat)
         ncols += (nd==2 ? size(a,2) : 1)
     end
 
-    nrows == size(H,1) || throw(DimensionMismatch("Destination rows must match"))
+    nrows == size(first(arrays),1) || throw(DimensionMismatch("Destination rows must match"))
     ncols == size(dest,2) || throw(DimensionMismatch("Destination columns must match"))
 
     pos = 1
@@ -242,22 +217,22 @@ function copyto!(dest::AbstractMatrix, H::Hcat)
     else
         for a in arrays
             p1 = pos+(isa(a,AbstractMatrix) ? size(a, 2) : 1)-1
-            dest[:, pos:p1] .= a
+            copyto!(view(dest,:, pos:p1), a)
             pos = p1+1
         end
     end
     return dest
 end
 
-function copyto!(dest::AbstractMatrix, H::Hcat{<:Any,Tuple{Vararg{<:AbstractVector}}})
+function hcat_copyto!(dest::AbstractMatrix, arrays::AbstractVector...)
     height = size(dest, 1)
-    for j = 1:length(H)
-        if length(H[j]) != height
+    for j = 1:length(arrays)
+        if length(arrays[j]) != height
             throw(DimensionMismatch("vectors must have same lengths"))
         end
     end
-    for j=1:length(H)
-        dest[i,:] .= H[j]
+    for j=1:length(arrays)
+        copyto!(view(dest,:,j), arrays[j])
     end
 
     dest
@@ -502,27 +477,36 @@ applylayout(::Type{typeof(vcat)}, ::A, ::ZerosLayout) where A = PaddedLayout{A}(
 cachedlayout(::A, ::ZerosLayout) where A = PaddedLayout{A}()
 
 
-paddeddata(A::CachedArray) = A.data
+paddeddata(A::CachedArray) = view(A.data,OneTo.(A.datasize)...)
 paddeddata(A::Vcat) = A.args[1]
 
 function ==(A::CachedVector{<:Any,<:Any,<:Zeros}, B::CachedVector{<:Any,<:Any,<:Zeros})
     length(A) == length(B) || return false
-    n = max(length(A.data), length(B.data))
+    n = max(A.datasize[1], B.datasize[1])
     resizedata!(A,n); resizedata!(B,n)
-    A.data == B.data
+    view(A.data,OneTo(n)) == view(B.data,OneTo(n))
 end
 
 # special copyto! since `similar` of a padded returns a cached
 for Typ in (:Number, :AbstractVector)
     @eval function copyto!(dest::CachedVector{T,Vector{T},<:Zeros{T,1}}, src::Vcat{<:Any,1,<:Tuple{<:$Typ,<:Zeros}}) where T
         length(src) ≤ length(dest)  || throw(BoundsError())
         a,_ = src.args
-        resizedata!(dest, length(a)) # make sure we are padded enough
-        copyto!(dest.data, a)
+        n = length(a)
+        resizedata!(dest, n) # make sure we are padded enough
+        copyto!(view(dest.data,OneTo(n)), a)
         dest
     end
 end
 
+function copyto!(dest::CachedVector{T,Vector{T},<:Zeros{T,1}}, src::CachedVector{V,Vector{V},<:Zeros{V,1}}) where {T,V}
+    length(src) ≤ length(dest)  || throw(BoundsError())
+    n = src.datasize[1]
+    resizedata!(dest, n)
+    copyto!(view(dest.data,OneTo(n)), view(src.data,OneTo(n)))
+    dest
+end
+
 struct Dot{StyleA,StyleB,ATyp,BTyp}
     A::ATyp
     B::BTyp
@@ -556,16 +540,17 @@ end
 # subarrays
 ###
 
-subarraylayout(::ApplyLayout{typeof(vcat)}, _) = 
-    ApplyLayout{typeof(vcat)}()
-subarraylayout(::ApplyLayout{typeof(hcat)}, _) = 
-    ApplyLayout{typeof(hcat)}()
+subarraylayout(::ApplyLayout{typeof(vcat)}, _) = ApplyLayout{typeof(vcat)}()
+subarraylayout(::ApplyLayout{typeof(hcat)}, _) = ApplyLayout{typeof(hcat)}()
 
 arguments(::ApplyLayout{typeof(vcat)}, V::SubArray{<:Any,2,<:Any,<:Tuple{<:Slice,<:Any}}) = 
     view.(arguments(parent(V)), Ref(:), Ref(parentindices(V)[2]))
 arguments(::ApplyLayout{typeof(hcat)}, V::SubArray{<:Any,2,<:Any,<:Tuple{<:Any,<:Slice}}) = 
     view.(arguments(parent(V)), Ref(parentindices(V)[1]), Ref(:))
 
+copyto!(dest::AbstractArray{T,N}, src::SubArray{T,N,<:Vcat{T,N}}) where {T,N} = vcat_copyto!(dest, arguments(src)...)
+copyto!(dest::AbstractMatrix{T}, src::SubArray{T,2,<:Hcat{T}}) where T = hcat_copyto!(dest, arguments(src)...)
+
 
 _vcat_lastinds(sz) = _vcat_cumsum(sz...)
 _vcat_firstinds(sz) = (1, (1 .+ most(_vcat_lastinds(sz)))...)
@@ -612,7 +597,7 @@ function sub_materialize(::ApplyLayout{typeof(vcat)}, V)
     _,jr = parentindices(V)
     for a in arguments(V)
         m = size(a,1)
-        view(ret,n+1:n+m,:) .= a
+        copyto!(view(ret,n+1:n+m,:), a)
         n += m
     end
     ret
@@ -624,9 +609,13 @@ function sub_materialize(::ApplyLayout{typeof(hcat)}, V)
     kr,_ = parentindices(V)
     for a in arguments(V)
         m = size(a,2)
-        view(ret,:,n+1:n+m) .= a
+        copyto!(view(ret,:,n+1:n+m), a)
         n += m
     end
     ret
 end
-    
+# temporarily allocate. In the future, we add a loop over arguments
+materialize!(M::MatMulMatAdd{<:AbstractColumnMajor,<:ApplyLayout{typeof(vcat)}}) = 
+    materialize!(MulAdd(M.α,M.A,Array(M.B),M.β,M.C))
+materialize!(M::MatMulVecAdd{<:AbstractColumnMajor,<:ApplyLayout{typeof(vcat)}}) = 
+    materialize!(MulAdd(M.α,M.A,Array(M.B),M.β,M.C))    

src/lazyoperations.jl

@@ -62,6 +62,10 @@ copyto!(R::AbstractMatrix, K::Kron{<:Any,2,<:Tuple{<:AbstractMatrix,<:AbstractMa
     _kron2!(R, K)
 copyto!(R::AbstractVector, K::Kron{<:Any,1,<:Tuple{<:AbstractVector,<:AbstractVector}}) =
     _kron2!(R, K)
+copyto!(R::AbstractMatrix{T}, K::Kron{T,2,<:Tuple{<:AbstractMatrix,<:AbstractMatrix}}) where T =
+    _kron2!(R, K)
+copyto!(R::AbstractVector{T}, K::Kron{T,1,<:Tuple{<:AbstractVector,<:AbstractVector}}) where T =
+    _kron2!(R, K)
 
 
 struct Diff{T, N, Arr} <: LazyArray{T, N}

src/linalg/add.jl

@@ -90,7 +90,8 @@ _view_tuple(a, b::Tuple) = view(a, b...)
 for op in (:+, :-)
     @eval begin
         subarraylayout(a::ApplyLayout{typeof($op)}, _) = a
-        arguments(a::SubArray{<:Any,N,<:ApplyArray{<:Any,N,typeof($op)}}) where N =
+        arguments(::ApplyLayout{typeof($op)}, a::SubArray) =
             _view_tuple.(arguments(parent(a)), Ref(parentindices(a)))
+        call(::ApplyLayout{typeof($op)}, a::SubArray) = $op            
     end
 end

src/linalg/diagonal.jl

@@ -1,12 +1,9 @@
-####
-# Diagonal
-####
-
-rowsupport(::Diagonal, k) = k:k
-colsupport(::Diagonal, j) = j:j
+diagonallayout(::LazyLayout) = DiagonalLayout{LazyLayout}()
+diagonallayout(::ApplyLayout) = DiagonalLayout{LazyLayout}()
+diagonallayout(::BroadcastLayout) = DiagonalLayout{LazyLayout}()    
 
-rowsupport(::DiagonalLayout, _, k) = k:k
-colsupport(::DiagonalLayout, _, j) = j:j
+rowsupport(::DiagonalLayout, _, k) = minimum(k):maximum(k)
+colsupport(::DiagonalLayout, _, j) = minimum(j):maximum(j)
 
 ###
 # Lmul