Optimizations on to_indices

docs/src/api.md

@@ -17,12 +17,12 @@ ArrayInterface.fast_scalar_indexing
 ArrayInterface.has_dimnames
 ArrayInterface.has_parent
 ArrayInterface.has_sparsestruct
-ArrayInterface.is_canonical
 ArrayInterface.is_column_major
 ArrayInterface.is_lazy_conjugate
 ArrayInterface.ismutable
 ArrayInterface.issingular
 ArrayInterface.isstructured
+ArrayInterface.is_splat_index
 ArrayInterface.known_first
 ArrayInterface.known_last
 ArrayInterface.known_length
@@ -31,6 +31,8 @@ ArrayInterface.known_offsets
 ArrayInterface.known_size
 ArrayInterface.known_step
 ArrayInterface.known_strides
+ArrayInterface.ndims_index
+ArrayInterface.ndims_shape
 ```
 
 ## Functions
@@ -43,7 +45,6 @@ ArrayInterface.axes
 ArrayInterface.axes_types
 ArrayInterface.broadcast_axis
 ArrayInterface.buffer
-ArrayInterface.canonicalize
 ArrayInterface.deleteat
 ArrayInterface.dense_dims
 ArrayInterface.findstructralnz

docs/src/index.md

@@ -29,6 +29,7 @@ end
 ```
 
 Most traits in `ArrayInterface` are a variant on this pattern.
+If the trait in question may be altered by a wrapper array, this pattern should be altered or may be inappropriate.
 
 ## Static Traits
 
