Improve inferability of offsets by using StaticInt for indices. Resolves #92

Author: chriselrod

Project.toml

@@ -1,6 +1,6 @@
 name = "ArrayInterface"
 uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
-version = "2.14.0"
+version = "2.14.1"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/dimensions.jl

@@ -151,22 +151,22 @@ julia> ArrayInterface.size(A)
 (StaticInt{3}(), StaticInt{4}())
 ```
 """
-size(A) = Base.size(A)
-size(A, d) = Base.size(A, to_dims(A, d))
-size(x::LinearAlgebra.Adjoint{T,V}) where {T, V <: AbstractVector{T}} = (One(), static_length(x))
-size(x::LinearAlgebra.Transpose{T,V}) where {T, V <: AbstractVector{T}} = (One(), static_length(x))
+@inline size(A) = Base.size(A)
+@inline size(A, d) = Base.size(A, to_dims(A, d))
+@inline size(x::LinearAlgebra.Adjoint{T,V}) where {T, V <: AbstractVector{T}} = (One(), static_length(x))
+@inline size(x::LinearAlgebra.Transpose{T,V}) where {T, V <: AbstractVector{T}} = (One(), static_length(x))
 
 """
     axes(A, d)
 
 Return a valid range that maps to each index along dimension `d` of `A`.
 """
-axes(A, d) = Base.axes(A, to_dims(A, d))
+@inline axes(A, d) = Base.axes(A, to_dims(A, d))
 
 """
     axes(A)
 
 Return a tuple of ranges where each range maps to each element along a dimension of `A`.
 """
-axes(A) = Base.axes(A)
+@inline axes(A) = Base.axes(A)
 

src/stridelayout.jl

@@ -337,10 +337,11 @@ end
 # @inline offsets(A::AbstractArray{<:Any,N}) where {N} = ntuple(n -> offsets(A, n), Val{N}())
 # Explicit tuple needed for inference.
 @generated function offsets(A::AbstractArray{<:Any,N}) where {N}
-    quote
-        $(Expr(:meta, :inline))
-        Base.Cartesian.@ntuple $N n -> offsets(A, n)
+    t = Expr(:tuple)
+    for n ∈ 1:N
+        push!(t.args, :(offsets(A, StaticInt{$n}())))
     end
+    Expr(:block, Expr(:meta, :inline), t)
 end
 
 @inline size(B::Union{Transpose{T,A},Adjoint{T,A}}) where {T,A<:AbstractMatrix{T}} = permute(size(parent(B)), Val{(2,1)}())