Recursive merge for named tuples

Specific Domains Probabilistic programming
7

Thanks for the pointer @JeffreySarnoff

This stuff gets really tricky to me, there are so many subtleties. I guess the important thing is to identify places where the compiler is really struggling, and generate simpler code to make its job easier.

For example, in the keyunion code, we might instead just do union(keys(x), keys(y)). That builds an array, and the compiler can’t reduce it too much:

julia> x = (a=1, b=2)
(a = 1, b = 2)

julia> y = (b=3, c=4)
(b = 3, c = 4)

julia> @code_typed union(keys(x), keys(y))
CodeInfo(
1 ─ %1 = $(Expr(:foreigncall, :(:jl_alloc_array_1d), Array{Symbol,1}, svec(Any, Int64), 0, :(:ccall), Array{Symbol,1}, 0, 0))::Array{Symbol,1}
│   %2 = (getfield)(sets, 1)::Tuple{Symbol,Symbol}
│   %3 = invoke Base.union!(%1::Array{Symbol,1}, _2::Tuple{Symbol,Symbol}, %2::Tuple{Symbol,Symbol})::Array{Symbol,1}
└──      return %3
) => Array{Symbol,1}

I think part of the weirdness here comes from the dynamic types. So in this case we can generate code to return a tuple, and it should do even better.

julia> @generated function keyunion(x::NamedTuple{Kx, Tx},y::NamedTuple{Ky,Ty}) where {Kx,Tx,Ky,Ty}
           Tuple(union(Kx,Ky))
       end

julia> @code_typed keyunion(x,y)
CodeInfo(
1 ─     return (:a, :b, :c)
) => Tuple{Symbol,Symbol,Symbol}

Can’t do much better than that :slight_smile:

But then sometimes the compiler has no trouble, e.g.

julia> fieldnames1(x::NamedTuple{N}) where {N} = N
fieldnames1 (generic function with 1 method)

julia> @code_llvm fieldnames1((a=1,b=2,c=3))
;  @ REPL[60]:1 within `fieldnames1'
define void @julia_fieldnames1_4718([3 x %jl_value_t*]* noalias nocapture sret %0, [3 x i64]* nocapture nonnull readonly dereferenceable(24) %1) {
top:
  %2 = bitcast [3 x %jl_value_t*]* %0 to i8*
  call void @llvm.memcpy.p0i8.p0i8.i64(i8* nonnull align 8 dereferenceable(24) %2, i8* nonnull align 16 dereferenceable(24) inttoptr (i64 140503769783536 to i8*), i64 24, i1 false)
  ret void
}

julia> @generated function fieldnames2(x::T) where {N,S, T<:NamedTuple{N,S}}
           N
       end
fieldnames2 (generic function with 1 method)

julia> @code_llvm fieldnames2((a=1,b=2,c=3))
;  @ REPL[62]:1 within `fieldnames2'
define void @julia_fieldnames2_4720([3 x %jl_value_t*]* noalias nocapture sret %0, [3 x i64]* nocapture nonnull readonly dereferenceable(24) %1) {
top:
; ┌ @ REPL[62]:1 within `macro expansion'
   %2 = bitcast [3 x %jl_value_t*]* %0 to i8*
   call void @llvm.memcpy.p0i8.p0i8.i64(i8* nonnull align 8 dereferenceable(24) %2, i8* nonnull align 16 dereferenceable(24) inttoptr (i64 140503769783536 to i8*), i64 24, i1 false)
   ret void
; └
}

Again, there are lots of subtleties here, honestly I’m still tuning my mental model.