refactor foreach_expr (JuliaArrays#21)

* refactor foreach_expr

* use integers for tuples and symbols otherwise

* test tuple push and setindex

* ups

* move to utils

Author: Pietro Vertechi

src/structarray.jl

@@ -83,27 +83,19 @@ end
     return :($(Expr(:tuple, [QuoteNode(Symbol("x$f")) for f in fieldnames(T)]...)))
 end
 
+@inline getfieldindex(v::Tuple, field::Symbol, index::Integer) = getfield(v, index)
+@inline getfieldindex(v, field::Symbol, index::Integer) = getproperty(v, field)
 
 @generated function Base.push!(s::StructArray{T, 1}, vals) where {T}
-    args = []
-    for key in fields(T)
-        field = Expr(:., :s, Expr(:quote, key))
-        val = Expr(:., :vals, Expr(:quote, key))
-        push!(args, :(push!($field, $val)))
-    end
-    push!(args, :s)
-    Expr(:block, args...)
+    exprs = foreach_expr((args...) -> Expr(:call, :push!, args...), T, :s, :vals)
+    push!(exprs, :s)
+    Expr(:block, exprs...)
 end
 
 @generated function Base.append!(s::StructArray{T, 1}, vals) where {T}
-    args = []
-    for key in fields(T)
-        field = Expr(:., :s, Expr(:quote, key))
-        val = Expr(:., :vals, Expr(:quote, key))
-        push!(args, :(append!($field, $val)))
-    end
-    push!(args, :s)
-    Expr(:block, args...)
+    exprs = foreach_expr((args...) -> Expr(:call, :append!, args...), T, :s, :vals)
+    push!(exprs, :s)
+    Expr(:block, exprs...)
 end
 
 function Base.cat(args::StructArray...; dims)

src/utils.jl

@@ -7,29 +7,29 @@ eltypes(::Type{NamedTuple{K, V}}) where {K, V} = eltypes(V)
 
 Base.@pure SkipConstructor(::Type) = false
 
-@generated function get_ith(s::StructArray{T}, I...) where {T}
-    args = []
-    for key in fields(T)
-        field = Expr(:., :s, Expr(:quote, key))
-        push!(args, :($field[I...]))
+function foreach_expr(f, T, args...)
+    exprs = []
+    for (ind, key) in enumerate(fields(T))
+        new_args = (Expr(:call, :getfieldindex, arg, Expr(:quote, key), ind) for arg in args)
+        push!(exprs, f(new_args...))
     end
+    exprs
+end
+
+@generated function get_ith(s::StructArray{T}, I...) where {T}
+    exprs = foreach_expr(field -> :($field[I...]), T, :s)
     return quote
         @boundscheck checkbounds(s, I...)
-        @inbounds $(Expr(:call, :createinstance, :T, args...))
+        @inbounds $(Expr(:call, :createinstance, :T, exprs...))
     end
 end
 
 @generated function set_ith!(s::StructArray{T}, vals, I...) where {T}
-    args = []
-    for key in fields(T)
-        field = Expr(:., :s, Expr(:quote, key))
-        val = Expr(:., :vals, Expr(:quote, key))
-        push!(args, :($field[I...] = $val))
-    end
-    push!(args, :s)
+    exprs = foreach_expr((field, val) -> :($field[I...] = $val), T, :s, :vals)
+    push!(exprs, :s)
     return quote
         @boundscheck checkbounds(s, I...)
-        @inbounds $(Expr(:block, args...))
+        @inbounds $(Expr(:block, exprs...))
     end
 end
 

test/runtests.jl

@@ -30,6 +30,12 @@ end
     @test Base.getproperty(s, 2) == ["test"]
     t = StructArray{Tuple{Int, Float64}}([1], [1.2])
     @test t[1] == (1, 1.2)
+
+    t[1] = (2, 3)
+    @test t[1] == (2, 3.0)
+    push!(t, (1, 2))
+    @test getproperty(t, 1) == [2, 1]
+    @test getproperty(t, 2) == [3.0, 2.0]
 end
 
 @testset "kwargs constructor" begin