add abstract types for views and permuted arrays (#255)

* add abstract types

* replace a.v with subarray function

* add AbstractReshapedDiskArray

* fix reshapeddiskarray docstring

* replace field access with function calls for permuteddiskarray

* import genperm

* more type parameters in AbstractPermutedDiskArray

---------

Co-authored-by: Rafael Schouten <[email protected]>

Author: tiemvanderdeure

src/DiskArrays.jl

@@ -1,6 +1,7 @@
 module DiskArrays
 
 import ConstructionBase
+import Base.PermutedDimsArrays: genperm
 
 using LRUCache: LRUCache, LRU
 

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,24 +48,24 @@ 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)
     # Permute the dest block and read from the true parent
     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
     DiskArrays.writeblock!(parent(a), PermutedDimsArray(v, iperm), inew...)
     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
 

src/reshape.jl

@@ -1,39 +1,45 @@
 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. 
 
 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

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)