Type inference when generating a NamedTuple

question

#1

This is an MWE simplified from a real-world problem — I have a T <: NamedTuple type, and would like to generate an value ::T from functions that generate fields based on type, to be named appropriately.

using Random

struct Gen{T <: NamedTuple, R <: AbstractRNG}
    rng::R # the original problem has a mutable field capturing a state, too
end

# convenience constructor
Gen{T}(rng::R) where {T <: NamedTuple, R <: AbstractRNG} = Gen{T,R}(rng)

# field generator: in the MWE, random
genfield(rng::AbstractRNG, ::Type{T}) where T <: Real = rand(rng, T)

# field generator for potentially missing
genfield(rng::AbstractRNG, ::Type{Union{Missing, T}}) where T =
    rand(rng) < 0.5 ? missing : genfield(rng, T)

# THIS is the problematic function
gen(g::Gen{T}) where T =
    T(ntuple(i -> genfield(g.rng, fieldtype(T, i)), fieldcount(T)))

g = Gen{NamedTuple{(:a, :b), Tuple{Int, Union{Missing, Float64}}}}(Random.GLOBAL_RNG)

@code_warntype genfield(Random.GLOBAL_RNG, Union{Missing, Float64})

@code_warntype gen(g)

With lot of the output cut,

julia> @code_warntype genfield(Random.GLOBAL_RNG, Union{Missing, Float64})
Body::Union{Missing, Float64}

julia> @code_warntype gen(g)
Body::Any

julia> VERSION
v"1.1.0-DEV.439"

#2

How about

gen(g::Gen{T} where T)  ::Union{Missing, T} =
    T(ntuple(i -> genfield(g.rng, fieldtype(T, i)), fieldcount(T)))

#3

Went with @generated, as in (the further simplified MWE)

struct Gen{T <: NamedTuple} end
genfield(::Type{T}) where {T <: Real} = zero(T)
genfield(::Type{Union{Missing, T}}) where T = rand() < 0.5 ? missing : genfield(T)

function gen(g::Gen{T}) where T
    if @generated
        _names = fieldnames(T)
        _types = map(n -> fieldtype(T, n), fieldnames(T))
        _vals = Any[:(genfield($fT)) for fT in _types]
        :( NamedTuple{$_names, Tuple{$(_types...)}}(($(_vals...),)) )
    else
        T(ntuple(i -> genfield(fieldtype(T, i)), fieldcount(T)))
    end
end

then

julia> g = Gen{NamedTuple{(:a, :b), Tuple{Int, Float64}}}()
Gen{NamedTuple{(:a, :b),Tuple{Int64,Float64}}}()

julia> @code_warntype gen(g)
Body::NamedTuple{(:a, :b),Tuple{Int64,Float64}}
 2 1 ─     return (a = 0, b = 0.0)                                                            │╻ macro expansion

#4

But for

g = Gen{NamedTuple{(:a, :b), Tuple{Int, Union{Missing, Float64}}}}()

I still get

julia> @code_warntype gen(g)
Body::Any
 2 1 ── %1  = Random.GLOBAL_RNG::Random.MersenneTwister                   │╻╷╷          macro expansion
   │    %2  = (Base.getfield)(%1, :idxF)::Int64                           ││┃││╷╷╷╷╷╷    genfield
   │    %3  = Random.MT_CACHE_F::Int64                                    │││┃││││││      rand
   │    %4  = (%2 === %3)::Bool                                           ││││┃││││││      rand
   └───       goto #3 if not %4                                           │││││┃│││         rand
   2 ── %6  = $(Expr(:gc_preserve_begin, :(%1)))                          ││││││┃│││╷        rand
   │    %7  = (Base.getfield)(%1, :state)::Random.DSFMT.DSFMT_state       │││││││╻            rand
   │    %8  = (Base.getfield)(%1, :vals)::Array{Float64,1}                ││││││││┃│││         reserve_1
   │    %9  = $(Expr(:foreigncall, :(:jl_array_ptr), Ptr{Float64}, svec(Any), :(:ccall), 1, :(%8)))::Ptr{Float64}
   │    %10 = (Base.getfield)(%1, :vals)::Array{Float64,1}                ││││││││││╻            macro expansion
   │    %11 = (Base.arraylen)(%10)::Int64                                 │││││││││││╻            length
   │          invoke Random.dsfmt_fill_array_close1_open2!(%7::Random.DSFMT.DSFMT_state, %9::Ptr{Float64}, %11::Int64)
   │          $(Expr(:gc_preserve_end, :(%6)))                            │││││││││││  
   └───       (Base.setfield!)(%1, :idxF, 0)                              │││││││││││╻╷           mt_setfull!
   3 ──       goto #4                                                     ││││││││╻            reserve_1
   4 ── %16 = (Base.getfield)(%1, :vals)::Array{Float64,1}                ││││││││╻╷╷          rand_inbounds
   │    %17 = (Base.getfield)(%1, :idxF)::Int64                           │││││││││┃│           mt_pop!
   │    %18 = (Base.add_int)(%17, 1)::Int64                               ││││││││││╻            +
   │          (Base.setfield!)(%1, :idxF, %18)                            ││││││││││╻            setproperty!
   │    %20 = (Base.arrayref)(false, %16, %18)::Float64                   ││││││││││╻            getindex
   └───       goto #5                                                     ││││││││     
   5 ──       goto #6                                                     │││││││      
   6 ── %23 = (Base.sub_float)(%20, 1.0)::Float64                         ││││││╻            -
   └───       goto #7                                                     ││││││       
   7 ──       goto #8                                                     │││││        
   8 ──       goto #9                                                     ││││         
   9 ── %27 = (Base.lt_float)(%23, 0.5)::Bool                             │││╻            <
   └───       goto #11 if not %27                                         │││          
   10 ─ %29 = Main.missing::Core.Compiler.Const(missing, false)           │││          
   └───       goto #12                                                    │││          
   11 ─       goto #12                                                    │││          
   12 ┄ %32 = φ (#10 => %29, #11 => 0.0)::Union{Missing, Float64}         ││           
   │    %33 = (Core.tuple)(0, %32)::Core.Compiler.PartialTuple(Tuple{Int64,Union{Missing, Float64}}, Any[Const(0, false), Union{Missing, Float64}])
   │    %34 = (NamedTuple{(:a, :b),Tuple{Int64,Union{Missing, Float64}}})(%33)::Any    
   └───       return %34                                                  ││           

so it did not really help much.