Integrte dimension into IndicesInfo

Author: Tokazama

lib/ArrayInterfaceCore/src/ArrayInterfaceCore.jl

@@ -671,51 +671,55 @@ ndims_shape(x) = ndims_shape(typeof(x))
 end
 
 """
-    IndicesInfo(T::Type{<:Tuple}) -> IndicesInfo{NI,NS,IS}()
+    IndicesInfo{N}(T::Type{<:Tuple}) -> IndicesInfo{N,NI,NS}()
 
 Provides basic trait information for each index type in in the tuple `T`. `NI`, `NS`, and
 `IS` are tuples of [`ndims_index`](@ref), [`ndims_shape`](@ref), and
 [`is_splat_index`](@ref) (respectively) for each field of `T`.
+
+# Examples
+
+```julia
+julia> using ArrayInterfaceCore: IndicesInfo
+
+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)}()
+
+```
 """
-struct IndicesInfo{NI,NS,IS} end
-IndicesInfo(x::Union{Tuple,SubArray}) = IndicesInfo(typeof(x))
-function IndicesInfo(@nospecialize T::Type{<:Tuple})
-    IndicesInfo{
-        map_tuple_type(ndims_index, T),
-        map_tuple_type(ndims_shape, T),
-        _find_first_true(map_tuple_type(is_splat_index, T))
-    }()
+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(@nospecialize T::Type{<:SubArray})
-    indices_to_dims(IndicesInfo(fieldtype(T, :indices)), parent_type(T))
+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
-
-@inline function indices_to_dims(@nospecialize(I::Type{<:Tuple}), @nospecialize(T::Type))
-    indices_to_dims(IndicesInfo(I), T)
+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
-@inline function indices_to_dims(::IndicesInfo{(1,),NS,nothing}, @nospecialize(T::Type)) where {NS}
+function _indices_info(::Val{nothing}, ::IndicesInfo{N,(1,),NS}) where {N,NS}
     ns1 = getfield(NS, 1)
-    nsd = ns1 > 1 ? ntuple(identity, ns1) : ns1
-    if ndims(T) === 1
-        IndicesInfo{(1,), (nsd,), nothing}()
-    else
-        IndicesInfo{(:,), (nsd,), nothing}()
-    end
+    IndicesInfo{N,(:,),(ns1 > 1 ? ntuple(identity, ns1) : ns1,)}()
 end
-@inline function indices_to_dims(::IndicesInfo{NI,NS,nothing}, @nospecialize(T::Type)) where {NI,NS}
-    if sum(NI) > ndims(T)
-        IndicesInfo{_replace_trailing(ndims(T), _accum_dims(cumsum(NI), NI)), _accum_dims(cumsum(NS), NS), nothing}()
+@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{_accum_dims(cumsum(NI), NI), _accum_dims(cumsum(NS), NS), nothing}()
+        IndicesInfo{N,_accum_dims(cumsum(NI), NI), _accum_dims(cumsum(NS), NS)}()
     end
 end
-@inline function indices_to_dims(::IndicesInfo{NI,NS,SI}, @nospecialize(T::Type)) where {NI,NS,SI}
-    nsplat = ndims(T) - sum(NI)
+@inline function _indices_info(::Val{SI}, ::IndicesInfo{N,NI,NS}) where {N,NI,NS,SI}
+    nsplat = N - sum(NI)
     if nsplat === 0
-        indices_to_dims(IndicesInfo{NI,NS,nothing}(), T)
+        _indices_info(Val{nothing}(), IndicesInfo{N,NI,NS}())
     else
         splatmul = max(0, nsplat + 1)
-        indices_to_dims(IndicesInfo{_map_splats(splatmul, SI, NI),_map_splats(splatmul, SI, NS),nothing}(), T)
+        _indices_info(Val{nothing}(), IndicesInfo{N,_map_splats(splatmul, SI, NI),_map_splats(splatmul, SI, NS)}())
     end
 end
 @inline function _map_splats(nsplat::Int, splat_index::Int, dims::Tuple{Vararg{Int}})

lib/ArrayInterfaceCore/test/runtests.jl

@@ -294,35 +294,35 @@ end
   @test !ArrayInterfaceCore.indices_do_not_alias(typeof(view(rand(4,4)', StepRangeLen(1,0,5), 1:2)))
 end
 
-@testset "indices_to_dims" begin
+@testset "IndicesInfo" begin
 
     struct SplatFirst 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,2),),nothing}()
+        IndicesInfo{1,(1,),((1,2),)}()
 
-    @test @inferred(indices_to_dims(Tuple{Vector{Int}}, Vector{Int})) == IndicesInfo{(1,), (1,), nothing}()
+    @test @inferred(IndicesInfo{1}((Tuple{Vector{Int}}))) == IndicesInfo{1, (1,), (1,)}()
 
