@inline StrideIndex (#343)

* Use generated `map_tuple_type`,`flatten_tuples` for pre-v1.8

Failure to elide certain calls on pre-v1.8 was causing performance problems downstream, so we keep the less desirable code in Julia versions that require it for performance

Co-authored-by: Zachary P. Christensen <[email protected]>

Author: chriselrod

Project.toml

@@ -1,6 +1,6 @@
 name = "ArrayInterface"
 uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
-version = "6.0.22"
+version = "6.0.23"
 
 [deps]
 ArrayInterfaceCore = "30b0a656-2188-435a-8636-2ec0e6a096e2"

lib/ArrayInterfaceCore/Project.toml

@@ -1,6 +1,6 @@
 name = "ArrayInterfaceCore"
 uuid = "30b0a656-2188-435a-8636-2ec0e6a096e2"
-version = "0.1.17"
+version = "0.1.18"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

lib/ArrayInterfaceCore/src/ArrayInterfaceCore.jl

@@ -35,8 +35,23 @@ julia> ArrayInterfaceCore.map_tuple_type(sqrt, Tuple{1,4,16})
 
 ```
 """
-@inline function map_tuple_type(f, @nospecialize(T::Type))
-    ntuple(i -> f(fieldtype(T, i)), Val{fieldcount(T)}())
+function map_tuple_type end
+if VERSION >= v"1.8"
+    @inline function map_tuple_type(f, @nospecialize(T::Type))
+        ntuple(i -> f(fieldtype(T, i)), Val{fieldcount(T)}())
+    end
+else
+    function map_tuple_type(f::F, ::Type{T}) where {F,T<:Tuple}
+        if @generated
+            t = Expr(:tuple)
+            for i in 1:fieldcount(T)
+                push!(t.args, :(f($(fieldtype(T, i)))))
+            end
+            Expr(:block, Expr(:meta, :inline), t)
+        else
+            Tuple(f(fieldtype(T, i)) for i in 1:fieldcount(T))
+        end
+    end
 end
 
 """
@@ -58,20 +73,45 @@ julia> ArrayInterfaceCore.flatten_tuples((1, (2, (3,))))
 
 ```
 """
-function flatten_tuples(t::Tuple)
-    fields = _new_field_positions(t)
-    ntuple(Val{nfields(fields)}()) do k
-        i, j = getfield(fields, k)
-        i = length(t) - i
-        @inbounds j === 0 ? getfield(t, i) : getfield(getfield(t, i), j)
+function flatten_tuples end
+if VERSION >= v"1.8"
+    function flatten_tuples(t::Tuple)
+        fields = _new_field_positions(t)
+        ntuple(Val{nfields(fields)}()) do k
+            i, j = getfield(fields, k)
+            i = length(t) - i
+            @inbounds j === 0 ? getfield(t, i) : getfield(getfield(t, i), j)
+        end
+    end
+    _new_field_positions(::Tuple{}) = ()
+    @nospecialize
+    _new_field_positions(x::Tuple) = (_fl1(x, x[1])..., _new_field_positions(Base.tail(x))...)
+    _fl1(x::Tuple, x1::Tuple) = ntuple(Base.Fix1(tuple, length(x) - 1), Val(length(x1)))
+    _fl1(x::Tuple, x1) = ((length(x) - 1, 0),)
+    @specialize
+else
+    @inline function flatten_tuples(t::Tuple)
+        if @generated
+            texpr = Expr(:tuple)
+            for i in 1:fieldcount(t)
+                p = fieldtype(t, i)
+                if p <: Tuple
+                    for j in 1:fieldcount(p)
+                        push!(texpr.args, :(@inbounds(getfield(getfield(t, $i), $j))))
+                    end
+                else
+                    push!(texpr.args, :(@inbounds(getfield(t, $i))))
+                end
+            end
+            Expr(:block, Expr(:meta, :inline), texpr)
+        else
+            _flatten(t)
+        end
     end
+    _flatten(::Tuple{}) = ()
+    @inline _flatten(t::Tuple{Any,Vararg{Any}}) = (getfield(t, 1), _flatten(Base.tail(t))...)
+    @inline _flatten(t::Tuple{Tuple,Vararg{Any}}) = (getfield(t, 1)..., _flatten(Base.tail(t))...)
 end
-_new_field_positions(::Tuple{}) = ()
-@nospecialize
-_new_field_positions(x::Tuple) = (_fl1(x, x[1])..., _new_field_positions(Base.tail(x))...)
-_fl1(x::Tuple, x1::Tuple) = ntuple(Base.Fix1(tuple, length(x) - 1), Val(length(x1)))
-_fl1(x::Tuple, x1) = ((length(x) - 1, 0),)
-@specialize
 
 """
     parent_type(::Type{T}) -> Type

src/array_index.jl

@@ -8,13 +8,13 @@ struct StrideIndex{N,R,C,S,O} <: ArrayIndex{N}
     strides::S
     offsets::O
 
-    function StrideIndex{N,R,C}(s::S, o::O) where {N,R,C,S,O}
+    @inline function StrideIndex{N,R,C}(s::S, o::O) where {N,R,C,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}
+    @inline function StrideIndex{N,R,C}(a::A) where {N,R,C,A}
         return StrideIndex{N,R,C}(strides(a), offsets(a))
     end
-    function StrideIndex(a::A) where {A}
+    @inline function StrideIndex(a::A) where {A}
         return StrideIndex{ndims(A),known(stride_rank(A)),known(contiguous_axis(A))}(a)
     end
 end