Add a guard that we start executing on the first PC

Methods with optional parameters don't always start executing at the
first PC, but we compile all methods assuming that they do.  This commit
adds a guard to ensure that we're actually starting at the first PC for
methods with optional params

5 files changed, 53 insertions, 3 deletions

diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
index 725eca5eb5..0ca293f59a 100644
--- a/bootstraptest/test_yjit.rb
+++ b/bootstraptest/test_yjit.rb
@@ -1,3 +1,14 @@
+# Make sure that optional param methods return the correct value
+assert_equal '1', %q{
+  def m(ary = [])
+    yield(ary)
+  end
+
+  # Warm the JIT with a 0 param call
+  2.times { m { } }
+  m(1) { |v| v }
+}
+
 # Test for topn
 assert_equal 'array', %q{
   def threequals(a)
diff --git a/yjit_codegen.c b/yjit_codegen.c
index 0ee795129a..7f12ab7b4f 100644
--- a/yjit_codegen.c
+++ b/yjit_codegen.c
@@ -318,12 +318,40 @@ yjit_side_exit(jitstate_t *jit, ctx_t *ctx)
     return yjit_gen_exit(jit, ctx, ocb);
 }
 
+// Generate a runtime guard that ensures the PC is at the start of the iseq,
+// otherwise take a side exit.  This is to handle the situation of optional
+// parameters.  When a function with optional parameters is called, the entry
+// PC for the method isn't necessarily 0, but we always generated code that
+// assumes the entry point is 0.
+static void
+yjit_pc_guard(const rb_iseq_t *iseq)
+{
+    RUBY_ASSERT(cb != NULL);
+
+    mov(cb, REG0, member_opnd(REG_CFP, rb_control_frame_t, pc));
+    mov(cb, REG1, const_ptr_opnd(iseq->body->iseq_encoded));
+    xor(cb, REG0, REG1);
+
+    // xor should impact ZF, so we can jz here
+    uint32_t pc_is_zero = cb_new_label(cb, "pc_is_zero");
+    jz_label(cb, pc_is_zero);
+
+    // We're not starting at the first PC, so we need to exit.
+    GEN_COUNTER_INC(cb, leave_start_pc_non_zero);
+    mov(cb, RAX, imm_opnd(Qundef));
+    ret(cb);
+
+    // PC should be at the beginning
+    cb_write_label(cb, pc_is_zero);
+    cb_link_labels(cb);
+}
+
 /*
 Compile an interpreter entry block to be inserted into an iseq
 Returns `NULL` if compilation fails.
 */
 uint8_t *
-yjit_entry_prologue(void)
+yjit_entry_prologue(const rb_iseq_t *iseq)
 {
     RUBY_ASSERT(cb != NULL);
 
@@ -352,6 +380,16 @@ yjit_entry_prologue(void)
     mov(cb, REG0, const_ptr_opnd(leave_exit_code));
     mov(cb, member_opnd(REG_CFP, rb_control_frame_t, jit_return), REG0);
 
+    // We're compiling iseqs that we *expect* to start at `insn_idx`. But in
+    // the case of optional parameters, the interpreter can set the pc to a
+    // different location depending on the optional parameters.  If an iseq
+    // has optional parameters, we'll add a runtime check that the PC we've
+    // compiled for is the same PC that the interpreter wants us to run with.
+    // If they don't match, then we'll take a side exit.
+    if (iseq->body->param.flags.has_opt) {
+        yjit_pc_guard(iseq);
+    }
+
     return code_ptr;
 }
 
diff --git a/yjit_codegen.h b/yjit_codegen.h
index 683765319e..7041a5a0a5 100644
--- a/yjit_codegen.h
+++ b/yjit_codegen.h
@@ -41,7 +41,7 @@ typedef enum codegen_status {
 // Code generation function signature
 typedef codegen_status_t (*codegen_fn)(jitstate_t* jit, ctx_t* ctx);
 
-uint8_t* yjit_entry_prologue();
+uint8_t* yjit_entry_prologue(const rb_iseq_t* iseq);
 
 void yjit_gen_block(block_t* block, rb_execution_context_t* ec);
 
diff --git a/yjit_core.c b/yjit_core.c
index 22527e24a6..988f034fc9 100644
--- a/yjit_core.c
+++ b/yjit_core.c
@@ -535,7 +535,7 @@ uint8_t* gen_entry_point(const rb_iseq_t *iseq, uint32_t insn_idx, rb_execution_
     blockid_t blockid = { iseq, insn_idx };
 
     // Write the interpreter entry prologue
-    uint8_t* code_ptr = yjit_entry_prologue();
+    uint8_t* code_ptr = yjit_entry_prologue(iseq);
 
     // Try to generate code for the entry block
     block_t* block = gen_block_version(blockid, &DEFAULT_CTX, ec);
diff --git a/yjit_iface.h b/yjit_iface.h
index 4f63cdb974..d2fad40194 100644
--- a/yjit_iface.h
+++ b/yjit_iface.h
@@ -48,6 +48,7 @@ YJIT_DECLARE_COUNTERS(
 
     leave_se_interrupt,
     leave_interp_return,
+    leave_start_pc_non_zero,
 
     getivar_se_self_not_heap,
     getivar_idx_out_of_range,