improve type stability and runtime performance (#194)

* specialize foreachfield on the function

* improve inference of index_type (necessary for creating views)

Author: ranocha

src/structarray.jl

@@ -32,6 +32,17 @@ function index_type(::Type{T}) where {T<:Tuple}
     S, U = tuple_type_head(T), tuple_type_tail(T)
     IndexStyle(S) isa IndexCartesian ? CartesianIndex{ndims(S)} : index_type(U)
 end
+# Julia v1.7.0-beta3 doesn't seem to specialize `index_type` as defined above
+# for tuple types with "many" elements (three or four, depending on the concrete
+# types). However, we can help the compiler for homogeneous types by defining
+# the specialization below.
+function index_type(::Type{<:NTuple{N, S}}) where {N, S}
+    if IndexStyle(S) isa IndexCartesian
+        return CartesianIndex{ndims(S)}
+    else
+        return Int
+    end
+end
 
 index_type(::Type{StructArray{T, N, C, I}}) where {T, N, C, I} = I
 
@@ -112,7 +123,7 @@ StructVector(args...; kwargs...) = StructArray(args...; kwargs...)
 """
     StructArray{T}(A::AbstractArray; dims, unwrap=FT->FT!=eltype(A))
 
-Construct a `StructArray` from slices of `A` along `dims`. 
+Construct a `StructArray` from slices of `A` along `dims`.
 
 The `unwrap` keyword argument is a function that determines whether to
 recursively convert fields of type `FT` to `StructArray`s.

src/utils.jl

@@ -87,7 +87,7 @@ end
 @generated foreachfield_gen(::S, f, xs::Vararg{Any, L}) where {S<:StructArray, L} =
     _foreachfield(array_names_types(S), L)
 
-foreachfield(f, x::StructArray, xs...) = foreachfield_gen(x, f, x, xs...)
+foreachfield(f::F, x::StructArray, xs::Vararg{Any, N}) where {F, N} = foreachfield_gen(x, f, x, xs...)
 
 """
     StructArrays.iscompatible(::Type{S}, ::Type{V}) where {S, V<:AbstractArray}
@@ -149,7 +149,7 @@ julia> s = StructArray(a=1:3, b = fill("string", 3));
 julia> s_pooled = StructArrays.replace_storage(s) do v
            isbitstype(eltype(v)) ? v : convert(PooledArray, v)
        end
-$(if VERSION < v"1.6-" 
+$(if VERSION < v"1.6-"
     "3-element StructArray(::UnitRange{Int64}, ::PooledArray{String,UInt32,1,Array{UInt32,1}}) with eltype NamedTuple{(:a, :b),Tuple{Int64,String}}:"
 else
         "3-element StructArray(::UnitRange{Int64}, ::PooledVector{String, UInt32, Vector{UInt32}}) with eltype NamedTuple{(:a, :b), Tuple{Int64, String}}:"

test/runtests.jl

@@ -70,6 +70,15 @@ end
     @test s[100] == s[10, 10] == (a=1, b=1)
     s[10, 10] = (a=0, b=0)
     @test s[100] == s[10, 10] == (a=0, b=0)
+
+    # inference for "many" types, both for linear ad Cartesian indexing
+    @inferred StructArrays.index_type(NTuple{2, Vector{Float64}})
+    @inferred StructArrays.index_type(NTuple{3, Matrix{Float64}})
+    @inferred StructArrays.index_type(NTuple{4, Array{Float64, 3}})
+
+    @inferred StructArrays.index_type(NTuple{2, SubArray{Float64, 1, Array{Float64, 2}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}})
+    @inferred StructArrays.index_type(NTuple{3, SubArray{Float64, 1, Array{Float64, 2}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}})
+    @inferred StructArrays.index_type(NTuple{4, SubArray{Float64, 1, Array{Float64, 2}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}})
 end
 
 @testset "replace_storage" begin
@@ -818,9 +827,8 @@ Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{MyArray}}, ::Type{El
 end
 
 @testset "staticarrays" begin
-
     # test that staticschema returns the right things
-    for StaticVectorType = [SVector, MVector, SizedVector]    
+    for StaticVectorType = [SVector, MVector, SizedVector]
         @test StructArrays.staticschema(StaticVectorType{2,Float64}) == Tuple{Float64,Float64}
     end
 
@@ -838,4 +846,11 @@ end
         @test StructArrays.components(x) == ([1., 2.], [1., 2.])
         @test x .+ y == StructArray([StaticArrayType{Tuple{1,2}}(3*ones(1,2) .+ 2*i) for i = 0:1])
     end
+
+    # test type stability of creating views with "many" homogeneous components
+    for n in 1:10
+        u = StructArray(randn(SVector{n, Float64}) for _ in 1:10, _ in 1:5)
+        @inferred view(u, :, 1)
+        @inferred view(u, 1, :)
+    end
 end