-    @test @inferred(indices_to_dims(Tuple{Vector{Int}}, Matrix{Int})) == IndicesInfo{(:,), (1,), nothing}()
+    @test @inferred(IndicesInfo{2}(Tuple{Vector{Int}})) == IndicesInfo{2, (:,), (1,)}()
 
-    @test @inferred(indices_to_dims(Tuple{SplatFirst}, Vector{Int})) == IndicesInfo{(1,), (1,), nothing}()
+    @test @inferred(IndicesInfo{1}(Tuple{SplatFirst})) == IndicesInfo{1, (1,), (1,)}()
 
-    @test @inferred(indices_to_dims(Tuple{SplatFirst}, Matrix{Int})) == IndicesInfo{((1,2),), ((1, 2),), nothing}()
+    @test @inferred(IndicesInfo{2}(Tuple{SplatFirst})) == IndicesInfo{2, ((1,2),), ((1, 2),)}()
 
-    @test @inferred(indices_to_dims(typeof((:,[CartesianIndex(1,1),CartesianIndex(1,1)], 1, ones(Int, 2, 2), :, 1)), Array{Int,5})) ==
-        IndicesInfo{(1, (2, 3), 4, 5, 0, 0), (1, 2, 0, (3, 4), 5, 0), nothing}()
+    @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(indices_to_dims(Tuple{Vararg{Int,10}}, Array{Int,10})) ==
-        IndicesInfo{(1, 2, 3, 4, 5, 6, 7, 8, 9, 10), (0, 0, 0, 0, 0, 0, 0, 0, 0, 0), nothing}()
+    @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(indices_to_dims(typeof((1, CartesianIndex(2, 1), 2, CartesianIndex(1, 2), 1, CartesianIndex(2, 1), 2)),Array{Int,10})) ==
-        IndicesInfo{(1, (2, 3), 4, (5, 6), 7, (8, 9), 10), (0, 0, 0, 0, 0, 0, 0), nothing}()
+    @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(indices_to_dims(typeof((fill(true, 4, 4), 2, fill(true, 4, 4), 2, 1, fill(true, 4, 4), 1)), Array{Int,10})) ==
-        IndicesInfo{((1, 2), 3, (4, 5), 6, 7, (8, 9), 10), (1, 0, 2, 0, 0, 3, 0), nothing}()
+    @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(indices_to_dims(typeof((1, SplatFirst(), 2, SplatFirst(), CartesianIndex(1, 1))), Array{Int,10})) ==
-        IndicesInfo{(1, (2, 3, 4, 5, 6), 7, 8, (9, 10)), (0, (1, 2, 3, 4, 5), 0, 6, 0), nothing}()
+    @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)}()
 end

src/ArrayInterface.jl

@@ -6,7 +6,7 @@ import ArrayInterfaceCore: allowed_getindex, allowed_setindex!, aos_to_soa, buff
     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, indices_to_dims
+    map_tuple_type, flatten_tuples, GetIndex
 
 # ArrayIndex subtypes and methods
 import ArrayInterfaceCore: ArrayIndex, MatrixIndex, VectorIndex, BidiagonalIndex, TridiagonalIndex

src/dimensions.jl

@@ -1,8 +1,7 @@
 
 
-_init_dimsmap(::Type{I}, ::Type{P}) where {I,P} = _init_dimsmap(indices_to_dims(I, P))
 _init_dimsmap(x) = _init_dimsmap(IndicesInfo(x))
-function _init_dimsmap(::IndicesInfo{pdims,cdims,nothing}) where {pdims,cdims}
+function _init_dimsmap(::IndicesInfo{N,pdims,cdims}) where {N,pdims,cdims}
     ntuple(i -> static(getfield(pdims, i)), length(pdims)),
     ntuple(i -> static(getfield(cdims, i)), length(pdims))
 end
