Permute NamedTuple based on names of another NamedTuple

General Usage
,
1

I would like to perform a calculation in type space: given a NamedTuple{M,K} type, permute its fields so that the names match a NamedTuple{N}. Example:

julia> p1(NamedTuple{(:a,:b)}, NamedTuple{(:b,:a),Tuple{Int,Float64}})
Tuple{Float64, Int64}

An example solution using generated functions, simplifications welcome:

@generated function p1(::Type{<:NamedTuple{N}},
                       ::Type{<:NamedTuple{M,K}}) where {N,M,K}
    # NOTE: solutions can assume that N is a permutation of M,
    # it is checked before
    _K = fieldtypes(K)
    S = map(N) do n
        i = findfirst(m -> m ≡ n, M)
        @assert i ≢ nothing
        _K[i]
    end
    :(Tuple{$(S...)})
end
2 
p2(::Type{<:NamedTuple{N}}, ::Type{T}) where {N, T<:NamedTuple} =
    Base.promote_op(nt -> values(nt[N]), T)

:slight_smile:
Generally super convenient approach for transforming any value-based computation to a type-based computation. Comes at a slight cost of depending on the inference.

1 Like
3

Thanks, but just to clarify, I am looking for solutions without this shortcut.

4

Another option here would be to use Base.@assume_effects to write a non-generated function but still have it constant-fold based on the types:

Base.@assume_effects :foldable function p1(::Type{<:NamedTuple{N}},
                                           ::Type{<:NamedTuple{M,K}}) where {N,M,K}
    # NOTE: solutions can assume that N is a permutation of M,
    # it is checked before
    _K = fieldtypes(K)
    S = map(N) do n
        i = findfirst(m -> m ≡ n, M)
        @assert i ≢ nothing
        _K[i]
    end
    Tuple{S...}
end

Not sure if this is really any nicer though.

Another option would be Tuple recursion:

function p1(::Type{<:NamedTuple{N}},
            ::Type{<:NamedTuple{M,K}}) where {N,M,K}
    K′ = inner(N, M, fieldtypes(K))
    Tuple{K′...}
end

inner((head, tail...), M, K, K′=()) = let i = findfirst(==(head), M)
    @assert i ≢ nothing
    inner(tail, M, K, (K′..., K[i]))
end
inner(::Tuple{}, _, _, K′) = K′

The compiler appears to be good at optimizing this away (as expected for Tuple recursion):

julia> @code_typed p1(@NamedTuple{a::Int, b::Float64}, @NamedTuple{b::ComplexF64, a::String})
CodeInfo(
1 ─     return Tuple{String, ComplexF64}
) => Type{Tuple{String, ComplexF64}}
1 Like