Cleanup indices info + better docs (#342)

Cleanup and document IndicesInfo 

Before making this public I wanted to ensure that we had control over
the construction of `IndicesInfo` so it was guaranteed to be valid. Now
it has inner struct constructors only. Also provided `ndims_index` and
`ndims_shape` so all information can be accessed without directly
looking at `IndicesInfo`'s parametric typing.

Author: Tokazama

docs/src/api.md

@@ -12,6 +12,8 @@ ArrayInterfaceCore.device
 ArrayInterfaceCore.defines_strides
 ArrayInterfaceCore.fast_matrix_colors
 ArrayInterfaceCore.fast_scalar_indexing
+ArrayInterfaceCore.indices_do_not_alias
+ArrayInterfaceCore.instances_do_not_alias
 ArrayInterfaceCore.is_forwarding_wrapper
 ArrayInterfaceCore.ismutable
 ArrayInterfaceCore.isstructured
@@ -84,6 +86,7 @@ ArrayInterface.from_parent_dims
 ArrayInterface.getindex
 ArrayInterface.indices
 ArrayInterface.insert
+ArrayInterface.length
 ArrayInterface.lazy_axes
 ArrayInterface.offset1
 ArrayInterface.offsets
@@ -106,8 +109,8 @@ ArrayInterface.BroadcastAxis
 ArrayInterface.LazyAxis
 ArrayInterface.OptionallyStaticStepRange
 ArrayInterface.OptionallyStaticUnitRange
-ArrayInteraface.SOneTo
-ArrayInteraface.SUnitRange
+ArrayInterface.SOneTo
+ArrayInterface.SUnitRange
 ArrayInterface.StrideIndex
 ```
 

lib/ArrayInterfaceCore/src/ArrayInterfaceCore.jl

@@ -806,74 +806,122 @@ ndims_shape(x) = ndims_shape(typeof(x))
 end
 
 """
+    IndicesInfo{N}(inds::Tuple) -> IndicesInfo{N}(typeof(inds))
     IndicesInfo{N}(T::Type{<:Tuple}) -> IndicesInfo{N,pdims,cdims}()
+    IndicesInfo(inds::Tuple) -> IndicesInfo(typeof(inds))
+    IndicesInfo(T::Type{<:Tuple}) -> IndicesInfo{maximum(pdims),pdims,cdims}()
 
-Provides basic trait information for each index type in in the tuple `T`. `pdims` and
-`cdims` are dimension mappings to the parent and child dimensions respectively.
+
+Maps a tuple of indices to `N` dimensions. The resulting `pdims` is a tuple where each
+field in `inds` (or field type in `T`) corresponds to the parent dimensions accessed.
+`cdims` similarly maps indices to the resulting child array produced after indexing with
+`inds`. If `N` is not provided then it is assumed that all indices are represented by parent
+dimensions and there are no trailing dimensions accessed. These may be accessed by through
+`parentdims(info::IndicesInfo)` and `childdims(info::IndicesInfo)`. If `N` is not provided,
+it is assumed that no indices are accessing trailing dimensions (which are represented as
+`0` in `parentdims(info)[index_position]`).
+
+The the fields and types of `IndicesInfo` should not be accessed directly.
+Instead [`parentdims`](@ref), [`childdims`](@ref), [`ndims_index`](@ref), and
+[`ndims_shape`](@ref) should be used to extract relevant information.
 
 # Examples
 
 ```julia
-julia> using ArrayInterfaceCore: IndicesInfo
+julia> using ArrayInterfaceCore: IndicesInfo, parentdims, childdims, ndims_index, ndims_shape
+
+julia> info = IndicesInfo{5}(typeof((:,[CartesianIndex(1,1),CartesianIndex(1,1)], 1, ones(Int, 2, 2), :, 1)));
+
+julia> parentdims(info)  # the last two indices access trailing dimensions
+(1, (2, 3), 4, 5, 0, 0)
+
+julia> childdims(info)
+(1, 2, 0, (3, 4), 5, 0)
+
+julia> childdims(info)[3]  # index 3 accesses a parent dimension but is dropped in the child array
+0
 
-julia> IndicesInfo{5}(typeof((:,[CartesianIndex(1,1),CartesianIndex(1,1)], 1, ones(Int, 2, 2), :, 1)))
-IndicesInfo{5, (1, (2, 3), 4, 5, 0, 0), (1, 2, 0, (3, 4), 5, 0)}()
+julia> ndims_index(info)
+5
+
+julia> ndims_shape(info)
+5
+
+julia> info = IndicesInfo(typeof((:,[CartesianIndex(1,1),CartesianIndex(1,1)], 1, ones(Int, 2, 2), :, 1)));
+
+julia> parentdims(info)  # assumed no trailing dimensions
+(1, (2, 3), 4, 5, 6, 7)
+
+julia> ndims_index(info)  # assumed no trailing dimensions
+7
 
 ```
 """
-struct IndicesInfo{N,NI,NS} end
-IndicesInfo(x::SubArray) = IndicesInfo{ndims(parent(x))}(typeof(x.indices))
-@inline function IndicesInfo(@nospecialize T::Type{<:SubArray})
-    IndicesInfo{ndims(parent_type(T))}(fieldtype(T, :indices))
-end
-function IndicesInfo{N}(@nospecialize(T::Type{<:Tuple})) where {N}
-    _indices_info(
-        Val{_find_first_true(map_tuple_type(is_splat_index, T))}(),
-        IndicesInfo{N,map_tuple_type(ndims_index, T),map_tuple_type(ndims_shape, T)}()
-    )
-end
-function _indices_info(::Val{nothing}, ::IndicesInfo{1,(1,),NS}) where {NS}
-    ns1 = getfield(NS, 1)
-    IndicesInfo{1,(1,), (ns1 > 1 ? ntuple(identity, ns1) : ns1,)}()
-end
-function _indices_info(::Val{nothing}, ::IndicesInfo{N,(1,),NS}) where {N,NS}
-    ns1 = getfield(NS, 1)
-    IndicesInfo{N,(:,),(ns1 > 1 ? ntuple(identity, ns1) : ns1,)}()
-end
-@inline function _indices_info(::Val{nothing}, ::IndicesInfo{N,NI,NS}) where {N,NI,NS}
-    if sum(NI) > N
-        IndicesInfo{N,_replace_trailing(N, _accum_dims(cumsum(NI), NI)), _accum_dims(cumsum(NS), NS)}()
-    else
-        IndicesInfo{N,_accum_dims(cumsum(NI), NI), _accum_dims(cumsum(NS), NS)}()
+struct IndicesInfo{Np,pdims,cdims,Nc}
+    function IndicesInfo{N}(@nospecialize(T::Type{<:Tuple})) where {N}
+        SI = _find_first_true(map_tuple_type(is_splat_index, T))
+        NI = map_tuple_type(ndims_index, T)
+        NS = map_tuple_type(ndims_shape, T)
+        if SI === nothing
+            ndi = NI
+            nds = NS
+        else
+            nsplat = N - sum(NI)
+            if nsplat === 0
+                ndi = NI
+                nds = NS
+            else
+                splatmul = max(0, nsplat + 1)
+                ndi = _map_splats(splatmul, SI, NI)
+                nds = _map_splats(splatmul, SI, NS)
+            end
+        end
+        if ndi === (1,) && N !== 1
+            ns1 = getfield(nds, 1)
+            new{N,(:,),(ns1 > 1 ? ntuple(identity, ns1) : ns1,),ns1}()
+        else
+            nds_cumsum = cumsum(nds)
+            if sum(ndi) > N
+                init_pdims = _accum_dims(cumsum(ndi), ndi)
+                pdims = ntuple(nfields(init_pdims)) do i
+                    dim_i = getfield(init_pdims, i)
+                    if dim_i isa Tuple
+                        ntuple(length(dim_i)) do j
+                            dim_i_j = getfield(dim_i, j)
+                            dim_i_j > N ? 0 : dim_i_j
+                        end
+                    else
+                        dim_i > N ? 0 : dim_i
+                    end
+                end
+                new{N, pdims, _accum_dims(nds_cumsum, nds), last(nds_cumsum)}()
+            else
+                new{N,_accum_dims(cumsum(ndi), ndi), _accum_dims(nds_cumsum, nds), last(nds_cumsum)}()
+            end
+        end
     end
-end
-@inline function _indices_info(::Val{SI}, ::IndicesInfo{N,NI,NS}) where {N,NI,NS,SI}
-    nsplat = N - sum(NI)
-    if nsplat === 0
-        _indices_info(Val{nothing}(), IndicesInfo{N,NI,NS}())
-    else
-        splatmul = max(0, nsplat + 1)
-        _indices_info(Val{nothing}(), IndicesInfo{N,_map_splats(splatmul, SI, NI),_map_splats(splatmul, SI, NS)}())
+    IndicesInfo{N}(@nospecialize(t::Tuple)) where {N} = IndicesInfo{N}(typeof(t))
+    function IndicesInfo(@nospecialize(T::Type{<:Tuple}))
+        ndi = map_tuple_type(ndims_index, T)
+        nds = map_tuple_type(ndims_shape, T)
+        ndi_sum = cumsum(ndi)
+        nds_sum = cumsum(nds)
+        nf = nfields(ndi_sum)
+        pdims = _accum_dims(ndi_sum, ndi)
+        cdims = _accum_dims(nds_sum, nds)
+        new{getfield(ndi_sum, nf),pdims,cdims,getfield(nds_sum, nf)}()
     end
+    IndicesInfo(@nospecialize t::Tuple) = IndicesInfo(typeof(t))
+    @inline function IndicesInfo(@nospecialize T::Type{<:SubArray})
+        IndicesInfo{ndims(parent_type(T))}(fieldtype(T, :indices))
+    end
+    IndicesInfo(x::SubArray) = IndicesInfo{ndims(parent(x))}(typeof(x.indices))
 end
 @inline function _map_splats(nsplat::Int, splat_index::Int, dims::Tuple{Vararg{Int}})
     ntuple(length(dims)) do i
         i === splat_index ? (nsplat * getfield(dims, i)) : getfield(dims, i)
     end
 end
-@inline function _replace_trailing(n::Int, dims::Tuple{Vararg{Any,N}}) where {N}
-    ntuple(N) do i
-        dim_i = getfield(dims, i)
-        if dim_i isa Tuple
-            ntuple(length(dim_i)) do j
-                dim_i_j = getfield(dim_i, j)
-                dim_i_j > n ? 0 : dim_i_j
-            end
-        else
-            dim_i > n ? 0 : dim_i
-        end
-    end
-end
 @inline function _accum_dims(csdims::NTuple{N,Int}, nd::NTuple{N,Int}) where {N}
     ntuple(N) do i
         nd_i = getfield(nd, i)
@@ -887,6 +935,30 @@ end
     end
 end
 
+_lower_info(::IndicesInfo{Np,pdims,cdims,Nc}) where {Np,pdims,cdims,Nc} = Np,pdims,cdims,Nc
+
+ndims_index(@nospecialize(info::IndicesInfo)) = getfield(_lower_info(info), 1)
+ndims_shape(@nospecialize(info::IndicesInfo)) = getfield(_lower_info(info), 4)
+
+"""
+    parentdims(::IndicesInfo) -> Tuple
+
+Returns the parent dimension mapping from `IndicesInfo`.
+
+See also: [`IndicesInfo`](@ref), [`childdims`](@ref)
+"""
+parentdims(@nospecialize info::IndicesInfo) = getfield(_lower_info(info), 2)
+
+"""
+    childdims(::IndicesInfo) -> Tuple
+
+Returns the child dimension mapping from `IndicesInfo`.
+
+See also: [`IndicesInfo`](@ref), [`parentdims`](@ref)
+"""
+childdims(@nospecialize info::IndicesInfo) = getfield(_lower_info(info), 3)
+
+
 """
     instances_do_not_alias(::Type{T}) -> Bool
 

lib/ArrayInterfaceCore/test/runtests.jl

@@ -300,29 +300,29 @@ end
 
     ArrayInterfaceCore.is_splat_index(::Type{SplatFirst}) = true
 
-    @test @inferred(IndicesInfo(SubArray{Float64, 2, Vector{Float64}, Tuple{Base.ReshapedArray{Int64, 2, UnitRange{Int64}, Tuple{}}}, true})) ==
-        IndicesInfo{1,(1,),((1,2),)}()
+    @test @inferred(IndicesInfo(SubArray{Float64, 2, Vector{Float64}, Tuple{Base.ReshapedArray{Int64, 2, UnitRange{Int64}, Tuple{}}}, true})) isa
+        IndicesInfo{1,(1,),((1,2),)}
 
-    @test @inferred(IndicesInfo{1}((Tuple{Vector{Int}}))) == IndicesInfo{1, (1,), (1,)}()
+    @test @inferred(IndicesInfo{1}((Tuple{Vector{Int}}))) isa IndicesInfo{1, (1,), (1,)}
 
-    @test @inferred(IndicesInfo{2}(Tuple{Vector{Int}})) == IndicesInfo{2, (:,), (1,)}()
+    @test @inferred(IndicesInfo{2}(Tuple{Vector{Int}})) isa IndicesInfo{2, (:,), (1,)}
 
-    @test @inferred(IndicesInfo{1}(Tuple{SplatFirst})) == IndicesInfo{1, (1,), (1,)}()
+    @test @inferred(IndicesInfo{1}(Tuple{SplatFirst})) isa IndicesInfo{1, (1,), (1,)}
 
-    @test @inferred(IndicesInfo{2}(Tuple{SplatFirst})) == IndicesInfo{2, ((1,2),), ((1, 2),)}()
+    @test @inferred(IndicesInfo{2}(Tuple{SplatFirst})) isa IndicesInfo{2, ((1,2),), ((1, 2),)}
 
-    @test @inferred(IndicesInfo{5}(typeof((:,[CartesianIndex(1,1),CartesianIndex(1,1)], 1, ones(Int, 2, 2), :, 1)))) ==
-        IndicesInfo{5, (1, (2, 3), 4, 5, 0, 0), (1, 2, 0, (3, 4), 5, 0)}()
+    @test @inferred(IndicesInfo{5}(typeof((:,[CartesianIndex(1,1),CartesianIndex(1,1)], 1, ones(Int, 2, 2), :, 1)))) isa
+        IndicesInfo{5, (1, (2, 3), 4, 5, 0, 0), (1, 2, 0, (3, 4), 5, 0)}
 
-    @test @inferred(IndicesInfo{10}(Tuple{Vararg{Int,10}})) ==
-        IndicesInfo{10, (1, 2, 3, 4, 5, 6, 7, 8, 9, 10), (0, 0, 0, 0, 0, 0, 0, 0, 0, 0)}()
+    @test @inferred(IndicesInfo{10}(Tuple{Vararg{Int,10}})) isa
+        IndicesInfo{10, (1, 2, 3, 4, 5, 6, 7, 8, 9, 10), (0, 0, 0, 0, 0, 0, 0, 0, 0, 0)}
 
-    @test @inferred(IndicesInfo{10}(typeof((1, CartesianIndex(2, 1), 2, CartesianIndex(1, 2), 1, CartesianIndex(2, 1), 2)))) ==
-        IndicesInfo{10, (1, (2, 3), 4, (5, 6), 7, (8, 9), 10), (0, 0, 0, 0, 0, 0, 0)}()
+    @test @inferred(IndicesInfo{10}(typeof((1, CartesianIndex(2, 1), 2, CartesianIndex(1, 2), 1, CartesianIndex(2, 1), 2)))) isa
+        IndicesInfo{10, (1, (2, 3), 4, (5, 6), 7, (8, 9), 10), (0, 0, 0, 0, 0, 0, 0)}
 
-    @test @inferred(IndicesInfo{10}(typeof((fill(true, 4, 4), 2, fill(true, 4, 4), 2, 1, fill(true, 4, 4), 1)))) ==
-        IndicesInfo{10, ((1, 2), 3, (4, 5), 6, 7, (8, 9), 10), (1, 0, 2, 0, 0, 3, 0)}()
+    @test @inferred(IndicesInfo{10}(typeof((fill(true, 4, 4), 2, fill(true, 4, 4), 2, 1, fill(true, 4, 4), 1)))) isa
+        IndicesInfo{10, ((1, 2), 3, (4, 5), 6, 7, (8, 9), 10), (1, 0, 2, 0, 0, 3, 0)}
 
-    @test @inferred(IndicesInfo{10}(typeof((1, SplatFirst(), 2, SplatFirst(), CartesianIndex(1, 1))))) ==
-        IndicesInfo{10, (1, (2, 3, 4, 5, 6), 7, 8, (9, 10)), (0, (1, 2, 3, 4, 5), 0, 6, 0)}()
+    @test @inferred(IndicesInfo{10}(typeof((1, SplatFirst(), 2, SplatFirst(), CartesianIndex(1, 1))))) isa
+        IndicesInfo{10, (1, (2, 3, 4, 5, 6), 7, 8, (9, 10)), (0, (1, 2, 3, 4, 5), 0, 6, 0)}
 end

src/ArrayInterface.jl

@@ -5,8 +5,8 @@ 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, ndims_shape, is_splat_index, is_forwarding_wrapper, IndicesInfo,
-    map_tuple_type, flatten_tuples, GetIndex, SetIndex!, defines_strides,
+    ndims_index, ndims_shape, is_splat_index, is_forwarding_wrapper, IndicesInfo, childdims,
+    parentdims, map_tuple_type, flatten_tuples, GetIndex, SetIndex!, defines_strides,
     stride_preserving_index
 
 # ArrayIndex subtypes and methods

src/dimensions.jl

@@ -1,7 +1,9 @@
 
 
 _init_dimsmap(x) = _init_dimsmap(IndicesInfo(x))
-function _init_dimsmap(::IndicesInfo{N,pdims,cdims}) where {N,pdims,cdims}
+function _init_dimsmap(@nospecialize info::IndicesInfo)
+    pdims = parentdims(info)
+    cdims = childdims(info)
     ntuple(i -> static(getfield(pdims, i)), length(pdims)),
     ntuple(i -> static(getfield(cdims, i)), length(pdims))
 end
@@ -48,7 +50,9 @@ function _sub_axis_map(@nospecialize(T::Type{<:SubArray}), x::Tuple{StaticInt{in
     end
 end
 
-function map_indices_info(::IndicesInfo{N,pdims,cdims}) where {N,pdims,cdims}
+function map_indices_info(@nospecialize info::IndicesInfo)
+    pdims = parentdims(info)
+    cdims = childdims(info)
     ntuple(i -> (static(i), static(getfield(pdims, i)), static(getfield(cdims, i))), length(pdims))
 end
 function sub_dimnames_map(dnames::Tuple, imap::Tuple)
@@ -82,7 +86,9 @@ from_parent_dims(@nospecialize T::Type{<:MatAdjTrans}) = (StaticInt(2), StaticIn
     from_parent_dims(IndicesInfo{ndims(parent_type(T))}(fieldtype(T, :indices)))
 end
 # TODO do I need to flatten_tuples here?
-function from_parent_dims(::IndicesInfo{N,pdims,cdims}) where {N,pdims,cdims}
+function from_parent_dims(@nospecialize(info::IndicesInfo))
+    pdims = parentdims(info)
+    cdims = childdims(info)
     ntuple(length(cdims)) do i
         pdim_i = getfield(pdims, i)
         cdim_i = static(getfield(cdims, i))