@@ -21,9 +20,9 @@ Returns the mapping from child dimensions to parent dimensions.
 """
 to_parent_dims(@nospecialize x) = to_parent_dims(typeof(x))
 @inline function to_parent_dims(@nospecialize T::Type{<:SubArray})
-    to_parent_dims(indices_to_dims(fieldtype(T, :indices), parent_type(T)))
+    to_parent_dims(IndicesInfo{ndims(parent_type(T))}(fieldtype(T, :indices)))
 end
-function to_parent_dims(::IndicesInfo{pdims,cdims,nothing}) where {pdims,cdims}
+function to_parent_dims(::IndicesInfo{N,pdims,cdims}) where {N,pdims,cdims}
     flatten_tuples(ntuple(length(cdims)) do i
         cdim_i = getfield(cdims, i)
         if cdim_i isa Tuple
@@ -48,7 +47,7 @@ end
 # Base will sometomes demote statically known slices in `SubArray` to `OneTo{Int}` so we
 # provide the parent mapping to check for static size info
 sub_axes_map(@nospecialize(T::Type{<:SubArray})) = _sub_axes_map(T, IndicesInfo(T))
-function _sub_axes_map(@nospecialize(T::Type{<:SubArray}), ::IndicesInfo{pdims}) where {pdims}
+function _sub_axes_map(@nospecialize(T::Type{<:SubArray}), ::IndicesInfo{N,pdims}) where {N,pdims}
     ntuple(length(pdims)) do i
         if fieldtype(fieldtype(T, :indices), i) <: Base.Slice{OneTo{Int}}
             sz = known_size(parent_type(T), getfield(pdims, i))
@@ -60,7 +59,7 @@ function _sub_axes_map(@nospecialize(T::Type{<:SubArray}), ::IndicesInfo{pdims})
 end
 
 sub_dimnames_map(@nospecialize T::Type{<:SubArray}) = _sub_dimnames_map(IndicesInfo(T))
-function _sub_dimnames_map(::IndicesInfo{pdims,cdims}) where {pdims,cdims}
+function _sub_dimnames_map(::IndicesInfo{N,pdims,cdims}) where {N,pdims,cdims}
     ntuple(length(pdims)) do i
         cdim_i = getfield(cdims, i)
         pdim_i = getfield(pdims, i)
@@ -87,10 +86,10 @@ from_parent_dims(@nospecialize x) = from_parent_dims(typeof(x))
 from_parent_dims(@nospecialize T::Type{<:PermutedDimsArray}) = getfield(_permdims(T), 2)
 from_parent_dims(@nospecialize T::Type{<:MatAdjTrans}) = (StaticInt(2), StaticInt(1))
 @inline function from_parent_dims(@nospecialize T::Type{<:SubArray})
-    from_parent_dims(indices_to_dims(fieldtype(T, :indices), parent_type(T)))
+    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{pdims,cdims,nothing}) where {pdims,cdims}
+function from_parent_dims(::IndicesInfo{N,pdims,cdims}) where {N,pdims,cdims}
     ntuple(length(cdims)) do i
         pdim_i = getfield(pdims, i)
         cdim_i = static(getfield(cdims, i))
@@ -130,7 +129,6 @@ function _known_sub_dimname(::Tuple, ::Pair{StaticSymbol{:underscore},StaticInt{
     ntuple(Compat.Returns(:_), StaticInt(N))
 end
 
-
 function known_dimnames(::Type{<:ReinterpretArray{T,N,S,A,IsReshaped}}) where {T,N,S,A,IsReshaped}
     pnames = known_dimnames(A)
     if IsReshaped

src/indexing.jl

@@ -81,7 +81,7 @@ This implementation differs from that of `Base.to_indices` in the following ways
 """
 to_indices(A, ::Tuple{}) = ()
 @inline function to_indices(a::A, inds::I) where {A,I}
-    flatten_tuples(map(IndexedMappedArray(a), inds, getfield(_init_dimsmap(I, A), 1)))
+    flatten_tuples(map(IndexedMappedArray(a), inds, getfield(_init_dimsmap(IndicesInfo{ndims(A)}(I)), 1)))
 end
 
 struct IndexedMappedArray{A}

src/stridelayout.jl

@@ -125,11 +125,11 @@ contiguous_axis(::Type{<:StrideIndex{N,R,nothing}}) where {N,R} = nothing
 function contiguous_axis(::Type{T}) where {T}
     is_forwarding_wrapper(T) ? contiguous_axis(parent_type(T)) : nothing
 end
-contiguous_axis(::Type{<:DenseArray}) = One()
+contiguous_axis(@nospecialize T::Type{<:DenseArray}) = One()
 contiguous_axis(::Type{<:BitArray}) = One()
-contiguous_axis(::Type{<:AbstractRange}) = One()
-contiguous_axis(::Type{<:Tuple}) = One()
-function contiguous_axis(::Type{T}) where {T<:VecAdjTrans}
+contiguous_axis(@nospecialize T::Type{<:AbstractRange}) = One()
+contiguous_axis(@nospecialize T::Type{<:Tuple}) = One()
+function contiguous_axis(@nospecialize T::Type{<:VecAdjTrans})
     c = contiguous_axis(parent_type(T))
     if c === nothing
         return nothing
@@ -139,7 +139,7 @@ function contiguous_axis(::Type{T}) where {T<:VecAdjTrans}
         return -One()
     end
 end
-function contiguous_axis(::Type{T}) where {T<:MatAdjTrans}
+function contiguous_axis(@nospecialize T::Type{<:MatAdjTrans})
     c = contiguous_axis(parent_type(T))
     if c === nothing
         return nothing
@@ -149,7 +149,7 @@ function contiguous_axis(::Type{T}) where {T<:MatAdjTrans}
         return StaticInt(3) - c
     end
 end
-function contiguous_axis(::Type{T}) where {T<:PermutedDimsArray}
+function contiguous_axis(@nospecialize T::Type{<:PermutedDimsArray})
     c = contiguous_axis(parent_type(T))
     if c === nothing
         return nothing