Base.@locals is slow and allocates

[MasonProtter]

Currently Base.@locals returns a Dict, but this is unfortunate because it's slow, allocating, and needs to box all of its contents. If Base.@locals carried no runtime cost to use, then it could be used to solve #34168.

It'd be really great if @locals (or if we don't want to change it a new macro @staticlocals) could be a NamedTuple which we should be able to construct in a type stable manner just from syntax alone.

To be concrete, if I write

function f(x)
    local y
    Base.@locals
end

this lowers to

 CodeInfo(
1 ─      Core.NewvarNode(:(y))
│   %2 = Core.apply_type(Base.Dict, Core.Symbol, Core.Any)
│   %3 = (%2)()
│   %4 = $(Expr(:isdefined, :(y)))
└──      goto #3 if not %4
2 ─      Base.setindex!(%3, y, :y)
3 ┄ %7 = $(Expr(:isdefined, :(x)))
└──      goto #5 if not %7
4 ─      Base.setindex!(%3, x, :x)
5 ┄      return %3
)

But I think that we could emit code that looks more like this:

function Base.setindex(nt::NamedTuple{names}, x, ::Val{n}) where {names, n}
    NamedTuple{(names..., n)}((values(nt)..., x))
end

function f(x)
    local y
    
    locals = NamedTuple()
    if @isdefined x
        locals = Base.setindex(locals, x, Val{:x}())
    end
    if @isdefined y
        locals = Base.setindex(locals, y, Val{:y}())
    end
    locals
end

which the optimizer has no problem with:

julia> @code_typed f(1)
CodeInfo(
1 ─      goto #3 if not true
2 ─ %2 = %new(NamedTuple{(:x,), Tuple{Int64}}, x)::NamedTuple{(:x,), Tuple{Int64}}
3 ┄ %3 = φ (#2 => %2)::NamedTuple{(:x,), Tuple{Int64}}
└──      return %3
) => NamedTuple{(:x,), Tuple{Int64}}

---

In code with very very many local variables, this appears to slow things down, but does still get resolved in a type stable manner.
E.g.

julia> function  Base.setindex(nt::NamedTuple{names}, x, ::Val{n}) where {names, n}
           NamedTuple{(names..., n)}((values(nt)..., x))
       end

julia> let N = 12, M = 21
           args = [Symbol(:arg, n) for n ∈ 1:N]
           vars = [Symbol(:var, m) for m ∈ 1:M]
       
           defargs = map(enumerate(vars)) do (m, var)
               if isodd(m)
                   :(local $var)
               else
                   :($var = $(rand((1, "hi", QuoteNode(:bye), 2.0, 3 => 4 + im))))
               end
           end
           defs = Expr(:block, defargs...)
       
           setindicesargs = map([args; vars]) do x
               :(if @isdefined $x
                     nt = Base.setindex(nt, $x, Val{$(QuoteNode(x))}())
                 end)
           end
           thelocals = Expr(:block, :(nt = NamedTuple()), setindicesargs..., :nt)
           @eval @time begin
           #quote
               function f($(args...),)
                   $defs
                   $thelocals
               end
               Core.Compiler.return_type(f, Tuple{$(rand((Int, String, Symbol, ComplexF64), N)...),})
           end
       end 
  0.056692 seconds (487.75 k allocations: 23.483 MiB, 92.06% compilation time)
NamedTuple{(:arg1, :arg2, :arg3, :arg4, :arg5, :arg6, :arg7, :arg8, :arg9, :arg10, :arg11, :arg12, :var2, :var4, :var6, :var8, :var10, :var12, :var14, :var16, :var18, :var20), Tuple{Int64, String, Int64, Int64, ComplexF64, ComplexF64, String, String, Symbol, Int64, ComplexF64, String, Float64, Float64, Symbol, Pair{Int64, Complex{Int64}}, Symbol, Pair{Int64, Complex{Int64}}, Symbol, Symbol, String, String}}

[KristofferC]

FWIW, named tuple was discussed briefly in the original PR: #29733 (comment)