@@ -174,3 +175,67 @@ Defining these two methods ensures that other array types that wrap `OffsetArray
 It is entirely optional to define `ArrayInterface.size` for `OffsetArray` because the size can be derived from the axes.
 However, in this particularly case we should also define
  `ArrayInterface.size(A::OffsetArray)  = ArrayInterface.size(parent(A))` because the relative offsets attached to `OffsetArray` do not change the size but may hide static sizes if using a relative offset that is defined with an `Int`.
+
+
+## Indexing Protocol
+
+### Defining New Indexers
+
+* An index type that maps to a single dimension, such as `Bool`.
+
+We don't need to define `ndims_index` because the default value for any type is one.
+We need to convert `Bool` to an integer if `true` but throw an error if it is `false`.
+`ndims_shape` needs to be defined so as to indicate that indexing by `Bool` doesn't produce a collection of accessed values.
+```julia
+function ArrayInterface.to_index(::IndexStyle, x, i::Bool)
+    start = first(x)
+    if i
+        return start
+    else  # return out of bounds index
+        return start - 1
+    end
+end
+ArrayInterface.ndims_shape(::Type{Bool}) = static(0)
+```
+
+* An index that maps to a single point along multiple dimensions,
+
+If we want to map multiple boolean values to multiple dimensions with a single type we need to subtype `AbstractCartesianIndex`.
+`ndims_index` is already defined for the super type, so we just need to provide a method for converting it to a tuple.
+following.
+```julia
+struct CartesianBool{N} <: Base.AbstractCartesianIndex{N}
+    I::NTuple{N,Bool}
+end
+Base.Tuple(x::CartesianBool) = x.I
+
+```
+
+Note that we don't need to define a `to_index` method because the each value of the tuple will be a `Bool` and already has this method defined.
+
+* Splat first index of each axis to fill in all dimensions.
+
+```julia
+struct SplatFirst end
+
+ArrayInterface.to_index(x, ::SplatFirst) = first(x)
+
+ArrayInterface.is_splat_index(::Type{SplatFirst}) = static(true)
+
+x = rand(4,4,4,4,4,4,4,4,4,4);
+
+i = (2, SplatFirst(), 2, SplatFirst(), CartesianIndex(2, 2))
+
+ArrayInterface.to_indices(x, i) == (2, 1, 1, 1, 1, 1, 2, 1, 2, 2)
+```
+
+* An array whose elements map to multiple dimensions, such as `AbstractArray{Bool,2}`
+
+```julia
+ArrayInterface.ndims_index(::Type{<:AbstractArray{Bool,N}}) where {N} = static(N)
+function ArrayInterface.to_indices(A, i::Tuple{AbstractArray{Bool,N}}) where {N}
+    (Base.LogicalIndex(getfield(i, 1)),)
+end
+
+```
+This works by passing the first element of `I` in `to_indices(A, I::Tuple{AbstractArray{Bool,2},Vararg{Any}}})` to `to_indices(CartesianIndices((axes(A, 1), axes(A, 2))), (I[1],))`.

src/dimensions.jl

@@ -22,20 +22,40 @@ function is_increasing(perm::Tuple{StaticInt{X},StaticInt{Y}}) where {X, Y}
 end
 is_increasing(::Tuple{StaticInt{X}}) where {X} = True()
 
-#=
+"""
     ndims_index(::Type{I})::StaticInt
 
-The number of dimensions an instance of `I` maps to when indexing an instance of `A`.
-=#
+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(i) = ndims_index(typeof(i))
 ndims_index(::Type{I}) where {I} = static(1)
-ndims_index(::Type{I}) where {N,I<:AbstractCartesianIndex{N}} = static(N)
-ndims_index(::Type{I}) where {I<:AbstractArray} = ndims_index(eltype(I))
-ndims_index(::Type{I}) where {I<:AbstractArray{Bool}} = static(ndims(I))
-ndims_index(::Type{I}) where {N,I<:LogicalIndex{<:Any,<:AbstractArray{Bool,N}}} = static(N)
+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(_get_tuple(I, i))
 ndims_index(::Type{I}) where {N,I<:Tuple{Vararg{Any,N}}} = eachop(_ndims_index, nstatic(Val(N)), I)
 
+"""
+    ndims_shape(::Type{I})::StaticInt
+
+Returns the number of dimensions that `I` maps to in the returned array from an indexing
+operation. For example, `Int` would map to 0 dimensions but `Vector{Int}` would map to one
+dimension.
+"""
+ndims_shape(i) = ndims_shape(typeof(i))
+ndims_shape(::Type{I}) where {I} = static(1)
+ndims_shape(::Type{<:Integer}) = static(0)
+ndims_shape(@nospecialize(x::Type{<:StaticInt})) = static(0)
+ndims_shape(::Type{<:AbstractCartesianIndex}) = static(0)
+ndims_shape(::Type{<:AbstractArray{T,N}}) where {T,N} = static(N)
+ndims_shape(::Type{<:AbstractArray{Bool}}) = static(1)
+ndims_shape(::Type{<:CartesianIndices{N}}) where {N} = static(N)
+_ndims_shape(::Type{I}, i::StaticInt) where {I} = ndims_shape(_get_tuple(I, i))
+ndims_shape(::Type{I}) where {N,I<:Tuple{Vararg{Any,N}}} = eachop(_ndims_shape, nstatic(Val(N)), I)
+
 """
     from_parent_dims(::Type{T}) -> Tuple{Vararg{Union{Int,StaticInt}}}
     from_parent_dims(::Type{T}, dim) -> Union{Int,StaticInt}
@@ -191,7 +211,8 @@ end
 This returns the dimension(s) of `x` corresponding to `d`.
 """
 to_dims(x, dim) = to_dims(typeof(x), dim)
-to_dims(::Type{T}, dim::Integer) where {T} = canonicalize(dim)
+to_dims(::Type{T}, dim::StaticInt) where {T} = dim
+to_dims(::Type{T}, dim::Integer) where {T} = Int(dim)
 to_dims(::Type{T}, dim::Colon) where {T} = dim
 function to_dims(::Type{T}, dim::StaticSymbol) where {T}
     i = find_first_eq(dim, dimnames(T))