Use is_forwarding_wrapper trait instead of parent_type(T) <: T (#296)

Using the heuristic parent_type(T) <: T sometimes gives incorrect results and may hide holes in the interface (as I found when making these changes).

Some problems I found:

* is_increasing(stride_rank(A)) didn't account for A having zero dimensions
* is_dense(A) didn't know how to handle when dense_dims(A) returned nothing.
* contiguous_axis(::Type{<:ReshapedArray}) would sometimes return nothing (meaning the contiguous axis is unknown) instead of StaticInt(-1) when we actually knew there couldn't be a contiguous axis because it didn't exist in the parent array.
* (known)_dimnames(A) was often returning incorrect results for ReinterpretArray and ReshapedArray.

Author: Tokazama

docs/src/index.md

@@ -34,15 +34,19 @@ Creating an array type with unique behavior in Julia is often accomplished by cr
 This allows the new array type to inherit functionality by redirecting methods to the parent array (e.g., `Base.size(x::Wrapper) = size(parent(x))`).
 Generic design limits the need to define an excessive number of methods like this.
 However, methods used to describe a type's traits often need to be explicitly defined for each trait method.
-`ArrayInterface` assists with this by providing information about the parent type using [`ArrayInterface.parent_type`](@ref).
-By default `ArrayInterface.parent_type(::Type{T})` returns `T` (analogous to `Base.parent(x) = x`).
-If any type other than `T` is returned we assume `T` wraps a parent structure, so methods know to unwrap instances of `T`.
-It is also assumed that if `T` has a parent type `Base.parent` is defined.
+If the the underlying data and access to it are unchanged by it's wrapper the [`ArrayInterface.is_forwarding_wrapper`](@ref) trait can signal to other trait methods to access its parent data structure.
+Supporting this for a new type only requires defines these methods:
+
+```julia
+ArrayInterface.is_forwarding_wrapper(::Type{<:NewType}) = true
+ArrayInterface.parent_type(::Type{<:NewType}) = NewTypeParent
+Base.parent(x::NewType) = x.parent
+```
 
 For those authoring new trait methods, this may change the default definition from `has_trait(::Type{T}) where {T} = false`, to:
 ```julia
 function has_trait(::Type{T}) where {T}
-    if parent_type(T) <:T
+    if is_forwarding_wrapper(T)
         return false
     else
         return has_trait(parent_type(T))

src/ArrayInterface.jl

@@ -5,7 +5,7 @@ import ArrayInterfaceCore: allowed_getindex, allowed_setindex!, aos_to_soa, buff
     parent_type, fast_matrix_colors, findstructralnz, has_sparsestruct,
     issingular, isstructured, matrix_colors, restructure, lu_instance,
     safevec, zeromatrix, ColoringAlgorithm,
-    fast_scalar_indexing, parameterless_type, ndims_index, is_splat_index
+    fast_scalar_indexing, parameterless_type, ndims_index, is_splat_index, is_forwarding_wrapper
 
 # ArrayIndex subtypes and methods
 import ArrayInterfaceCore: ArrayIndex, MatrixIndex, VectorIndex, BidiagonalIndex, TridiagonalIndex
@@ -34,6 +34,8 @@ using LinearAlgebra
 
 import Compat
 
+n_of_x(::StaticInt{N}, x::X) where {N,X} = ntuple(Compat.Returns(x), Val{N}())
+
 @generated function merge_tuple_type(::Type{X}, ::Type{Y}) where {X<:Tuple,Y<:Tuple}
     Tuple{X.parameters...,Y.parameters...}
 end
@@ -48,7 +50,7 @@ Base.size(A::AbstractArray2) = map(Int, ArrayInterface.size(A))
 Base.size(A::AbstractArray2, dim) = Int(ArrayInterface.size(A, dim))
 
 function Base.axes(A::AbstractArray2)
-    !(parent_type(A) <: typeof(A)) && return ArrayInterface.axes(parent(A))
+    is_forwarding_wrapper(A) && return ArrayInterface.axes(parent(A))
     throw(ArgumentError("Subtypes of `AbstractArray2` must define an axes method"))
 end
 function Base.axes(A::AbstractArray2, dim::Union{Symbol,StaticSymbol})
@@ -62,11 +64,7 @@ end
 Base.strides(A::AbstractArray2, dim) = Int(ArrayInterface.strides(A, dim))
 
 function Base.IndexStyle(::Type{T}) where {T<:AbstractArray2}
-    if parent_type(T) <: T
-        return IndexCartesian()
-    else
-        return IndexStyle(parent_type(T))
-    end
+    is_forwarding_wrapper(T) ? IndexStyle(parent_type(T)) : IndexCartesian()
 end
 
 function Base.length(A::AbstractArray2)

src/axes.jl

@@ -23,10 +23,10 @@ end
 axes_types(x) = axes_types(typeof(x))
 axes_types(::Type{T}) where {T<:Array} = NTuple{ndims(T),OneTo{Int}}
 @inline function axes_types(::Type{T}) where {T}
-    if parent_type(T) <: T
-        return NTuple{ndims(T),OptionallyStaticUnitRange{One,Int}}
-    else
+    if is_forwarding_wrapper(T)
         return axes_types(parent_type(T))
+    else
+        return NTuple{ndims(T),OptionallyStaticUnitRange{One,Int}}
     end
 end
 axes_types(::Type{<:LinearIndices{N,R}}) where {N,R} = R
@@ -212,7 +212,7 @@ end
 
 Base.keys(x::LazyAxis) = keys(parent(x))
 
-Base.IndexStyle(::Type{<:LazyAxis}) = IndexStyle(parent_type(T))
+Base.IndexStyle(T::Type{<:LazyAxis}) = IndexStyle(parent_type(T))
 
 ArrayInterfaceCore.can_change_size(@nospecialize T::Type{<:LazyAxis}) = can_change_size(fieldtype(T, :parent))
 

src/dimensions.jl

@@ -35,6 +35,7 @@ function is_increasing(perm::Tuple{StaticInt{X},StaticInt{Y}}) where {X, Y}
     end
 end
 is_increasing(::Tuple{StaticInt{X}}) where {X} = True()
+is_increasing(::Tuple{}) = True()
 
 """
     from_parent_dims(::Type{T}) -> Tuple{Vararg{Union{Int,StaticInt}}}
@@ -164,13 +165,53 @@ have a name.
 """
 @inline known_dimnames(x, dim) = _known_dimname(known_dimnames(x), canonicalize(dim))
 known_dimnames(x) = known_dimnames(typeof(x))
-known_dimnames(::Type{T}) where {T} = _known_dimnames(T, parent_type(T))
-_known_dimnames(::Type{T}, ::Type{T}) where {T} = _unknown_dimnames(Base.IteratorSize(T))
-_unknown_dimnames(::Base.HasShape{N}) where {N} = ntuple(Compat.Returns(:_), Val(N))
+function known_dimnames(@nospecialize T::Type{<:VecAdjTrans})
+    (:_, getfield(known_dimnames(parent_type(T)), 1))
+end
+function known_dimnames(@nospecialize T::Type{<:Union{MatAdjTrans,PermutedDimsArray,SubArray}})
+    eachop(_inbounds_known_dimname, to_parent_dims(T), known_dimnames(parent_type(T)))
+end
+function known_dimnames(::Type{<:ReinterpretArray{T,N,S,A,IsReshaped}}) where {T,N,S,A,IsReshaped}
+    pnames = known_dimnames(A)
+    if IsReshaped
+        if sizeof(S) === sizeof(T)
+            return pnames
+        elseif sizeof(S) > sizeof(T)
+            return (:_, pnames...)
+        else
+            return tail(pnames)
+        end
+    else
+        return pnames
+    end
+end
+
+@inline function known_dimnames(@nospecialize T::Type{<:Base.ReshapedArray})
+    if ndims(T) === ndims(parent_type(T))
+        return known_dimnames(parent_type(T))
+    elseif ndims(T) > ndims(parent_type(T))
+        return (known_dimnames(parent_type(T))..., n_of_x(StaticInt(ndims(T) - ndims(parent_type(T))), :_)...)
+    else
+        return n_of_x(StaticInt(ndims(T)), :_)
+    end
+end
+@inline function known_dimnames(::Type{T}) where {T}
+    if is_forwarding_wrapper(T)
+        return known_dimnames(parent_type(T))
+    else
+        return _unknown_dimnames(Base.IteratorSize(T))
+    end
+end
+
+_unknown_dimnames(::Base.HasShape{N}) where {N} = n_of_x(StaticInt(N), :_)
 _unknown_dimnames(::Any) = (:_,)
+
+#=
+_known_dimnames(::Type{T}, ::Type{T}) where {T} = _unknown_dimnames(Base.IteratorSize(T))
 function _known_dimnames(::Type{C}, ::Type{P}) where {C,P}
     eachop(_inbounds_known_dimname, to_parent_dims(C), known_dimnames(P))
 end
+=#
 @inline function _known_dimname(x::Tuple{Vararg{Any,N}}, dim::CanonicalInt) where {N}
     # we cannot have `@boundscheck`, else this will depend on bounds checking being enabled
     (dim > N || dim < 1) && return :_
@@ -186,11 +227,41 @@ Return the names of the dimensions for `x`. `:_` is used to indicate a dimension
 have a name.
 """
 @inline dimnames(x, dim) = _dimname(dimnames(x), canonicalize(dim))
-@inline dimnames(x) = _dimnames(has_parent(x), x)
-@inline function _dimnames(::True, x)
-    eachop(_inbounds_dimname, to_parent_dims(x), dimnames(parent(x)))
+@inline function dimnames(x::Union{MatAdjTrans,PermutedDimsArray,SubArray})
+    eachop(_inbounds_known_dimname, to_parent_dims(x), dimnames(parent(x)))
+end
+dimnames(x::VecAdjTrans) = (static(:_), getfield(dimnames(parent(x)), 1))
+@inline function dimnames(x::ReinterpretArray{T,N,S,A,IsReshaped}) where {T,N,S,A,IsReshaped}
+    pnames = dimnames(parent(x))
+    if IsReshaped
+        if sizeof(S) === sizeof(T)
+            return pnames
+        elseif sizeof(S) > sizeof(T)
+            return (static(:_), pnames...)
+        else
+            return tail(pnames)
+        end
+    else
+        return pnames
+    end
+end
+@inline function dimnames(x::Base.ReshapedArray)
+    p = parent(x)
+    if ndims(x) === ndims(p)
+        return dimnames(p)
+    elseif ndims(x) > ndims(p)
+        return (dimnames(p)..., n_of_x(StaticInt(ndims(x) - ndims(p)), static(:_))...)
+    else
+        return n_of_x(StaticInt(ndims(x)), static(:_))
+    end
+end
+@inline function dimnames(x::X) where {X}
+    if is_forwarding_wrapper(X)
+        return dimnames(parent(x))
+    else
+        return n_of_x(StaticInt(ndims(x)), static(:_))
+    end
 end
-_dimnames(::False, x) = ntuple(_->static(:_), Val(ndims(x)))
 @inline function _dimname(x::Tuple{Vararg{Any,N}}, dim::CanonicalInt) where {N}
     # we cannot have `@boundscheck`, else this will depend on bounds checking being enabled
     # for calls such as `dimnames(view(x, :, 1, :))`

src/indexing.jl

@@ -318,14 +318,14 @@ end
 
 ## unsafe_getindex ##
 function unsafe_getindex(a::A) where {A}
-    parent_type(A) <: A && throw(MethodError(unsafe_getindex, (A,)))
+    is_forwarding_wrapper(A) || throw(MethodError(unsafe_getindex, (A,)))
     unsafe_getindex(parent(a))
 end
 
 # TODO Need to manage index transformations between nested layers of arrays
 function unsafe_getindex(a::A, i::CanonicalInt) where {A}
     if IndexStyle(A) === IndexLinear()
-        parent_type(A) <: A && throw(MethodError(unsafe_getindex, (A, i)))
+        is_forwarding_wrapper(A) || throw(MethodError(unsafe_getindex, (A, i)))
         return unsafe_getindex(parent(a), i)
     else
         return unsafe_getindex(a, _to_cartesian(a, i)...)
@@ -335,7 +335,7 @@ function unsafe_getindex(a::A, i::CanonicalInt, ii::Vararg{CanonicalInt}) where
     if IndexStyle(A) === IndexLinear()
         return unsafe_getindex(a, _to_linear(a, (i, ii...)))
     else
-        parent_type(A) <: A && throw(MethodError(unsafe_getindex, (A, i)))
+        is_forwarding_wrapper(A) || throw(MethodError(unsafe_getindex, (A, i)))
         return unsafe_getindex(parent(a), i, ii...)
     end
 end
@@ -415,13 +415,13 @@ end
 
 ## unsafe_setindex! ##
 function unsafe_setindex!(a::A, v) where {A}
-    parent_type(A) <: A && throw(MethodError(unsafe_setindex!, (A, v)))
+    is_forwarding_wrapper(A) || throw(MethodError(unsafe_setindex!, (A, v)))
     return unsafe_setindex!(parent(a), v)
 end
 # TODO Need to manage index transformations between nested layers of arrays
 function unsafe_setindex!(a::A, v, i::CanonicalInt) where {A}
     if IndexStyle(A) === IndexLinear()
-        parent_type(A) <: A && throw(MethodError(unsafe_setindex!, (A, v, i)))
+        is_forwarding_wrapper(A) || throw(MethodError(unsafe_setindex!, (A, v, i)))
         return unsafe_setindex!(parent(a), v, i)
     else
         return unsafe_setindex!(a, v, _to_cartesian(a, i)...)
@@ -431,7 +431,7 @@ function unsafe_setindex!(a::A, v, i::CanonicalInt, ii::Vararg{CanonicalInt}) wh
     if IndexStyle(A) === IndexLinear()
         return unsafe_setindex!(a, v, _to_linear(a, (i, ii...)))
     else
-        parent_type(A) <: A && throw(MethodError(unsafe_setindex!, (A, v, i, ii...)))
+        is_forwarding_wrapper(A) || throw(MethodError(unsafe_setindex!, (A, v, i, ii...)))
         return unsafe_setindex!(parent(a), v, i, ii...)
     end
 end

src/size.jl

@@ -16,7 +16,13 @@ julia> ArrayInterface.size(A)
 (static(3), static(4))
 ```
 """
-size(a::A) where {A} = _maybe_size(Base.IteratorSize(A), a)
+@inline function size(a::A) where {A}
+    if is_forwarding_wrapper(A)
+        return size(parent(a))
+    else
+        return _maybe_size(Base.IteratorSize(A), a)
+    end
+end
 size(a::Base.Broadcast.Broadcasted) = map(length, axes(a))
 
 _maybe_size(::Base.HasShape{N}, a::A) where {N,A} = map(length, axes(a))
@@ -56,15 +62,15 @@ end
 size(a, dim) = size(a, to_dims(a, dim))
 size(a::Array, dim::CanonicalInt) = Base.arraysize(a, convert(Int, dim))
 function size(a::A, dim::CanonicalInt) where {A}
-    if parent_type(A) <: A
+    if is_forwarding_wrapper(A)
+        return size(parent(a), dim)
+    else
         len = known_size(A, dim)
         if len === nothing
             return Int(length(axes(a, dim)))
         else
             return StaticInt(len)
         end
-    else
-        return size(a)[dim]
     end
 end
 function size(A::SubArray, dim::CanonicalInt)
@@ -86,6 +92,17 @@ compile time. If a dimension does not have a known size along a dimension then `
 returned in its position.
 """
 known_size(x) = known_size(typeof(x))
+@inline function known_size(::Type{T}) where {T}
+    if is_forwarding_wrapper(T)
+        return known_size(parent_type(T))
+    else
+        return _maybe_known_size(Base.IteratorSize(T), T)
+    end
+end
+function _maybe_known_size(::Base.HasShape{N}, ::Type{T}) where {N,T}
+    eachop(_known_size, nstatic(Val(N)), axes_types(T))
+end
+_maybe_known_size(::Base.IteratorSize, ::Type{T}) where {T} = (known_length(T),)
 function known_size(::Type{T}) where {T<:AbstractRange}
     (_range_length(known_first(T), known_step(T), known_last(T)),)
 end
@@ -109,12 +126,6 @@ end
     dynamic(reduce_tup(_promote_shape, eachop(_unzip_size, nstatic(Val(known_length(T))), T)))
 end
 _unzip_size(::Type{T}, n::StaticInt{N}) where {T,N} = known_size(field_type(T, n))
-
-known_size(::Type{T}) where {T} = _maybe_known_size(Base.IteratorSize(T), T)
-function _maybe_known_size(::Base.HasShape{N}, ::Type{T}) where {N,T}
-    eachop(_known_size, nstatic(Val(N)), axes_types(T))
-end
-_maybe_known_size(::Base.IteratorSize, ::Type{T}) where {T} = (known_length(T),)
 _known_size(::Type{T}, dim::StaticInt) where {T} = known_length(field_type(T, dim))
 @inline known_size(x, dim) = known_size(typeof(x), dim)
 @inline known_size(::Type{T}, dim) where {T} = known_size(T, to_dims(T, dim))