implement support for current_scope (#605)

* implement support for `current_scope`

* add `current_scope` support for compiled frames

Author: aviatesk

bin/generate_builtins.jl

@@ -11,6 +11,7 @@ const RECENTLY_ADDED = Core.Builtin[
     Core.finalizer, Core._compute_sparams, Core._svec_ref,
     Core.compilerbarrier,
     Core.memoryref, Core.memoryref_isassigned, Core.memoryrefget, Core.memoryrefoffset, Core.memoryrefset!,
+    #=Core.current_scope=#
 ]
 # Builtins present in 1.6, not builtins (potentially still normal functions) anymore
 const RECENTLY_REMOVED = GlobalRef.(Ref(Core), [
@@ -187,13 +188,13 @@ function maybe_evaluate_builtin(frame, call_expr, expand::Bool)
             print(io,
 """
     $head f === $fstr
-            if !expand
-                argswrapped = getargs(args, frame)
-                return Some{Any}($fstr(argswrapped...))
-            end
-            # This uses the original arguments to avoid looking them up twice
-            # See #442
-            return Expr(:call, invoke, args[2:end]...)
+        if !expand
+            argswrapped = getargs(args, frame)
+            return Some{Any}($fstr(argswrapped...))
+        end
+        # This uses the original arguments to avoid looking them up twice
+        # See #442
+        return Expr(:call, invoke, args[2:end]...)
 """)
             continue
         elseif f === Core._call_latest
@@ -211,6 +212,22 @@ function maybe_evaluate_builtin(frame, call_expr, expand::Bool)
         end
         return maybe_recurse_expanded_builtin(frame, new_expr)
 """)
+            continue
+        elseif f === Core.current_scope
+            print(io,
+"""
+    elseif @static isdefined(Core, :current_scope) && f === Core.current_scope
+        if nargs == 0
+            if isempty(frame.framedata.current_scopes)
+                return Some{Any}(nothing)
+            else
+                return Some{Any}(frame.framedata.current_scopes[end])
+            end
+        else
+            return Some{Any}(Core.current_scope(getargs(args, frame)...))
+        end
+""")
+            continue
         end
 
         id = findfirst(isequal(f), Core.Compiler.T_FFUNC_KEY)

src/builtins.jl

@@ -94,8 +94,6 @@ function maybe_evaluate_builtin(frame, call_expr, expand::Bool)
             push!(new_expr.args, QuoteNode(x))
         end
         return maybe_recurse_expanded_builtin(frame, new_expr)
-    elseif f === Core._call_latest
-        return Some{Any}(Core._call_latest(getargs(args, frame)...))
     elseif @static isdefined(Core, :_compute_sparams) && f === Core._compute_sparams
         return Some{Any}(Core._compute_sparams(getargs(args, frame)...))
     elseif f === Core._equiv_typedef
@@ -126,6 +124,16 @@ function maybe_evaluate_builtin(frame, call_expr, expand::Bool)
         else
             return Some{Any}(Core.compilerbarrier(getargs(args, frame)...))
         end
+    elseif @static isdefined(Core, :current_scope) && f === Core.current_scope
+        if nargs == 0
+            if isempty(frame.framedata.current_scopes)
+                return Some{Any}(nothing)
+            else
+                return Some{Any}(frame.framedata.current_scopes[end])
+            end
+        else
+            return Some{Any}(Core.current_scope(getargs(args, frame)...))
+        end
     elseif @static isdefined(Core, :donotdelete) && f === Core.donotdelete
         return Some{Any}(Core.donotdelete(getargs(args, frame)...))
     elseif @static isdefined(Core, :finalizer) && f === Core.finalizer
@@ -227,13 +235,13 @@ function maybe_evaluate_builtin(frame, call_expr, expand::Bool)
             return Some{Any}(getglobal(getargs(args, frame)...))
         end
     elseif f === invoke
-            if !expand
-                argswrapped = getargs(args, frame)
-                return Some{Any}(invoke(argswrapped...))
-            end
-            # This uses the original arguments to avoid looking them up twice
-            # See #442
-            return Expr(:call, invoke, args[2:end]...)
+        if !expand
+            argswrapped = getargs(args, frame)
+            return Some{Any}(invoke(argswrapped...))
+        end
+        # This uses the original arguments to avoid looking them up twice
+        # See #442
+        return Expr(:call, invoke, args[2:end]...)
     elseif f === isa
         if nargs == 2
             return Some{Any}(isa(@lookup(frame, args[2]), @lookup(frame, args[3])))

src/construct.jl

@@ -269,7 +269,7 @@ function prepare_framedata(framecode, argvals::Vector{Any}, lenv::SimpleVector=e
     if length(junk_framedata) > 0
         olddata = pop!(junk_framedata)
         locals, ssavalues, sparams = olddata.locals, olddata.ssavalues, olddata.sparams
-        exception_frames, last_reference = olddata.exception_frames, olddata.last_reference
+        exception_frames, current_scopes, last_reference = olddata.exception_frames, olddata.current_scopes, olddata.last_reference
         last_exception = olddata.last_exception
         callargs = olddata.callargs
         resize!(locals, ns)
@@ -279,13 +279,15 @@ function prepare_framedata(framecode, argvals::Vector{Any}, lenv::SimpleVector=e
         # for check_isdefined to work properly, we need sparams to start out unassigned
         resize!(sparams, 0)
         empty!(exception_frames)
+        empty!(current_scopes)
         resize!(last_reference, ns)
         last_exception[] = _INACTIVE_EXCEPTION.instance
     else
         locals = Vector{Union{Nothing,Some{Any}}}(nothing, ns)
         ssavalues = Vector{Any}(undef, ng)
         sparams = Vector{Any}(undef, 0)
         exception_frames = Int[]
+        current_scopes = Any[]
         last_reference = Vector{Int}(undef, ns)
         callargs = Any[]
         last_exception = Ref{Any}(_INACTIVE_EXCEPTION.instance)
@@ -314,7 +316,8 @@ function prepare_framedata(framecode, argvals::Vector{Any}, lenv::SimpleVector=e
         isa(T, TypeVar) && continue  # only fill concrete types
         sparams[i] = T
     end
-    FrameData(locals, ssavalues, sparams, exception_frames, last_exception, caller_will_catch_err, last_reference, callargs)
+    return FrameData(locals, ssavalues, sparams, exception_frames, current_scopes,
+                     last_exception, caller_will_catch_err, last_reference, callargs)
 end
 
 """
@@ -331,6 +334,7 @@ end
 function prepare_frame_caller(caller::Frame, framecode::FrameCode, args::Vector{Any}, lenv::SimpleVector)
     caller_will_catch_err = !isempty(caller.framedata.exception_frames) || caller.framedata.caller_will_catch_err
     caller.callee = frame = prepare_frame(framecode, args, lenv, caller_will_catch_err)
+    copy!(frame.framedata.current_scopes, caller.framedata.current_scopes)
     frame.caller = caller
     return frame
 end

src/interpret.jl

@@ -222,6 +222,19 @@ function bypass_builtins(@nospecialize(recurse), frame, call_expr, pc)
     return nothing
 end
 
+function native_call(fargs::Vector{Any}, frame::Frame)
+    f = popfirst!(fargs) # now it's really just `args`
+    if !isempty(frame.framedata.current_scopes)
+        newscope = Core.current_scope()
+        for scope in frame.framedata.current_scopes
+            newscope = Scope(newscope, scope.values...)
+        end
+        ex = Expr(:tryfinally, :($f($fargs...)), nothing, newscope)
+        return Core.eval(moduleof(frame), ex)
+    end
+    return Base.invokelatest(f, fargs...)
+end
+
 function evaluate_call_compiled!(::Compiled, frame::Frame, call_expr::Expr; enter_generated::Bool=false)
     # @assert !enter_generated
     pc = frame.pc
@@ -230,9 +243,7 @@ function evaluate_call_compiled!(::Compiled, frame::Frame, call_expr::Expr; ente
     ret = maybe_evaluate_builtin(frame, call_expr, false)
     isa(ret, Some{Any}) && return ret.value
     fargs = collect_args(Compiled(), frame, call_expr)
-    f = fargs[1]
-    popfirst!(fargs)  # now it's really just `args`
-    return f(fargs...)
+    return native_call(fargs, frame)
 end
 
 function evaluate_call_recurse!(@nospecialize(recurse), frame::Frame, call_expr::Expr; enter_generated::Bool=false)
@@ -263,8 +274,7 @@ function evaluate_call_recurse!(@nospecialize(recurse), frame::Frame, call_expr:
         framecode, lenv = get_call_framecode(fargs, frame.framecode, frame.pc; enter_generated=enter_generated)
         if lenv === nothing
             if isa(framecode, Compiled)
-                f = popfirst!(fargs)  # now it's really just `args`
-                return Base.invokelatest(f, fargs...)
+                return native_call(fargs, frame)
             end
             return framecode  # this was a Builtin
         end
@@ -415,9 +425,6 @@ function eval_rhs(@nospecialize(recurse), frame, node::Expr)
     elseif head === :copyast
         val = (node.args[1]::QuoteNode).value
         return isa(val, Expr) ? copy(val) : val
-    elseif head === :enter
-        # XXX This seems to be dead code
-        return length(frame.framedata.exception_frames)
     elseif head === :boundscheck
         return true
     elseif head === :meta || head === :inbounds || head === :loopinfo ||
@@ -504,8 +511,12 @@ function step_expr!(@nospecialize(recurse), frame, @nospecialize(node), istoplev
                     for i = 1:length(node.args)
                         targ = node.args[i]
                         targ === nothing && continue
-                        frame.framecode.src.code[(targ::SSAValue).id] === nothing && continue
+                        enterstmt = frame.framecode.src.code[(targ::SSAValue).id]
+                        enterstmt === nothing && continue
                         pop!(data.exception_frames)
+                        if isdefined(enterstmt, :scope)
+                            pop!(data.current_scopes)
+                        end
                     end
                 end
             elseif node.head === :pop_exception
@@ -586,6 +597,9 @@ function step_expr!(@nospecialize(recurse), frame, @nospecialize(node), istoplev
         elseif @static (isdefined(Core.IR, :EnterNode) && true) && isa(node, Core.IR.EnterNode)
             rhs = node.catch_dest
             push!(data.exception_frames, rhs)
+            if isdefined(node, :scope)
+                push!(data.current_scopes, @lookup(frame, node.scope))
+            end
         else
             rhs = @lookup(frame, node)
         end

src/optimize.jl

@@ -241,7 +241,7 @@ function build_compiled_foreigncall!(stmt::Expr, code::CodeInfo, sparams::Vector
     return nothing
 end
 
-function replace_coretypes!(src; rev::Bool=false)
+function replace_coretypes!(@nospecialize(src); rev::Bool=false)
     if isa(src, CodeInfo)
         replace_coretypes_list!(src.code; rev=rev)
     elseif isa(src, Expr)
@@ -286,6 +286,13 @@ function replace_coretypes_list!(list::AbstractVector; rev::Bool=false)
             if rval !== val
                 list[i] = ReturnNode(rval)
             end
+        elseif @static (isdefined(Core.IR, :EnterNode) && true) && isa(stmt, Core.IR.EnterNode)
+            if isdefined(stmt, :scope)
+                rscope = rep(stmt.scope, rev)
+                if rscope !== stmt.scope
+                    list[i] = Core.IR.EnterNode(stmt.catch_dest, rscope)
+                end
+            end
         elseif isa(stmt, Expr)
             replace_coretypes!(stmt; rev=rev)
         end

src/packagedef.jl

@@ -43,6 +43,12 @@ if !isdefined(Base, Symbol("@something"))
     end
 end
 
+if isdefined(Base, :ScopedValues)
+    using Base: ScopedValues.Scope
+else
+    const Scope = Any
+end
+
 include("types.jl")
 include("utils.jl")
 include("construct.jl")

src/types.jl

@@ -200,6 +200,7 @@ struct FrameData
     ssavalues::Vector{Any}
     sparams::Vector{Any}
     exception_frames::Vector{Int}
+    current_scopes::Vector{Scope}
     last_exception::Base.RefValue{Any}
     caller_will_catch_err::Bool
     last_reference::Vector{Int}

test/interpret.jl

@@ -997,3 +997,31 @@ end
 # test for using generic functions that were previously builtin
 func_arrayref(a, i) = Core.arrayref(true, a, i)
 @test 2 == @interpret func_arrayref([1,2,3], 2)
+
+@static if isdefined(Base, :ScopedValues)
+const sval = ScopedValue(1)
+@test 1 == @interpret getindex(sval)
+
+# current_scope support for interpretation
+sval_func1() = @with sval => 2 begin
+    return sval[]
+end
+@test 2 == @interpret sval_func1()
+
+# current_scope support for compiled calls
+_sval_func2() = sval[]
+sval_func2() = @with sval => 2 begin
+    return _sval_func2()
+end
+let m = only(methods(_sval_func2))
+    push!(JuliaInterpreter.compiled_methods, m)
+    try
+        @test 2 == @interpret sval_func2()
+    finally
+        delete!(JuliaInterpreter.compiled_methods, m)
+    end
+end
+let frame = Frame(@__MODULE__, only(code_lowered(sval_func2)))
+    @test 2 == JuliaInterpreter.finish_and_return!(Compiled(), frame)
+end
+end