Make EnterNode save/restore dynamic scope (#52309)

As discussed in #51352, this gives `EnterNode` the ability to set (and
restore on leave or catch edge) jl_current_task->scope. Manual
modifications of the task field after the task has started are
considered undefined behavior. In addition, we gain a new intrinsic to
access current_task->scope and both inference and the optimizer will
forward scopes from EnterNodes to this intrinsic
(non-interprocedurally). Together with #51993 this is sufficient to
fully optimize ScopedValues (non-interprocedurally at least).

Author: Keno

base/boot.jl

@@ -460,6 +460,7 @@ eval(Core, quote
     ReturnNode() = $(Expr(:new, :ReturnNode)) # unassigned val indicates unreachable
     GotoIfNot(@nospecialize(cond), dest::Int) = $(Expr(:new, :GotoIfNot, :cond, :dest))
     EnterNode(dest::Int) = $(Expr(:new, :EnterNode, :dest))
+    EnterNode(dest::Int, @nospecialize(scope)) = $(Expr(:new, :EnterNode, :dest, :scope))
     LineNumberNode(l::Int) = $(Expr(:new, :LineNumberNode, :l, nothing))
     function LineNumberNode(l::Int, @nospecialize(f))
         isa(f, String) && (f = Symbol(f))
@@ -966,7 +967,8 @@ arraysize(a::Array, i::Int) = sle_int(i, nfields(a.size)) ? getfield(a.size, i)
 export arrayref, arrayset, arraysize, const_arrayref
 
 # For convenience
-EnterNode(old::EnterNode, new_dest::Int) = EnterNode(new_dest)
+EnterNode(old::EnterNode, new_dest::Int) = isdefined(old, :scope) ?
+    EnterNode(new_dest, old.scope) : EnterNode(new_dest)
 
 include(Core, "optimized_generics.jl")
 

base/compiler/abstractinterpretation.jl

@@ -3270,6 +3270,19 @@ function typeinf_local(interp::AbstractInterpreter, frame::InferenceState)
                 elseif isa(stmt, EnterNode)
                     ssavaluetypes[currpc] = Any
                     add_curr_ssaflag!(frame, IR_FLAG_NOTHROW)
+                    if isdefined(stmt, :scope)
+                        scopet = abstract_eval_value(interp, stmt.scope, currstate, frame)
+                        handler = frame.handlers[frame.handler_at[frame.currpc+1][1]]
+                        @assert handler.scopet !== nothing
+                        if !⊑(𝕃ᵢ, scopet, handler.scopet)
+                            handler.scopet = tmerge(𝕃ᵢ, scopet, handler.scopet)
+                            if isdefined(handler, :scope_uses)
+                                for bb in handler.scope_uses
+                                    push!(W, bb)
+                                end
+                            end
+                        end
+                    end
                     @goto fallthrough
                 elseif isexpr(stmt, :leave)
                     ssavaluetypes[currpc] = Any

base/compiler/inferencestate.jl

@@ -205,8 +205,10 @@ const CACHE_MODE_VOLATILE = 0x01 << 2 # not cached, optimization allowed
 
 mutable struct TryCatchFrame
     exct
+    scopet
     const enter_idx::Int
-    TryCatchFrame(@nospecialize(exct), enter_idx::Int) = new(exct, enter_idx)
+    scope_uses::Vector{Int}
+    TryCatchFrame(@nospecialize(exct), @nospecialize(scopet), enter_idx::Int) = new(exct, scopet, enter_idx)
 end
 
 mutable struct InferenceState
@@ -364,7 +366,7 @@ function compute_trycatch(code::Vector{Any}, ip::BitSet)
         stmt = code[pc]
         if isa(stmt, EnterNode)
             l = stmt.catch_dest
-            push!(handlers, TryCatchFrame(Bottom, pc))
+            push!(handlers, TryCatchFrame(Bottom, isdefined(stmt, :scope) ? Bottom : nothing, pc))
             handler_id = length(handlers)
             handler_at[pc + 1] = (handler_id, 0)
             push!(ip, pc + 1)

base/compiler/ssair/ir.jl

@@ -1414,6 +1414,7 @@ function process_node!(compact::IncrementalCompact, result_idx::Int, inst::Instr
             result_idx += 1
         end
     elseif cfg_transforms_enabled && isa(stmt, EnterNode)
+        stmt = renumber_ssa2!(stmt, ssa_rename, used_ssas, new_new_used_ssas, late_fixup, result_idx, do_rename_ssa, mark_refined!)::EnterNode
         label = bb_rename_succ[stmt.catch_dest]
         @assert label > 0
         ssa_rename[idx] = SSAValue(result_idx)

base/compiler/ssair/passes.jl

@@ -1069,6 +1069,29 @@ function fold_ifelse!(compact::IncrementalCompact, idx::Int, stmt::Expr)
     return false
 end
 
+function fold_current_scope!(compact::IncrementalCompact, idx::Int, stmt::Expr, lazydomtree::LazyDomtree)
+    domtree = get!(lazydomtree)
+
+    # The frontend enforces the invariant that any :enter dominates its active
+    # region, so all we have to do here is walk the domtree to find it.
+    dombb = block_for_inst(compact, SSAValue(idx))
+
+    local bbterminator
+    while true
+        dombb = domtree.idoms_bb[dombb]
+
+        # Did not find any dominating :enter - scope is inherited from the outside
+        dombb == 0 && return nothing
+
+        bbterminator = compact[SSAValue(last(compact.cfg_transform.result_bbs[dombb].stmts))][:stmt]
+        isa(bbterminator, EnterNode) || continue
+        isdefined(bbterminator, :scope) || continue
+        compact[idx] = bbterminator.scope
+        return nothing
+    end
+end
+
+
 # NOTE we use `IdSet{Int}` instead of `BitSet` for in these passes since they work on IR after inlining,
 # which can be very large sometimes, and program counters in question are often very sparse
 const SPCSet = IdSet{Int}
@@ -1201,6 +1224,8 @@ function sroa_pass!(ir::IRCode, inlining::Union{Nothing,InliningState}=nothing)
             elseif is_known_invoke_or_call(stmt, Core.OptimizedGenerics.KeyValue.get, compact)
                 2 == (length(stmt.args) - (isexpr(stmt, :invoke) ? 2 : 1)) || continue
                 lift_keyvalue_get!(compact, idx, stmt, 𝕃ₒ)
+            elseif is_known_call(stmt, Core.current_scope, compact)
+                fold_current_scope!(compact, idx, stmt, lazydomtree)
             elseif isexpr(stmt, :new)
                 refine_new_effects!(𝕃ₒ, compact, idx, stmt)
             end

base/compiler/ssair/show.jl

@@ -69,6 +69,10 @@ function print_stmt(io::IO, idx::Int, @nospecialize(stmt), used::BitSet, maxleng
     # given control flow information, we prefer to print these with the basic block #, instead of the ssa %
     elseif isa(stmt, EnterNode)
         print(io, "enter #", stmt.catch_dest, "")
+        if isdefined(stmt, :scope)
+            print(io, " with scope ")
+            show_unquoted(io, stmt.scope, indent)
+        end
     elseif stmt isa GotoNode
         print(io, "goto #", stmt.label)
     elseif stmt isa PhiNode

base/compiler/ssair/verify.jl

@@ -2,7 +2,7 @@
 
 function maybe_show_ir(ir::IRCode)
     if isdefined(Core, :Main)
-        Core.Main.Base.display(ir)
+        invokelatest(Core.Main.Base.display, ir)
     end
 end
 

base/compiler/tfuncs.jl

@@ -2488,6 +2488,19 @@ function builtin_effects(𝕃::AbstractLattice, @nospecialize(f::Builtin), argty
         return Effects(EFFECTS_TOTAL;
             consistent = (isa(setting, Const) && setting.val === :conditional) ? ALWAYS_TRUE : ALWAYS_FALSE,
             nothrow = compilerbarrier_nothrow(setting, nothing))
+    elseif f === Core.current_scope
+        nothrow = true
+        if length(argtypes) != 0
+            if length(argtypes) != 1 || !isvarargtype(argtypes[1])
+                return EFFECTS_THROWS
+            end
+            nothrow = false
+        end
+        return Effects(EFFECTS_TOTAL;
+            consistent = ALWAYS_FALSE,
+            notaskstate = false,
+            nothrow
+        )
     else
         if contains_is(_CONSISTENT_BUILTINS, f)
             consistent = ALWAYS_TRUE
@@ -2554,6 +2567,32 @@ function memoryop_noub(@nospecialize(f), argtypes::Vector{Any})
     return false
 end
 
+function current_scope_tfunc(interp::AbstractInterpreter, sv::InferenceState)
+    pc = sv.currpc
+    while true
+        handleridx = sv.handler_at[pc][1]
+        if handleridx == 0
+            # No local scope available - inherited from the outside
+            return Any
+        end
+        pchandler = sv.handlers[handleridx]
+        # Remember that we looked at this handler, so we get re-scheduled
+        # if the scope information changes
+        isdefined(pchandler, :scope_uses) || (pchandler.scope_uses = Int[])
+        pcbb = block_for_inst(sv.cfg, pc)
+        if findfirst(==(pcbb), pchandler.scope_uses) === nothing
+            push!(pchandler.scope_uses, pcbb)
+        end
+        scope = pchandler.scopet
+        if scope !== nothing
+            # Found the scope - forward it
+            return scope
+        end
+        pc = pchandler.enter_idx
+    end
+end
+current_scope_tfunc(interp::AbstractInterpreter, sv) = Any
+
 """
     builtin_nothrow(𝕃::AbstractLattice, f::Builtin, argtypes::Vector{Any}, rt) -> Bool
 
@@ -2568,9 +2607,6 @@ end
 function builtin_tfunction(interp::AbstractInterpreter, @nospecialize(f), argtypes::Vector{Any},
                            sv::Union{AbsIntState, Nothing})
     𝕃ᵢ = typeinf_lattice(interp)
-    if f === tuple
-        return tuple_tfunc(𝕃ᵢ, argtypes)
-    end
     if isa(f, IntrinsicFunction)
         if is_pure_intrinsic_infer(f) && all(@nospecialize(a) -> isa(a, Const), argtypes)
             argvals = anymap(@nospecialize(a) -> (a::Const).val, argtypes)
@@ -2596,6 +2632,16 @@ function builtin_tfunction(interp::AbstractInterpreter, @nospecialize(f), argtyp
         end
         tf = T_IFUNC[iidx]
     else
+        if f === tuple
+            return tuple_tfunc(𝕃ᵢ, argtypes)
+        elseif f === Core.current_scope
+            if length(argtypes) != 0
+                if length(argtypes) != 1 || !isvarargtype(argtypes[1])
+                    return Bottom
+                end
+            end
+            return current_scope_tfunc(interp, sv)
+        end
         fidx = find_tfunc(f)
         if fidx === nothing
             # unknown/unhandled builtin function

base/compiler/validation.jl

@@ -13,7 +13,7 @@ const VALID_EXPR_HEADS = IdDict{Symbol,UnitRange{Int}}(
     :new => 1:typemax(Int),
     :splatnew => 2:2,
     :the_exception => 0:0,
-    :enter => 1:1,
+    :enter => 1:2,
     :leave => 1:typemax(Int),
     :pop_exception => 1:1,
     :inbounds => 1:1,
@@ -160,6 +160,13 @@ function validate_code!(errors::Vector{InvalidCodeError}, c::CodeInfo, is_top_le
                 push!(errors, InvalidCodeError(INVALID_CALL_ARG, x.cond))
             end
             validate_val!(x.cond)
+        elseif isa(x, EnterNode)
+            if isdefined(x, :scope)
+                if !is_valid_argument(x.scope)
+                    push!(errors, InvalidCodeError(INVALID_CALL_ARG, x.scope))
+                end
+                validate_val!(x.scope)
+            end
         elseif isa(x, ReturnNode)
             if isdefined(x, :val)
                 if !is_valid_return(x.val)

base/scopedvalues.jl

@@ -81,13 +81,6 @@ function Scope(scope, pair1::Pair{<:ScopedValue}, pair2::Pair{<:ScopedValue}, pa
 end
 Scope(::Nothing) = nothing
 
-"""
-    current_scope()::Union{Nothing, Scope}
-
-Return the current dynamic scope.
-"""
-current_scope() = current_task().scope::Union{Nothing, Scope}
-
 function Base.show(io::IO, scope::Scope)
     print(io, Scope, "(")
     first = true
@@ -116,8 +109,7 @@ return `nothing`. Otherwise returns `Some{T}` with the current
 value.
 """
 function get(val::ScopedValue{T}) where {T}
-    # Inline current_scope to avoid doing the type assertion twice.
-    scope = current_task().scope
+    scope = Core.current_scope()::Union{Scope, Nothing}
     if scope === nothing
         isassigned(val) && return Some{T}(val.default)
         return nothing
@@ -151,25 +143,6 @@ function Base.show(io::IO, val::ScopedValue)
     print(io, ')')
 end
 
-"""
-    with(f, (var::ScopedValue{T} => val::T)...)
-
-Execute `f` in a new scope with `var` set to `val`.
-"""
-function with(f, pair::Pair{<:ScopedValue}, rest::Pair{<:ScopedValue}...)
-    @nospecialize
-    ct = Base.current_task()
-    current_scope = ct.scope::Union{Nothing, Scope}
-    ct.scope = Scope(current_scope, pair, rest...)
-    try
-        return f()
-    finally
-        ct.scope = current_scope
-    end
-end
-
-with(@nospecialize(f)) = f()
-
 """
     @with vars... expr
 
@@ -187,18 +160,18 @@ macro with(exprs...)
     else
         error("@with expects at least one argument")
     end
-    for expr in exprs
-        if expr.head !== :call || first(expr.args) !== :(=>)
-            error("@with expects arguments of the form `A => 2` got $expr")
-        end
-    end
     exprs = map(esc, exprs)
-    quote
-        ct = $(Base.current_task)()
-        current_scope = ct.scope::$(Union{Nothing, Scope})
-        ct.scope = $(Scope)(current_scope, $(exprs...))
-        $(Expr(:tryfinally, esc(ex), :(ct.scope = current_scope)))
-    end
+    Expr(:tryfinally, esc(ex), :(), :(Scope(Core.current_scope()::Union{Nothing, Scope}, $(exprs...))))
 end
 
+"""
+    with(f, (var::ScopedValue{T} => val::T)...)
+
+Execute `f` in a new scope with `var` set to `val`.
+"""
+function with(f, pair::Pair{<:ScopedValue}, rest::Pair{<:ScopedValue}...)
+    @with(pair, rest..., f())
+end
+with(@nospecialize(f)) = f()
+
 end # module ScopedValues