From 9ae93806d9facca8871ed9eb68b48702e5be7c7e Mon Sep 17 00:00:00 2001
From: Cody Tapscott <topolarity@tapscott.me>
Date: Tue, 19 Nov 2024 13:44:38 -0500
Subject: [PATCH] codegen: manually restore `ct->scope` when no exception
 handler is emitted

This fixes a bug introduced by #55907, which was neglecting that it's possible
for `Expr(:enter, ...)` to have no `catch` destination and still have a scope,
especially if the compiler has decided that the body is `nothrow` and chooses
to optimize away the `catch` destination.
---
 src/codegen.cpp      | 25 ++++++++++++++++++++-----
 test/scopedvalues.jl |  7 +++++++
 2 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/src/codegen.cpp b/src/codegen.cpp
index 761c3e7eb87583..d10438a179c7e6 100644
--- a/src/codegen.cpp
+++ b/src/codegen.cpp
@@ -1955,7 +1955,7 @@ class jl_codectx_t {
     // local var info. globals are not in here.
     SmallVector<jl_varinfo_t, 0> slots;
     std::map<int, jl_varinfo_t> phic_slots;
-    std::map<int, Value*> scope_tokens;
+    std::map<int, std::pair<Value*, Value*> > scope_restore;
     SmallVector<jl_cgval_t, 0> SAvalues;
     SmallVector<std::tuple<jl_cgval_t, BasicBlock *, AllocaInst *, PHINode *, SmallVector<PHINode*,0>, jl_value_t *>, 0> PhiNodes;
     SmallVector<bool, 0> ssavalue_assigned;
@@ -6573,6 +6573,7 @@ static void emit_stmtpos(jl_codectx_t &ctx, jl_value_t *expr, int ssaval_result)
     }
     else if (head == jl_leave_sym) {
         int hand_n_leave = 0;
+        Value *scope_to_restore = nullptr;
         for (size_t i = 0; i < jl_expr_nargs(ex); ++i) {
             jl_value_t *arg = args[i];
             if (arg == jl_nothing)
@@ -6582,7 +6583,7 @@ static void emit_stmtpos(jl_codectx_t &ctx, jl_value_t *expr, int ssaval_result)
             jl_value_t *enter_stmt = jl_array_ptr_ref(ctx.code, enter_idx);
             if (enter_stmt == jl_nothing)
                 continue;
-            if (ctx.scope_tokens.count(enter_idx)) {
+            if (ctx.scope_restore.count(enter_idx)) {
                 // TODO: The semantics of `gc_preserve` are not perfect here. An `Expr(:enter, ...)` block may
                 //       have multiple exits, but effects of `preserve_end` are only extended to the end of the
                 //       dominance of each `Expr(:leave, ...)`.
@@ -6594,15 +6595,25 @@ static void emit_stmtpos(jl_codectx_t &ctx, jl_value_t *expr, int ssaval_result)
                 //
                 //       This is correct as-is anyway - it just means the scope lives longer than it needs to
                 //       if the `Expr(:enter, ...)` has multiple exits.
-                ctx.builder.CreateCall(prepare_call(gc_preserve_end_func), {ctx.scope_tokens[enter_idx]});
+                ctx.builder.CreateCall(prepare_call(gc_preserve_end_func), {ctx.scope_restore[enter_idx].first});
             }
             if (jl_enternode_catch_dest(enter_stmt)) {
                 // We're not actually setting up the exception frames for these, so
                 // we don't need to exit them.
                 hand_n_leave += 1;
             }
+            if (!jl_enternode_catch_dest(enter_stmt ) && ctx.scope_restore.count(enter_idx)) {
+                // We are leaving an `:enter` context that we didn't create a handler for, so
+                // the scope needs to be manually restored.
+                scope_to_restore = ctx.scope_restore[enter_idx].second;
+            }
         }
         ctx.builder.CreateCall(prepare_call(jlleave_noexcept_func), {get_current_task(ctx), ConstantInt::get(getInt32Ty(ctx.builder.getContext()), hand_n_leave)});
+        if (scope_to_restore) {
+            Value *scope_ptr = get_scope_field(ctx);
+            jl_aliasinfo_t::fromTBAA(ctx, ctx.tbaa().tbaa_gcframe).decorateInst(
+                ctx.builder.CreateAlignedStore(scope_to_restore, scope_ptr, ctx.types().alignof_ptr));
+        }
     }
     else if (head == jl_pop_exception_sym) {
         jl_cgval_t excstack_state = emit_expr(ctx, jl_exprarg(expr, 0));
@@ -9652,12 +9663,16 @@ static jl_llvm_functions_t
                     continue;
                 }
                 Value *scope_boxed = boxed(ctx, scope);
-                StoreInst *scope_store = ctx.builder.CreateAlignedStore(scope_boxed, get_scope_field(ctx), ctx.types().alignof_ptr);
+                Value *scope_ptr = get_scope_field(ctx);
+                LoadInst *current_scope = ctx.builder.CreateAlignedLoad(ctx.types().T_prjlvalue, scope_ptr, ctx.types().alignof_ptr);
+                StoreInst *scope_store = ctx.builder.CreateAlignedStore(scope_boxed, scope_ptr, ctx.types().alignof_ptr);
+                jl_aliasinfo_t::fromTBAA(ctx, ctx.tbaa().tbaa_gcframe).decorateInst(current_scope);
                 jl_aliasinfo_t::fromTBAA(ctx, ctx.tbaa().tbaa_gcframe).decorateInst(scope_store);
                 // GC preserve the scope, since it is not rooted in the `jl_handler_t *`
                 // and may be removed from jl_current_task by any nested block and then
                 // replaced later
-                ctx.scope_tokens[cursor] = ctx.builder.CreateCall(prepare_call(gc_preserve_begin_func), {scope_boxed});
+                Value *scope_token = ctx.builder.CreateCall(prepare_call(gc_preserve_begin_func), {scope_boxed});
+                ctx.scope_restore[cursor] = std::make_pair(scope_token, current_scope);
             }
         }
         else {
diff --git a/test/scopedvalues.jl b/test/scopedvalues.jl
index 2c38a0642ce247..174bc690ac0a2e 100644
--- a/test/scopedvalues.jl
+++ b/test/scopedvalues.jl
@@ -175,3 +175,10 @@ const inlineable_const_sv = ScopedValue(1)
 @test fully_eliminated(; retval=(inlineable_const_sv => 1)) do
     inlineable_const_sv => 1
 end
+
+# Handle nothrow scope bodies correctly (#56609)
+@eval function nothrow_scope()
+    $(Expr(:tryfinally, :(), nothing, 1))
+    @test Core.current_scope() === nothing
+end
+nothrow_scope()