consolidate and clean-up code

Author: Tokazama

lib/ArrayInterfaceCore/src/ArrayInterfaceCore.jl

@@ -678,15 +678,76 @@ Provides basic trait information for each index type in in the tuple `T`. `NI`,
 [`is_splat_index`](@ref) (respectively) for each field of `T`.
 """
 struct IndicesInfo{NI,NS,IS} end
-IndicesInfo(x::Tuple) = IndicesInfo(typeof(x))
+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))
-        #findfirst(==(1), map_tuple_type(is_splat_index, T))
     }()
 end
+function IndicesInfo(@nospecialize T::Type{<:SubArray})
+    indices_to_dims(IndicesInfo(fieldtype(T, :indices)), parent_type(T))
+end
+
+@inline function indices_to_dims(@nospecialize(I::Type{<:Tuple}), @nospecialize(T::Type))
+    indices_to_dims(IndicesInfo(I), T)
+end
+@inline function indices_to_dims(::IndicesInfo{(1,),NS,nothing}, @nospecialize(T::Type)) where {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
+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}()
+    else
+        IndicesInfo{_accum_dims(cumsum(NI), NI), _accum_dims(cumsum(NS), NS), nothing}()
+    end
+end
+@inline function indices_to_dims(::IndicesInfo{NI,NS,SI}, @nospecialize(T::Type)) where {NI,NS,SI}
+    nsplat = ndims(T) - sum(NI)
+    if nsplat === 0
+        indices_to_dims(IndicesInfo{NI,NS,nothing}(), T)
+    else
+        splatmul = max(0, nsplat + 1)
+        indices_to_dims(IndicesInfo{_map_splats(splatmul, SI, NI),_map_splats(splatmul, SI, NS),nothing}(), T)
+    end
+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)
+        if nd_i === 0
+            0
+        elseif nd_i === 1
+            getfield(csdims, i)
+        else
+            ntuple(Base.Fix1(+, getfield(csdims, i) - nd_i), nd_i)
+        end
+    end
+end
 
 """
     instances_do_not_alias(::Type{T}) -> Bool

lib/ArrayInterfaceCore/test/runtests.jl

@@ -1,7 +1,7 @@
 using ArrayInterfaceCore
 using ArrayInterfaceCore: zeromatrix
 import ArrayInterfaceCore: has_sparsestruct, findstructralnz, fast_scalar_indexing, lu_instance,
-        parent_type, zeromatrix
+        parent_type, zeromatrix, IndicesInfo, indices_to_dims
 using Base: setindex
 using LinearAlgebra
 using Random
@@ -293,3 +293,36 @@ end
   @test !ArrayInterfaceCore.indices_do_not_alias(typeof(view(fill(rand(4,4),4,4)', 2:3, 1:2)))
   @test !ArrayInterfaceCore.indices_do_not_alias(typeof(view(rand(4,4)', StepRangeLen(1,0,5), 1:2)))
 end
+
+@testset "indices_to_dims" 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}()
+
+    @test @inferred(indices_to_dims(Tuple{Vector{Int}}, Vector{Int})) == IndicesInfo{(1,), (1,), nothing}()
+
+    @test @inferred(indices_to_dims(Tuple{Vector{Int}}, Matrix{Int})) == IndicesInfo{(:,), (1,), nothing}()
+
+    @test @inferred(indices_to_dims(Tuple{SplatFirst}, Vector{Int})) == IndicesInfo{(1,), (1,), nothing}()
+
+    @test @inferred(indices_to_dims(Tuple{SplatFirst}, Matrix{Int})) == IndicesInfo{((1,2),), ((1, 2),), nothing}()
+
+    @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(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(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(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(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}()
+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
+    map_tuple_type, flatten_tuples, GetIndex, indices_to_dims
 
 # ArrayIndex subtypes and methods
 import ArrayInterfaceCore: ArrayIndex, MatrixIndex, VectorIndex, BidiagonalIndex, TridiagonalIndex

src/axes.jl

@@ -114,16 +114,10 @@ axes(A::ReshapedArray) = Base.axes(A)
     map(GetIndex{false}(axes(parent(x))), to_parent_dims(x))
 end
 axes(A::VecAdjTrans) = (SOneTo{1}(), axes(parent(A), 1))
-@inline function axes(x::SubArray)
-    flatten_tuples(map((i, d) -> _sub_axis(parent(x), i, d), x.indices, dimsmap(x)))
-end
-_sub_axis(x, idx, ::Tuple{StaticInt{0}}) = ()
-function _sub_axis(x, idx::Base.Slice{Base.OneTo{Int}}, dm::Tuple{StaticInt,Any})
-    sz = known_size(x, getfield(dm, 2))
-    sz === nothing ? axes(idx) : StaticInt(1):StaticInt(sz)
-end
-_sub_axis(x, idx, ::Tuple{Any,Any}) = axes(idx)
 
+@inline axes(x::SubArray) = flatten_tuples(map(Base.Fix1(_sub_axes, x), sub_axes_map(typeof(x))))
+@inline _sub_axes(x::SubArray, ::Pair{StaticSymbol{:parent},StaticInt{s}}) where {s} = StaticInt(1):StaticInt(s)
+_sub_axes(x::SubArray, ::StaticInt{index}) where {index} = axes(getfield(x.indices, index))
 
 @inline axes(A, dim) = _axes(A, to_dims(A, dim))
 @inline _axes(A, dim::Int) = dim > ndims(A) ? OneTo(1) : getfield(axes(A), dim)