Skip to content

Base.@locals is slow and allocates #46358

@MasonProtter

Description

@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)
5return %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}}

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions