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

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)

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;
 }
 

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);
 

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);

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,