Move ndims_index and is_splat_index to core and put docs back in

Author: Tokazama

lib/ArrayInterfaceCore/src/ArrayInterfaceCore.jl

@@ -511,4 +511,25 @@ known_step(x) = known_step(typeof(x))
 known_step(::Type{T}) where {T} = parent_type(T) <: T ? nothing : known_step(parent_type(T))
 known_step(::Type{<:AbstractUnitRange}) = 1
 
+"""
+    is_splat_index(::Type{T}) -> Bool
+Returns `static(true)` if `T` is a type that splats across multiple dimensions. 
+"""
+is_splat_index(@nospecialize(x)) = is_splat_index(typeof(x))
+is_splat_index(T::Type) = false
+
+"""
+    ndims_index(::Type{I}) -> Int
+
+Returns the number of dimension that an instance of `I` maps to when indexing. For example,
+`CartesianIndex{3}` maps to 3 dimensions. If this method is not explicitly defined, then `1`
+is returned.
+"""
+ndims_index(@nospecialize(i)) = ndims_index(typeof(i))
+ndims_index(::Type{I}) where {I} = 1
+ndims_index(::Type{<:Base.AbstractCartesianIndex{N}}) where {N} = N
+ndims_index(::Type{<:AbstractArray{T}}) where {T} = ndims_index(T)
+ndims_index(::Type{<:AbstractArray{Bool,N}}) where {N} = N
+ndims_index(::Type{<:Base.LogicalIndex{<:Any,<:AbstractArray{Bool,N}}}) where {N} = N
+
 end # module

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, merge_tuple_type,
-    fast_scalar_indexing, parameterless_type, _is_reshaped
+    fast_scalar_indexing, parameterless_type, _is_reshaped, ndims_index, is_splat_index
 
 # ArrayIndex subtypes and methods
 import ArrayInterfaceCore: ArrayIndex, MatrixIndex, VectorIndex, BidiagonalIndex, TridiagonalIndex

src/dimensions.jl

@@ -36,6 +36,12 @@ function is_increasing(perm::Tuple{StaticInt{X},StaticInt{Y}}) where {X, Y}
 end
 is_increasing(::Tuple{StaticInt{X}}) where {X} = True()
 
+"""
+    from_parent_dims(::Type{T}) -> Tuple{Vararg{Union{Int,StaticInt}}}
+    from_parent_dims(::Type{T}, dim) -> Union{Int,StaticInt}
+
+Returns the mapping from parent dimensions to child dimensions.
+"""
 from_parent_dims(x) = from_parent_dims(typeof(x))
 from_parent_dims(::Type{T}) where {T} = nstatic(Val(ndims(T)))
 from_parent_dims(::Type{T}) where {T<:VecAdjTrans} = (StaticInt(2),)
@@ -76,7 +82,6 @@ Compat.@constprop :aggressive function from_parent_dims(::Type{T}, dim::Int)::In
         throw_dim_error(T, dim)
     end
 end
-
 function from_parent_dims(::Type{T}, ::StaticInt{dim}) where {T,dim}
     if dim > ndims(T)
         return static(ndims(parent_type(T)) + dim - ndims(T))
@@ -87,6 +92,12 @@ function from_parent_dims(::Type{T}, ::StaticInt{dim}) where {T,dim}
     end
 end
 
+"""
+    to_parent_dims(::Type{T}) -> Tuple{Vararg{Union{Int,StaticInt}}}
+    to_parent_dims(::Type{T}, dim) -> Union{Int,StaticInt}
+
+Returns the mapping from child dimensions to parent dimensions.
+"""
 to_parent_dims(x) = to_parent_dims(typeof(x))
 to_parent_dims(::Type{T}) where {T} = nstatic(Val(ndims(T)))
 to_parent_dims(::Type{T}) where {T<:Union{Transpose,Adjoint}} = (StaticInt(2), One())
@@ -135,10 +146,21 @@ function to_parent_dims(::Type{T}, ::StaticInt{dim}) where {T,dim}
     end
 end
 
