From 8286544dc5518f59bd607f6b6f8870fd22b3bac7 Mon Sep 17 00:00:00 2001
From: Jimmy Miller <jimmy.miller@shopify.com>
Date: Thu, 23 Mar 2023 13:11:46 -0400
Subject: YJIT: Use starting context for status === CantCompile (#7583)

---
 bootstraptest/test_yjit.rb | 42 ++++++++++++++++++++++++++++++++++++++++++
 yjit/src/codegen.rs        | 11 ++++++++---
 2 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
index 14ed13035d..f428407195 100644
--- a/bootstraptest/test_yjit.rb
+++ b/bootstraptest/test_yjit.rb
@@ -3737,3 +3737,45 @@ assert_normal_exit %q{
   )
   foo 1
 }
+
+# Regression test for CantCompile not using starting_ctx
+assert_normal_exit %q{
+  class Integer
+    def ===(other)
+      false
+    end
+  end
+
+  def my_func(x)
+    case x
+    when 1
+      1
+    when 2
+      2
+    else
+      3
+    end
+  end
+
+  my_func(1)
+}
+
+# Regression test for CantCompile not using starting_ctx
+assert_equal "ArgumentError", %q{
+  def literal(*args, &block)
+    s = ''.dup
+    args = [1, 2, 3]
+    literal_append(s, *args, &block)
+    s
+  end
+
+  def literal_append(sql, v)
+    [sql.inspect, v.inspect]
+  end
+
+  begin
+    literal("foo")
+  rescue ArgumentError
+    "ArgumentError"
+  end
+}
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 0301b2da8d..79267481fe 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -806,6 +806,10 @@ pub fn gen_single_block(
     // For each instruction to compile
     // NOTE: could rewrite this loop with a std::iter::Iterator
     while insn_idx < iseq_size {
+        // Set the starting_ctx so we can use it for side_exiting
+        // if status == CantCompile
+        let starting_ctx = ctx.clone();
+
         // Get the current pc and opcode
         let pc = unsafe { rb_iseq_pc_at_idx(iseq, insn_idx.into()) };
         // try_into() call below is unfortunate. Maybe pick i32 instead of usize for opcodes.
@@ -869,9 +873,10 @@ pub fn gen_single_block(
             }
 
             // TODO: if the codegen function makes changes to ctx and then return YJIT_CANT_COMPILE,
-            // the exit this generates would be wrong. We could save a copy of the entry context
-            // and assert that ctx is the same here.
-            gen_exit(jit.pc, &ctx, &mut asm);
+            // the exit this generates would be wrong. There are some changes that are safe to make
+            // like popping from the stack (but not writing). We are assuming here that only safe
+            // changes were made and using the starting_ctx.
+            gen_exit(jit.pc, &starting_ctx, &mut asm);
 
             // If this is the first instruction in the block, then
             // the entry address is the address for block_entry_exit
-- 
cgit v1.2.1