Fix inference issues of StructArray constructor on Julia 1.7+ (#209)

* Add fully inferred implementation of _map_params

* Add new inference test

* Add consistency test for _map_params

* Add similar inference fix for map_params

* Add map_params test

* `_map_params_recursive` -> `_map_params_fallback`

Co-authored-by: Pietro Vertechi <pietro.vertechi@protonmail.com>

* `map_params_recursive` -> `map_params_fallback`

Co-authored-by: Pietro Vertechi <pietro.vertechi@protonmail.com>

* Use `fieldtypes` and remove `types_to_tuple`

Co-authored-by: Pietro Vertechi <pietro.vertechi@protonmail.com>

Author: jipolanco

src/utils.jl

@@ -13,13 +13,24 @@ julia> StructArrays.map_params(T -> Complex{T}, Tuple{Int32,Float64})
 Tuple{Complex{Int32},Complex{Float64}}
 ```
 """
-map_params(f, ::Type{Tuple{}}) = Tuple{}
-function map_params(f, ::Type{T}) where {T<:Tuple}
-    tuple_type_cons(f(tuple_type_head(T)), map_params(f, tuple_type_tail(T)))
-end
 map_params(f, ::Type{NamedTuple{names, types}}) where {names, types} =
     NamedTuple{names, map_params(f, types)}
 
+function map_params(f, ::Type{T}) where {T<:Tuple}
+    if @generated
+        types = fieldtypes(T)
+        ex = :(Tuple{})
+        for t ∈ types
+            push!(ex.args, :(f($t)))
+        end
+        ex
+    else
+        map_params_fallback(f, T)
+    end
+end
+
+map_params_fallback(f, ::Type{T}) where {T<:Tuple} = Tuple{map(f, fieldtypes(T))...}
+
 """
     StructArrays._map_params(f, T)
 
@@ -31,13 +42,24 @@ julia> StructArrays._map_params(T -> Complex{T}, Tuple{Int32,Float64})
 (Complex{Int32}, Complex{Float64})
 ```
 """
-_map_params(f, ::Type{Tuple{}}) = ()
-function _map_params(f, ::Type{T}) where {T<:Tuple}
-    (f(tuple_type_head(T)), _map_params(f, tuple_type_tail(T))...)
-end
 _map_params(f::F, ::Type{NamedTuple{names, types}}) where {names, types, F} =
     NamedTuple{names}(_map_params(f, types))
 
+function _map_params(f::F, ::Type{T}) where {T<:Tuple, F}
+    if @generated
+        types = fieldtypes(T)
+        ex = :()
+        for t ∈ types
+            push!(ex.args, :(f($t)))
+        end
+        ex
+    else
+        _map_params_fallback(f, T)
+    end
+end
+
+_map_params_fallback(f, ::Type{T}) where {T<:Tuple} = map(f, fieldtypes(T))
+
 buildfromschema(initializer::F, ::Type{T}) where {T, F} = buildfromschema(initializer, T, staticschema(T))
 
 """

test/runtests.jl

@@ -342,6 +342,7 @@ f_infer() = StructArray{ComplexF64}((rand(2,2), rand(2,2)))
 g_infer() = StructArray([(a=(b="1",), c=2)], unwrap = t -> t <: NamedTuple)
 tup_infer() = StructArray([(1, 2), (3, 4)])
 cols_infer() = StructArray(([1, 2], [1.2, 2.3]))
+nt_infer(nt) = StructArray{typeof(nt)}(undef, 4)
 
 @testset "inferrability" begin
     @inferred f_infer()
@@ -352,6 +353,7 @@ cols_infer() = StructArray(([1, 2], [1.2, 2.3]))
     @test s[1] == (1, 2)
     @test s[2] == (3, 4)
     @inferred cols_infer()
+    @inferred nt_infer((x = 3, y = :a, z = :b))
 end
 
 @testset "propertynames" begin
@@ -879,3 +881,22 @@ end
         @inferred view(u, 1, :)
     end
 end
+
+# Test fallback (non-@generated) variant of _map_params
+@testset "_map_params" begin
+    v = StructArray(rand(ComplexF64, 2, 2))
+    f(T) = similar(v, T)
+    types = Tuple{Int, Float64, ComplexF32, String}
+    A = @inferred StructArrays._map_params(f, types)
+    B = StructArrays._map_params_fallback(f, types)
+    @test typeof(A) === typeof(B)
+end
+
+# Same for map_params
+@testset "map_params" begin
+    types = Tuple{Int, Float64, Int32}
+    f(T) = Complex{T}
+    A = @inferred StructArrays.map_params(f, types)
+    B = StructArrays.map_params_fallback(f, types)
+    @test A === B
+end