-@inline function has_dimnames(x)
-    static(known_dimnames(x) !== ntuple(Compat.Returns(:_), Val(ndims(x))))
-end
+"""
+    has_dimnames(::Type{T}) -> Bool
 
+Returns `true` if `x` has on or more named dimensions. If all dimensions correspond
+to `:_`, then `false` is returned.
+"""
+@inline has_dimnames(x) = static(known_dimnames(x) !== ntuple(Compat.Returns(:_), Val(ndims(x))))
+
+"""
+    known_dimnames(::Type{T}) -> Tuple{Vararg{Union{Symbol,Nothing}}}
+    known_dimnames(::Type{T}, dim::Union{Int,StaticInt}) -> Union{Symbol,Nothing}
+
+Return the names of the dimensions for `x`. `:_` is used to indicate a dimension does not
+have a name.
+"""
 @inline known_dimnames(x, dim::Integer) = _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))
@@ -155,6 +177,13 @@ end
 end
 @inline _inbounds_known_dimname(x, dim) = @inbounds(_known_dimname(x, dim))
 
+"""
+    dimnames(x) -> Tuple{Vararg{Union{Symbol,StaticSymbol}}}
+    dimnames(x, dim::Union{Int,StaticInt}) -> Union{Symbol,StaticSymbol}
+
+Return the names of the dimensions for `x`. `:_` is used to indicate a dimension does not
+have a name.
+"""
 @inline dimnames(x, dim::Integer) = _dimname(dimnames(x), canonicalize(dim))
 @inline dimnames(x) = _dimnames(has_parent(x), x)
 @inline function _dimnames(::True, x)
@@ -169,6 +198,11 @@ _dimnames(::False, x) = ntuple(_->static(:_), Val(ndims(x)))
 end
 @inline _inbounds_dimname(x, dim) = @inbounds(_dimname(x, dim))
 
+"""
+    to_dims(x, dim) -> Union{Int,StaticInt}
+
+This returns the dimension(s) of `x` corresponding to `dim`.
+"""
 to_dims(x, dim::Colon) = dim
 to_dims(x, dim::Integer) = canonicalize(dim)
 to_dims(x, dim::Union{StaticSymbol,Symbol}) = _to_dim(dimnames(x), dim)

src/indexing.jl

@@ -1,22 +1,13 @@
-is_splat_index(@nospecialize(x)) = is_splat_index(typeof(x))
-is_splat_index(::Type{T}) where {T} = static(false)
-_is_splat(::Type{I}, i::StaticInt) where {I} = is_splat_index(field_type(I, i))
 
-"""
-    ndims_index(::Type{I}) -> StaticInt
-
-Returns the number of dimension that an instance of `I` maps to when indexing. For example,
-`CartesianIndex{3}` maps to 3 dimensions. If this method is not explicitly defined, then `1`
-is returned.
+function _is_splat(::Type{I}, i::StaticInt) where {I}
+    if dynamic(is_splat_index(field_type(I, i)))
+        True()
+    else
+        False()
+    end
+end
 
-"""
-ndims_index(@nospecialize(i)) = ndims_index(typeof(i))
-ndims_index(::Type{I}) where {I} = static(1)
-ndims_index(::Type{<:AbstractCartesianIndex{N}}) where {N} = static(N)
-ndims_index(::Type{<:AbstractArray{T}}) where {T} = ndims_index(T)
-ndims_index(::Type{<:AbstractArray{Bool,N}}) where {N} = static(N)
-ndims_index(::Type{<:LogicalIndex{<:Any,<:AbstractArray{Bool,N}}}) where {N} = static(N)
-_ndims_index(::Type{I}, i::StaticInt) where {I} = ndims_index(field_type(I, i))
+_ndims_index(::Type{I}, i::StaticInt) where {I} = StaticInt(ndims_index(field_type(I, i)))
 
 """
     to_indices(A, I::Tuple) -> Tuple
@@ -238,7 +229,7 @@ indices calling [`to_axis`](@ref).
 end
 # drop this dimension
 to_axes(A, a::Tuple, i::Tuple{<:Integer,Vararg{Any}}) = to_axes(A, tail(a), tail(i))
-to_axes(A, a::Tuple, i::Tuple{I,Vararg{Any}}) where {I} = _to_axes(ndims_index(I), A, a, i)
+to_axes(A, a::Tuple, i::Tuple{I,Vararg{Any}}) where {I} = _to_axes(StaticInt(ndims_index(I)), A, a, i)
 function _to_axes(::StaticInt{1}, A, axs::Tuple, inds::Tuple)
     return (to_axis(first(axs), first(inds)), to_axes(A, tail(axs), tail(inds))...)
 end