Cleanup/fix range and dimension related traits/methods (#196)

axes_types was erroneously assuming that the the parent axis always had an offfset of 1
known_first and known_last work for CartesianIndices now
defined contiguous_axis, stride_rank, and offset1 for StrideIndex
has_dimnames returns StaticBool now. Not breaking because no downstream dependencies yet.

Author: Tokazama

Project.toml

@@ -1,6 +1,6 @@
 name = "ArrayInterface"
 uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
-version = "3.1.27"
+version = "3.1.28"
 
 [deps]
 IfElse = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173"

src/ArrayInterface.jl

@@ -14,9 +14,6 @@ using Base: @propagate_inbounds, tail, OneTo, LogicalIndex, Slice, ReinterpretAr
 
 
 ## utilites for internal use only ##
-_int_or_static_int(::Nothing) = Int
-_int_or_static_int(x::Int) = StaticInt{x}
-
 @static if VERSION >= v"1.7.0-DEV.421"
     using Base: @aggressive_constprop
 else

src/array_index.jl

@@ -194,7 +194,7 @@ struct StrideIndex{N,R,C,S,O} <: ArrayIndex{N}
     offsets::O
 
     function StrideIndex{N,R,C}(s::S, o::O) where {N,R,C,S,O}
-        return new{N,R::NTuple{N,Int},C::Int,S,O}(s, o)
+        return new{N,R::NTuple{N,Int},C,S,O}(s, o)
     end
     function StrideIndex{N,R,C}(a::A) where {N,R,C,A}
         return StrideIndex{N,R,C}(strides(a), offsets(a))

src/axes.jl

@@ -1,4 +1,10 @@
 
+_static_range_type(::Any, ::Any) = OptionallyStaticUnitRange{Int,Int}
+_static_range_type(start::Int, ::Nothing) = OptionallyStaticUnitRange{StaticInt{start},Int}
+function _static_range_type(start::Int, size::Int)
+    OptionallyStaticUnitRange{StaticInt{start},StaticInt{(size - 1) + start}}
+end
+
 """
     axes_types(::Type{T}) -> Type{Tuple{Vararg{AbstractUnitRange{Int}}}}
     axes_types(::Type{T}, dim) -> Type{AbstractUnitRange{Int}}
@@ -54,11 +60,9 @@ end
 @inline function axes_types(::Type{T}) where {N,P,I,T<:SubArray{<:Any,N,P,I}}
     return eachop_tuple(_sub_axis_type, to_parent_dims(T), T)
 end
+
 @inline function _sub_axis_type(::Type{A}, dim::StaticInt) where {T,N,P,I,A<:SubArray{T,N,P,I}}
-    return OptionallyStaticUnitRange{
-        _int_or_static_int(known_first(axes_types(P, dim))),
-        _int_or_static_int(known_length(_get_tuple(I, dim)))
-    }
+    _static_range_type(known_first(axes_types(P, dim)),known_length(_get_tuple(I, dim)))
 end
 
 function axes_types(::Type{R}) where {T,N,S,A,R<:ReinterpretArray{T,N,S,A}}

src/dimensions.jl

@@ -46,11 +46,11 @@ 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),)
 from_parent_dims(::Type{T}) where {T<:MatAdjTrans} = (StaticInt(2), One())
-from_parent_dims(::Type{<:SubArray{T,N,A,I}}) where {T,N,A,I} = _from_sub_dims(A, I)
-@generated function _from_sub_dims(::Type{A}, ::Type{I}) where {A,I<:Tuple}
+from_parent_dims(::Type{<:SubArray{T,N,A,I}}) where {T,N,A,I} = _from_sub_dims(I)
+@generated function _from_sub_dims(::Type{I}) where {I<:Tuple}
     out = Expr(:tuple)
     dim_i = 1
-    for i in 1:ndims(A)
+    for i in 1:length(I.parameters)
         p = I.parameters[i]
         if p <: Integer
             push!(out.args, :(StaticInt(0)))
@@ -103,8 +103,8 @@ 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())
 to_parent_dims(::Type{<:PermutedDimsArray{T,N,I}}) where {T,N,I} = static(Val(I))
-to_parent_dims(::Type{<:SubArray{T,N,A,I}}) where {T,N,A,I} = _to_sub_dims(A, I)
-@generated function _to_sub_dims(::Type{A}, ::Type{I}) where {A,N,I<:Tuple{Vararg{Any,N}}}
+to_parent_dims(::Type{<:SubArray{T,N,A,I}}) where {T,N,A,I} = _to_sub_dims(I)
+@generated function _to_sub_dims(::Type{I}) where {I<:Tuple}
     out = Expr(:tuple)
     n = 1
     for p in I.parameters
@@ -150,16 +150,12 @@ end
 """
     has_dimnames(::Type{T}) -> Bool
 
-Returns `true` if `x` has names for each dimension.
+Returns `static(true)` if `x` has on or more named dimensions.
 """
-@inline has_dimnames(x) = has_dimnames(typeof(x))
-function has_dimnames(::Type{T}) where {T}
-    if parent_type(T) <: T
-        return false
-    else
-        return has_dimnames(parent_type(T))
-    end
-end
+has_dimnames(x) = has_dimnames(typeof(x))
+@inline has_dimnames(::Type{T}) where {T} = _has_dimnames(dimnames(T))
+_has_dimnames(::Tuple{Vararg{StaticSymbol{:_}}}) = static(false)
+_has_dimnames(::Tuple) = static(true)
 
 # this takes the place of dimension names that aren't defined
 const SUnderscore = StaticSymbol(:_)

src/ranges.jl

@@ -1,4 +1,7 @@
 
+_cartesian_index(i::Tuple{Vararg{Int}}) = CartesianIndex(i)
+_cartesian_index(::Any) = nothing
+
 """
     known_first(::Type{T}) -> Union{Int,Nothing}
 
@@ -22,6 +25,10 @@ function known_first(::Type{T}) where {T}
     end
 end
 known_first(::Type{Base.OneTo{T}}) where {T} = one(T)
+function known_first(::Type{T}) where {N,R,T<:CartesianIndices{N,R}}
+    _cartesian_index(ntuple(i -> known_first(R.parameters[i]), Val(N)))
+end
+
 
 """
     known_last(::Type{T}) -> Union{Int,Nothing}
@@ -46,6 +53,9 @@ function known_last(::Type{T}) where {T}
         return known_last(parent_type(T))
     end
 end
+function known_last(::Type{T}) where {N,R,T<:CartesianIndices{N,R}}
+    _cartesian_index(ntuple(i -> known_last(R.parameters[i]), Val(N)))
+end
 
 """
     known_step(::Type{T}) -> Union{Int,Nothing}

src/stridelayout.jl

@@ -64,6 +64,9 @@ function _offsets(x::X, dim::StaticInt{D}) where {X,D}
         return static(start)
     end
 end
+# we can't generate an axis for `StrideIndex` so this is performed manually here
+@inline offsets(x::StrideIndex, dim::Int) = getfield(offsets(x), dim)
+@inline offsets(x::StrideIndex, ::StaticInt{dim}) where {dim} = getfield(offsets(x), dim)
 
 """
     known_offset1(::Type{T}) -> Union{Int,Nothing}
@@ -105,6 +108,8 @@ If no axis is contiguous, it returns a `StaticInt{-1}`.
 If unknown, it returns `nothing`.
 """
 contiguous_axis(x) = contiguous_axis(typeof(x))
+contiguous_axis(::Type{<:StrideIndex{N,R,C}}) where {N,R,C} = static(C)
+contiguous_axis(::Type{<:StrideIndex{N,R,nothing}}) where {N,R} = nothing
 function contiguous_axis(::Type{T}) where {T}
     if parent_type(T) <: T
         return nothing
@@ -197,6 +202,7 @@ function rank_to_sortperm(R::Tuple{Vararg{StaticInt,N}}) where {N}
     return sp
 end
 
+stride_rank(::Type{<:StrideIndex{N,R}}) where {N,R} = static(R)
 stride_rank(x) = stride_rank(typeof(x))
 function stride_rank(::Type{T}) where {T}
     if parent_type(T) <: T

test/array_index.jl

@@ -0,0 +1,20 @@
+
+A = zeros(3, 4, 5);
+A[:] = 1:60
+Ap = @view(PermutedDimsArray(A,(3,1,2))[:,1:2,1])';
+
+ap_index = ArrayInterface.StrideIndex(Ap)
+for x_i in axes(Ap, 1)
+    for y_i in axes(Ap, 2)
+        @test ap_index[x_i, y_i] == ap_index[x_i, y_i]
+    end
+end
+@test @inferred(ArrayInterface.known_offsets(ap_index)) === ArrayInterface.known_offsets(Ap)
+@test @inferred(ArrayInterface.known_offset1(ap_index)) === ArrayInterface.known_offset1(Ap)
+@test @inferred(ArrayInterface.offsets(ap_index, 1)) === ArrayInterface.offset1(Ap)
+@test @inferred(ArrayInterface.offsets(ap_index, static(1))) === ArrayInterface.offset1(Ap)
+@test @inferred(ArrayInterface.known_strides(ap_index)) === ArrayInterface.known_strides(Ap)
+@test @inferred(ArrayInterface.contiguous_axis(ap_index)) == 1
+@test @inferred(ArrayInterface.contiguous_axis(ArrayInterface.StrideIndex{2,(1,2),nothing,NTuple{2,Int},NTuple{2,Int}})) == nothing
+@test @inferred(ArrayInterface.stride_rank(ap_index)) == (1, 3)
+

test/dimensions.jl

@@ -10,7 +10,6 @@ struct NamedDimsWrapper{L,T,N,P<:AbstractArray{T,N}} <: ArrayInterface.AbstractA
     NamedDimsWrapper{L}(p) where {L} = new{L,eltype(p),ndims(p),typeof(p)}(p)
 end
 ArrayInterface.parent_type(::Type{T}) where {P,T<:NamedDimsWrapper{<:Any,<:Any,<:Any,P}} = P
-ArrayInterface.has_dimnames(::Type{T}) where {T<:NamedDimsWrapper} = true
 ArrayInterface.dimnames(::Type{T}) where {L,T<:NamedDimsWrapper{L}} = static(Val(L))
 function ArrayInterface.dimnames(::Type{T}, dim) where {L,T<:NamedDimsWrapper{L}}
     if ndims(T) < dim
@@ -19,7 +18,6 @@ function ArrayInterface.dimnames(::Type{T}, dim) where {L,T<:NamedDimsWrapper{L}
         return static(L[dim])
     end
 end
-ArrayInterface.has_dimnames(::Type{T}) where {T<:NamedDimsWrapper} = true
 Base.parent(x::NamedDimsWrapper) = x.parent
 
 @testset "dimension permutations" begin
@@ -90,18 +88,17 @@ end
     @test_throws ErrorException ArrayInterface.order_named_inds(n2, (x=30, y=20, z=40))
 end
 
-val_has_dimnames(x) = Val(ArrayInterface.has_dimnames(x))
 
 @testset "dimnames" begin
     d = (static(:x), static(:y))
     x = NamedDimsWrapper{d}(ones(2,2));
     y = NamedDimsWrapper{(:x,)}(ones(2));
     dnums = ntuple(+, length(d))
-    @test @inferred(val_has_dimnames(x)) === Val(true)
-    @test @inferred(ArrayInterface.has_dimnames(ones(2,2))) === false
-    @test @inferred(ArrayInterface.has_dimnames(Array{Int,2})) === false
-    @test @inferred(val_has_dimnames(typeof(x))) === Val(true)
-    @test @inferred(val_has_dimnames(typeof(view(x, :, 1, :)))) === Val(true)
+    @test @inferred(ArrayInterface.has_dimnames(x)) == true
+    @test @inferred(ArrayInterface.has_dimnames(ones(2,2))) == false
+    @test @inferred(ArrayInterface.has_dimnames(Array{Int,2})) == false
+    @test @inferred(ArrayInterface.has_dimnames(typeof(x))) == true
+    @test @inferred(ArrayInterface.has_dimnames(typeof(view(x, :, 1, :)))) == true
     @test @inferred(dimnames(x)) === d
     @test @inferred(dimnames(parent(x))) === (static(:_), static(:_))
     @test @inferred(dimnames(x')) === reverse(d)

test/ranges.jl

@@ -57,6 +57,15 @@
     @test isnothing(@inferred(ArrayInterface.known_last(typeof(1:4))))
     @test isone(@inferred(ArrayInterface.known_last(typeof(StaticInt(-1):StaticInt(2):StaticInt(1)))))
 
+    # CartesianIndices
+    CI = CartesianIndices((2, 2))
+    @test @inferred(ArrayInterface.known_first(typeof(CI))) == CartesianIndex(1, 1)
+    @test @inferred(ArrayInterface.known_last(typeof(CI))) == nothing
+
+    CI = CartesianIndices((static(1):static(2), static(1):static(2)))
+    @test @inferred(ArrayInterface.known_first(typeof(CI))) == CartesianIndex(1, 1)
+    @test @inferred(ArrayInterface.known_last(typeof(CI))) == CartesianIndex(2, 2)
+
     @test isnothing(@inferred(ArrayInterface.known_step(typeof(1:0.2:4))))
     @test isone(@inferred(ArrayInterface.known_step(1:4)))
     @test isone(@inferred(ArrayInterface.known_step(typeof(1:4))))