specialize axes(S, dim) to return SOneTo

[jishnub]

Currently

julia> s = SA[1,2];

julia> axes(s,2)
Base.OneTo(1)

After this PR

julia> axes(s,2)
SOneTo(1)

As a consequence,

julia> reshape(s, axes(s,1), axes(s,2))
2×1 SMatrix{2, 1, Int64, 2} with indices SOneTo(2)×SOneTo(1):
 1
 2

returns a StaticArray

[mateuszbaran]

Thanks, this looks like a good idea.

[jishnub]

Gentle bump

[mateuszbaran]

I'd prefer to have a second opinion here before merging this -- if the axis number isn't constant-propagated, this might be considered a less desirable behavior than the current one.

[chriselrod]

For reference, the behavior of StrideArrays.jl:

julia> using StrideArrays, Static, ArrayInterface

julia> A = @StrideArray rand(10, 10);

julia> B = @StrideArray rand(8, 10);

julia> ArrayInterface.axes(A,1)
Base.OneTo(static(10))

julia> ArrayInterface.axes(B,1)
Base.OneTo(8)

julia> ArrayInterface.axes(B,StaticInt(1))
Base.OneTo(static(8))

julia> ArrayInterface.axes(B,StaticInt(2))
Base.OneTo(static(10))

For A, it'll return a static axis because both axes are the same, meaning there are no type stability issues.
For B, it'll return a dynamic axis, unless the index itself is static, guaranteeing type stability.

So one possibility is to return an SOneTo if an array is "square", or if the argument is a Static.StaticInt.

[mateuszbaran]

Good point, restricting this to arrays where all axes have the same length removes the issue with type stability.

[mcabbott]

The top message's case is for d > ndims(A), where being square won't save you:

julia> ArrayInterface.axes(A,3) |> dump
ArrayInterface.OptionallyStaticUnitRange{StaticInt{1}, StaticInt{1}}
  start: StaticInt{1} static(1)
  stop: StaticInt{1} static(1)

julia> ArrayInterface.axes(A,2) |> dump
Base.OneTo{StaticInt{10}}
  stop: StaticInt{10} static(10)

I have a vague memory that this PR grew out of thinking about adjoint and permutedims, which when acting on a vector might be viewed as putting its (trivial) 2nd axis first.

[jishnub]

Yes indeed, this came out of permutedims. Perhaps I'm missing something, but I don't see how this changes much from a type-instability point-of-view. The current implementation falls back to axes(A, d) = axes(A)[d] which we replace by axes(A, d) = _axes(A)[d]. If d is not constant-propagated, both the present and the proposed behavior suffer from similar type-instabilities, if anything the proposed behavior may be seen as somewhat better in the special case where at least one axis has a length of 1. The difference proposed here replaces the return type of the trivial axis from Base.OneTo{Int} to SOneTo{1}. Is this the concern here?

Using the nightly build,
On master

julia> @code_warntype axes(a, 1)
MethodInstance for axes(::SMatrix{1, 2, Int64, 2}, ::Int64)
  from axes(A::AbstractArray{T, N}, d) where {T, N} in Base at abstractarray.jl:72
Static Parameters
  T = Int64
  N = 2
Arguments
  #self#::Core.Const(axes)
  A::SMatrix{1, 2, Int64, 2}
  d::Int64
Body::Union{SOneTo{1}, SOneTo{2}, Base.OneTo{Int64}}
1 ─      nothing
│   %2 = Core.typeassert(d, Base.Integer)::Int64
│   %3 = (%2 <= $(Expr(:static_parameter, 2)))::Bool
└──      goto #3 if not %3
2 ─ %5 = Base.axes(A)::Core.Const((SOneTo(1), SOneTo(2)))
│   %6 = Base.getindex(%5, d)::Union{SOneTo{1}, SOneTo{2}}
└──      return %6
3 ─ %8 = Base.OneTo(1)::Core.Const(Base.OneTo(1))
└──      return %8

This PR

julia> @code_warntype axes(a, 1)
MethodInstance for axes(::SMatrix{1, 2, Int64, 2}, ::Int64)
  from axes(s::StaticArray, d) in StaticArrays at /home/jishnu/Dropbox/JuliaPackages/StaticArrays.jl/src/abstractarray.jl:15
Arguments
  #self#::Core.Const(axes)
  s::SMatrix{1, 2, Int64, 2}
  d::Int64
Body::Union{SOneTo{1}, SOneTo{2}}
1 ─ %1 = StaticArrays.ndims(s)::Core.Const(2)
│   %2 = (d <= %1)::Bool
└──      goto #3 if not %2
2 ─ %4 = StaticArrays.Size(s)::Core.Const(Size(1, 2))
│   %5 = StaticArrays._axes(%4)::Core.Const((SOneTo(1), SOneTo(2)))
│   %6 = Base.getindex(%5, d)::Union{SOneTo{1}, SOneTo{2}}
└──      return %6
3 ─ %8 = Core.apply_type(StaticArrays.SOneTo, 1)::Core.Const(SOneTo{1})
│   %9 = (%8)()::Core.Const(SOneTo(1))
└──      return %9

The return type inferred is narrower in this case (as certain axes have the same type as the trivial ones).

In general

on master

julia> @code_warntype axes(a, 1)
MethodInstance for axes(::SMatrix{3, 2, Int64, 6}, ::Int64)
  from axes(A::AbstractArray{T, N}, d) where {T, N} in Base at abstractarray.jl:72
Static Parameters
  T = Int64
  N = 2
Arguments
  #self#::Core.Const(axes)
  A::SMatrix{3, 2, Int64, 6}
  d::Int64
Body::Union{SOneTo{3}, SOneTo{2}, Base.OneTo{Int64}}
1 ─      nothing
│   %2 = Core.typeassert(d, Base.Integer)::Int64
│   %3 = (%2 <= $(Expr(:static_parameter, 2)))::Bool
└──      goto #3 if not %3
2 ─ %5 = Base.axes(A)::Core.Const((SOneTo(3), SOneTo(2)))
│   %6 = Base.getindex(%5, d)::Union{SOneTo{3}, SOneTo{2}}
└──      return %6
3 ─ %8 = Base.OneTo(1)::Core.Const(Base.OneTo(1))
└──      return %8

This PR

julia> @code_warntype axes(a, 1)
MethodInstance for axes(::SMatrix{3, 2, Int64, 6}, ::Int64)
  from axes(s::StaticArray, d) in StaticArrays at /home/jishnu/Dropbox/JuliaPackages/StaticArrays.jl/src/abstractarray.jl:15
Arguments
  #self#::Core.Const(axes)
  s::SMatrix{3, 2, Int64, 6}
  d::Int64
Body::Union{SOneTo{3}, SOneTo{2}, SOneTo{1}}
1 ─ %1 = StaticArrays.ndims(s)::Core.Const(2)
│   %2 = (d <= %1)::Bool
└──      goto #3 if not %2
2 ─ %4 = StaticArrays.Size(s)::Core.Const(Size(3, 2))
│   %5 = StaticArrays._axes(%4)::Core.Const((SOneTo(3), SOneTo(2)))
│   %6 = Base.getindex(%5, d)::Union{SOneTo{3}, SOneTo{2}}
└──      return %6
3 ─ %8 = Core.apply_type(StaticArrays.SOneTo, 1)::Core.Const(SOneTo{1})
│   %9 = (%8)()::Core.Const(SOneTo(1))
└──      return %9

The type-instability is similar in both master and in this PR.