Merge pull request #167 from JuliaArrays/static_axis_1_for_adj_or_trans

Static axis 1 for adj or trans

Author: chriselrod

Project.toml

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

src/ArrayInterface.jl

@@ -1163,8 +1163,9 @@ function __init__()
         function stride_rank(::Type{A}) where {A<:OffsetArrays.OffsetArray}
             return stride_rank(parent_type(A))
         end
-        ArrayInterface.axes(A::OffsetArrays.OffsetArray) = Base.axes(A)
-        ArrayInterface.axes(A::OffsetArrays.OffsetArray, dim::Integer) = Base.axes(A, dim)
+        @inline axes(A::OffsetArrays.OffsetArray) = Base.axes(A)
+        @inline _axes(A::OffsetArrays.OffsetArray, dim::Integer) = Base.axes(A, dim)
+        @inline axes(A::OffsetArrays.OffsetArray{T,N}, ::StaticInt{M}) where {T,M,N} = _axes(A, StaticInt{M}(), gt(StaticInt{M}(),StaticInt{N}()))
     end
 end
 

src/axes.jl

@@ -120,35 +120,41 @@ similar_type(::Type{OptionallyStaticUnitRange{One,StaticInt{N1}}}, ::Type{Int},
 
 Return a valid range that maps to each index along dimension `d` of `A`.
 """
-axes(a, dim) = axes(a, to_dims(a, dim))
-axes(a, dims::Tuple{Vararg{Any,K}}) where {K} = (axes(a, first(dims)), axes(a, tail(dims))...)
-axes(a, dims::Tuple{T}) where {T} = (axes(a, first(dims)), )
-axes(a, ::Tuple{}) = ()
-function axes(a::A, dim::Integer) where {A}
+@inline axes(a, dim) = axes(a, to_dims(a, dim))
+@inline axes(a, dims::Tuple{Vararg{Any,K}}) where {K} = (axes(a, first(dims)), axes(a, tail(dims))...)
+@inline axes(a, dims::Tuple{T}) where {T} = (axes(a, first(dims)), )
+@inline axes(a, ::Tuple{}) = ()
+@inline function _axes(a::A, dim::Integer) where {A}
     if parent_type(A) <: A
         return Base.axes(a, Int(dim))
     else
-        return axes(parent(a), to_parent_dims(A, dim))
+        return _axes(parent(a), to_parent_dims(A, dim))
     end
 end
-function axes(A::CartesianIndices{N}, dim::Integer) where {N}
+@inline function _axes(A::CartesianIndices{N}, dim::Integer) where {N}
     if dim > N
         return static(1):static(1)
     else
         return getfield(axes(A), Int(dim))
     end
 end
-function axes(A::LinearIndices{N}, dim::Integer) where {N}
+@inline function _axes(A::LinearIndices{N}, dim::Integer) where {N}
     if dim > N
         return static(1):static(1)
     else
         return getfield(axes(A), Int(dim))
     end
 end
+@inline _axes(::LinearAlgebra.AdjOrTrans{T,V}, ::One) where {T,V<:AbstractVector} = One():One()
+@inline axes(A::AbstractArray, dim::Integer) = _axes(A, dim, False())
+@inline axes(A::AbstractArray{T,N}, ::StaticInt{M}) where {T,N,M} = _axes(A, StaticInt{M}(), gt(StaticInt{M}(),StaticInt{N}()))
+@inline _axes(::Any, ::Any, ::True) = One():One()
+@inline _axes(A::AbstractArray, dim, ::False) = _axes(A, dim)
 
-axes(A::SubArray, dim::Integer) = Base.axes(A, Int(dim))  # TODO implement ArrayInterface version
-axes(A::ReinterpretArray, dim::Integer) = Base.axes(A, Int(dim))  # TODO implement ArrayInterface version
-axes(A::Base.ReshapedArray, dim::Integer) = Base.axes(A, Int(dim))  # TODO implement ArrayInterface version
+
+@inline _axes(A::SubArray, dim::Integer) = Base.axes(A, Int(dim))  # TODO implement ArrayInterface version
+@inline _axes(A::ReinterpretArray, dim::Integer) = Base.axes(A, Int(dim))  # TODO implement ArrayInterface version
+@inline _axes(A::Base.ReshapedArray, dim::Integer) = Base.axes(A, Int(dim))  # TODO implement ArrayInterface version
 
 """
     axes(A)

test/runtests.jl

@@ -835,6 +835,12 @@ end
     @test axes(Base.Slice(StaticInt(2):4)) === (Base.IdentityUnitRange(StaticInt(2):4),)
     @test Base.axes1(ArrayInterface.indices(ones(2,2))) === StaticInt(1):4
     @test Base.axes1(Base.Slice(StaticInt(2):4)) === Base.IdentityUnitRange(StaticInt(2):4)
+
+    x = vec(A23); y = vec(A32);
+    @test ArrayInterface.indices((x',y'),StaticInt(1)) === Base.Slice(StaticInt(1):StaticInt(1))
+    @test ArrayInterface.axes(x',StaticInt(1)) === StaticInt(1):StaticInt(1)
+    @test ArrayInterface.indices((x,y),StaticInt(2)) === Base.Slice(StaticInt(1):StaticInt(1))
+    @test ArrayInterface.axes(x,StaticInt(2)) === StaticInt(1):StaticInt(1)
 end
 
 @testset "insert/deleteat" begin