diff --git a/src/DiskArrays.jl b/src/DiskArrays.jl index 5443bd0..d623397 100644 --- a/src/DiskArrays.jl +++ b/src/DiskArrays.jl @@ -1,6 +1,7 @@ module DiskArrays import ConstructionBase +import Base.PermutedDimsArrays: genperm using LRUCache: LRUCache, LRU diff --git a/src/permute.jl b/src/permute.jl index 9eae918..95e86ee 100644 --- a/src/permute.jl +++ b/src/permute.jl @@ -1,9 +1,16 @@ """ - PermutedDiskArray <: AbstractDiskArray + AbstractPermutedDiskArray <: AbstractDiskArray + +Abstract supertype for diskarray with permuted dimensions. +""" +abstract type AbstractPermutedDiskArray{T,N,perm,iperm,A} <: AbstractDiskArray{T,N} end + +""" + PermutedDiskArray <: AbstractPermutedDiskArray A lazily permuted disk array returned by `permutedims(diskarray, permutation)`. """ -struct PermutedDiskArray{T,N,perm,iperm,A<:AbstractArray{T,N}} <: AbstractDiskArray{T,N} +struct PermutedDiskArray{T,N,perm,iperm,A<:AbstractArray{T,N}} <: AbstractPermutedDiskArray{T,N,perm,iperm,A} parent::A end # We use PermutedDimsArray internals instead of duplicating them, @@ -41,7 +48,7 @@ function eachchunk(a::PermutedDiskArray) # Return permuted GridChunks return GridChunks(genperm(gridchunks.chunks, perm)...) end -function DiskArrays.readblock!(a::PermutedDiskArray, aout, i::OrdinalRange...) +function DiskArrays.readblock!(a::AbstractPermutedDiskArray, aout, i::OrdinalRange...) iperm = _getiperm(a) # Permute the indices inew = genperm(i, iperm) @@ -49,7 +56,7 @@ function DiskArrays.readblock!(a::PermutedDiskArray, aout, i::OrdinalRange...) DiskArrays.readblock!(parent(a), PermutedDimsArray(aout, iperm), inew...) return nothing end -function DiskArrays.writeblock!(a::PermutedDiskArray, v, i::OrdinalRange...) +function DiskArrays.writeblock!(a::AbstractPermutedDiskArray, v, i::OrdinalRange...) iperm = _getiperm(a) inew = genperm(i, iperm) # Permute the dest block and write from the true parent @@ -57,8 +64,8 @@ function DiskArrays.writeblock!(a::PermutedDiskArray, v, i::OrdinalRange...) return nothing end -_getperm(::PermutedDiskArray{<:Any,<:Any,perm}) where {perm} = perm -_getiperm(::PermutedDiskArray{<:Any,<:Any,<:Any,iperm}) where {iperm} = iperm +_getperm(::AbstractPermutedDiskArray{<:Any,<:Any,perm}) where {perm} = perm +_getiperm(::AbstractPermutedDiskArray{<:Any,<:Any,<:Any,iperm}) where {iperm} = iperm # Implementation macro diff --git a/src/reshape.jl b/src/reshape.jl index 58234ee..e36913f 100644 --- a/src/reshape.jl +++ b/src/reshape.jl @@ -1,13 +1,18 @@ import Base: _throw_dmrs -import Base.PermutedDimsArrays: genperm """ - ReshapedDiskArray <: AbstractDiskArray + AbstractReshapedDiskArray <: AbstractDiskArray + +Abstract supertype for a replacements of `Base.ReshapedArray` for `AbstractDiskArray`s` +""" +abstract type AbstractReshapedDiskArray{T,N,P,M} <: AbstractDiskArray{T,N} end + +""" + ReshapedDiskArray <: AbstractReshapedDiskArray A replacement for `Base.ReshapedArray` for disk arrays, returned by `reshape`. - Reshaping is really not trivial, because the access pattern would completely change for reshaped arrays, rectangles would not remain rectangles in the parent array. @@ -15,25 +20,26 @@ rectangles in the parent array. However, we can support the case where only singleton dimensions are added, later we could allow more special cases like joining two dimensions to one """ -struct ReshapedDiskArray{T,N,P<:AbstractArray{T},M} <: AbstractDiskArray{T,N} +struct ReshapedDiskArray{T,N,P<:AbstractArray{T},M} <: AbstractReshapedDiskArray{T,N,P,M} parent::P keepdim::NTuple{M,Int} newsize::NTuple{N,Int} end # Base methods +Base.size(r::AbstractReshapedDiskArray) = r.newsize +Base.parent(r::AbstractReshapedDiskArray) = r.parent -Base.parent(r::ReshapedDiskArray) = r.parent -Base.size(r::ReshapedDiskArray) = r.newsize +keepdim(r::AbstractReshapedDiskArray) = r.keepdim # DiskArrays interface -haschunks(a::ReshapedDiskArray) = haschunks(a.parent) -function eachchunk(a::ReshapedDiskArray{<:Any,N}) where {N} - pchunks = eachchunk(a.parent) +haschunks(a::AbstractReshapedDiskArray) = haschunks(parent(a)) +function eachchunk(a::AbstractReshapedDiskArray{<:Any,N}) where {N} + pchunks = eachchunk(parent(a)) inow::Int = 0 outchunks = ntuple(N) do idim - if in(idim, a.keepdim) + if in(idim, keepdim(a)) inow += 1 pchunks.chunks[inow] else @@ -42,14 +48,14 @@ function eachchunk(a::ReshapedDiskArray{<:Any,N}) where {N} end return GridChunks(outchunks...) end -function DiskArrays.readblock!(a::ReshapedDiskArray, aout, i::OrdinalRange...) - inew = tuple_tuple_getindex(i, a.keepdim) - DiskArrays.readblock!(a.parent, reshape(aout, map(length, inew)), inew...) +function DiskArrays.readblock!(a::AbstractReshapedDiskArray, aout, i::OrdinalRange...) + inew = tuple_tuple_getindex(i, keepdim(a)) + DiskArrays.readblock!(parent(a), reshape(aout, map(length, inew)), inew...) return nothing end -function DiskArrays.writeblock!(a::ReshapedDiskArray, v, i::OrdinalRange...) - inew = tuple_tuple_getindex(i, a.keepdim) - DiskArrays.writeblock!(a.parent, reshape(v, map(length, inew)), inew...) +function DiskArrays.writeblock!(a::AbstractReshapedDiskArray, v, i::OrdinalRange...) + inew = tuple_tuple_getindex(i, keepdim(a)) + DiskArrays.writeblock!(parent(a), reshape(v, map(length, inew)), inew...) return nothing end function reshape_disk(parent, dims) @@ -93,6 +99,6 @@ macro implement_reshape(t) end # For ambiguity -function Base._reshape(A::DiskArrays.AbstractDiskArray{<:Any,1}, dims::Tuple{Int64}) +function Base._reshape(A::AbstractDiskArray{<:Any,1}, dims::Tuple{Int64}) return reshape_disk(A, dims) end diff --git a/src/subarray.jl b/src/subarray.jl index 847752a..46faf6b 100644 --- a/src/subarray.jl +++ b/src/subarray.jl @@ -1,32 +1,44 @@ +""" + SubDiskArray <: AbstractDiskArray + +Abstract supertype for a view of an AbstractDiskArray +""" +abstract type AbstractSubDiskArray{T,N,P,I,L} <: AbstractDiskArray{T,N} end + """ SubDiskArray <: AbstractDiskArray A replacement for `Base.SubArray` for disk arrays, returned by `view`. """ -struct SubDiskArray{T,N,P,I,L} <: AbstractDiskArray{T,N} +struct SubDiskArray{T,N,P,I,L} <: AbstractSubDiskArray{T,N,P,I,L} v::SubArray{T,N,P,I,L} end # Base methods -Base.view(a::SubDiskArray, i...) = SubDiskArray(view(a.v, i...)) -Base.view(a::SubDiskArray, i::CartesianIndices) = view(a, i.indices...) -Base.size(a::SubDiskArray) = size(a.v) -Base.parent(a::SubDiskArray) = a.v.parent +subarray(a::SubDiskArray) = a.v +function Base.view(a::T, i...) where T<:AbstractSubDiskArray + basetype = Base.typename(T).wrapper + basetype(view(subarray(a), i...)) +end +Base.view(a::AbstractSubDiskArray, i::CartesianIndices) = view(a, i.indices...) +Base.size(a::AbstractSubDiskArray) = size(subarray(a)) +Base.parent(a::AbstractSubDiskArray) = parent(subarray(a)) +Base.parentindices(a::AbstractSubDiskArray) = parentindices(subarray(a)) _replace_colon(s, ::Colon) = Base.OneTo(s) _replace_colon(s, r) = r # Diskarrays.jl interface -function readblock!(a::SubDiskArray, aout, i::OrdinalRange...) - pinds = parentindices(view(a.v, i...)) - getindex_disk!(aout, parent(a.v), pinds...) +function readblock!(a::AbstractSubDiskArray, aout, i::OrdinalRange...) + pinds = parentindices(view(a, i...)) + getindex_disk!(aout, parent(a), pinds...) end -function writeblock!(a::SubDiskArray, v, i::OrdinalRange...) - pinds = parentindices(view(a.v, i...)) - setindex_disk!(parent(a.v), v, pinds...) +function writeblock!(a::AbstractSubDiskArray, v, i::OrdinalRange...) + pinds = parentindices(view(a, i...)) + setindex_disk!(parent(a), v, pinds...) end -haschunks(a::SubDiskArray) = haschunks(parent(a.v)) -eachchunk(a::SubDiskArray) = eachchunk_view(haschunks(a.v.parent), a.v) +haschunks(a::AbstractSubDiskArray) = haschunks(parent(a)) +eachchunk(a::AbstractSubDiskArray) = eachchunk_view(haschunks(parent(a)), a) function eachchunk_view(::Chunked, vv) pinds = parentindices(vv)