YJIT: Always abandon the block when gen_branch() or defer_compilation() fails

In [1], we started checking for gen_branch failures, but I made two
crucial mistakes. One, defer_compilation() had the same issue as
gen_branch() but wasn't checked. Two, returning None from a codegen
function does not throw away the block. Checking how gen_single_block()
handles codegen functions, you can see that None terminates the block
with an exit, but does not overall return an Err. This handling is fine
for unimplemented instructions, for example, but incorrect in case
gen_branch() fails. The missing branch essentially corrupts the
block; adding more code after a missing branch doesn't correct the code.

Always abandon the block when defer_compilation() or gen_branch() fails.

[1]: cb661d7d82984cdb54485ea3f4af01ac21960882
Fixup: [1]

yjit/src/codegen.rs

@@ -195,6 +195,34 @@ impl<'a> JITState<'a> {
         self.outlined_code_block
     }
 
+    /// Leave a code stub to re-enter the compiler at runtime when the compiling program point is
+    /// reached. Should always be used in tail position like `return jit.defer_compilation(asm);`.
+    #[must_use]
+    fn defer_compilation(&mut self, asm: &mut Assembler) -> Option<CodegenStatus> {
+        if crate::core::defer_compilation(self, asm).is_err() {
+            // If we can't leave a stub, the block isn't usable and we have to bail.
+            self.block_abandoned = true;
+        }
+        Some(EndBlock)
+    }
+
+    /// Generate a branch with either end possibly stubbed out
+    fn gen_branch(
+        &mut self,
+        asm: &mut Assembler,
+        target0: BlockId,
+        ctx0: &Context,
+        target1: Option<BlockId>,
+        ctx1: Option<&Context>,
+        gen_fn: BranchGenFn,
+    ) {
+        if crate::core::gen_branch(self, asm, target0, ctx0, target1, ctx1, gen_fn).is_none() {
+            // If we can't meet the request for a branch, the code is
+            // essentially corrupt and we have to discard the block.
+            self.block_abandoned = true;
+        }
+    }
+
     /// Return true if the current ISEQ could escape an environment.
     ///
     /// As of vm_push_frame(), EP is always equal to BP. However, after pushing
@@ -1538,8 +1566,7 @@ fn fuse_putobject_opt_ltlt(
             return None;
         }
         if !jit.at_compile_target() {
-            defer_compilation(jit, asm);
+            return jit.defer_compilation(asm);
-            return Some(EndBlock);
         }
 
         let lhs = jit.peek_at_stack(&asm.ctx, 0);
@@ -1661,8 +1688,7 @@ fn gen_opt_plus(
     let two_fixnums = match asm.ctx.two_fixnums_on_stack(jit) {
         Some(two_fixnums) => two_fixnums,
         None => {
-            defer_compilation(jit, asm);
+            return jit.defer_compilation(asm);
-            return Some(EndBlock);
         }
     };
 
@@ -1802,8 +1828,7 @@ fn gen_splatkw(
 ) -> Option<CodegenStatus> {
     // Defer compilation so we can specialize on a runtime hash operand
     if !jit.at_compile_target() {
-        defer_compilation(jit, asm);
+        return jit.defer_compilation(asm);
-        return Some(EndBlock);
     }
 
     let comptime_hash = jit.peek_at_stack(&asm.ctx, 1);
@@ -2176,8 +2201,7 @@ fn gen_expandarray(
 
     // Defer compilation so we can specialize on a runtime `self`
     if !jit.at_compile_target() {
-        defer_compilation(jit, asm);
+        return jit.defer_compilation(asm);
-        return Some(EndBlock);
     }
 
     let comptime_recv = jit.peek_at_stack(&asm.ctx, 0);
@@ -2718,10 +2742,7 @@ fn jit_chain_guard(
             idx: jit.insn_idx,
         };
 
-        // Bail if we can't generate the branch
+        jit.gen_branch(asm, bid, &deeper, None, None, target0_gen_fn);
-        if gen_branch(jit, asm, bid, &deeper, None, None, target0_gen_fn).is_none() {
-            jit.block_abandoned = true;
-        }
     } else {
         target0_gen_fn.call(asm, Target::side_exit(counter), None);
     }
@@ -2895,8 +2916,7 @@ fn gen_getinstancevariable(
 ) -> Option<CodegenStatus> {
     // Defer compilation so we can specialize on a runtime `self`
     if !jit.at_compile_target() {
-        defer_compilation(jit, asm);
+        return jit.defer_compilation(asm);
-        return Some(EndBlock);
     }
 
     let ivar_name = jit.get_arg(0).as_u64();
@@ -2959,8 +2979,7 @@ fn gen_setinstancevariable(
 ) -> Option<CodegenStatus> {
     // Defer compilation so we can specialize on a runtime `self`
     if !jit.at_compile_target() {
-        defer_compilation(jit, asm);
+        return jit.defer_compilation(asm);
-        return Some(EndBlock);
     }
 
     let ivar_name = jit.get_arg(0).as_u64();
@@ -3270,8 +3289,7 @@ fn gen_definedivar(
 ) -> Option<CodegenStatus> {
     // Defer compilation so we can specialize base on a runtime receiver
     if !jit.at_compile_target() {
-        defer_compilation(jit, asm);
+        return jit.defer_compilation(asm);
-        return Some(EndBlock);
     }
 
     let ivar_name = jit.get_arg(0).as_u64();
@@ -3500,8 +3518,7 @@ fn gen_fixnum_cmp(
         Some(two_fixnums) => two_fixnums,
         None => {
             // Defer compilation so we can specialize based on a runtime receiver
-            defer_compilation(jit, asm);
+            return jit.defer_compilation(asm);
-            return Some(EndBlock);
         }
     };
 
@@ -3680,8 +3697,7 @@ fn gen_opt_eq(
         Some(specialized) => specialized,
         None => {
             // Defer compilation so we can specialize base on a runtime receiver
-            defer_compilation(jit, asm);
+            return jit.defer_compilation(asm);
-            return Some(EndBlock);
         }
     };
 
@@ -3718,8 +3734,7 @@ fn gen_opt_aref(
 
     // Defer compilation so we can specialize base on a runtime receiver
     if !jit.at_compile_target() {
-        defer_compilation(jit, asm);
+        return jit.defer_compilation(asm);
-        return Some(EndBlock);
     }
 
     // Specialize base on compile time values
@@ -3819,8 +3834,7 @@ fn gen_opt_aset(
 ) -> Option<CodegenStatus> {
     // Defer compilation so we can specialize on a runtime `self`
     if !jit.at_compile_target() {
-        defer_compilation(jit, asm);
+        return jit.defer_compilation(asm);
-        return Some(EndBlock);
     }
 
     let comptime_recv = jit.peek_at_stack(&asm.ctx, 2);
@@ -3951,8 +3965,7 @@ fn gen_opt_and(
         Some(two_fixnums) => two_fixnums,
         None => {
             // Defer compilation so we can specialize on a runtime `self`
-            defer_compilation(jit, asm);
+            return jit.defer_compilation(asm);
-            return Some(EndBlock);
         }
     };
 
@@ -3990,8 +4003,7 @@ fn gen_opt_or(
         Some(two_fixnums) => two_fixnums,
         None => {
             // Defer compilation so we can specialize on a runtime `self`
-            defer_compilation(jit, asm);
+            return jit.defer_compilation(asm);
-            return Some(EndBlock);
         }
     };
 
@@ -4029,8 +4041,7 @@ fn gen_opt_minus(
         Some(two_fixnums) => two_fixnums,
         None => {
             // Defer compilation so we can specialize on a runtime `self`
-            defer_compilation(jit, asm);
+            return jit.defer_compilation(asm);
-            return Some(EndBlock);
         }
     };
 
@@ -4069,8 +4080,7 @@ fn gen_opt_mult(
     let two_fixnums = match asm.ctx.two_fixnums_on_stack(jit) {
         Some(two_fixnums) => two_fixnums,
         None => {
-            defer_compilation(jit, asm);
+            return jit.defer_compilation(asm);
-            return Some(EndBlock);
         }
     };
 
@@ -4121,8 +4131,7 @@ fn gen_opt_mod(
         Some(two_fixnums) => two_fixnums,
         None => {
             // Defer compilation so we can specialize on a runtime `self`
-            defer_compilation(jit, asm);
+            return jit.defer_compilation(asm);
-            return Some(EndBlock);
         }
     };
 
@@ -4459,8 +4468,7 @@ fn gen_opt_case_dispatch(
     // hash lookup, at least for small hashes, but it's worth revisiting this
     // assumption in the future.
     if !jit.at_compile_target() {
-        defer_compilation(jit, asm);
+        return jit.defer_compilation(asm);
-        return Some(EndBlock);
     }
 
     let case_hash = jit.get_arg(0);
@@ -4572,15 +4580,14 @@ fn gen_branchif(
 
         // Generate the branch instructions
         let ctx = asm.ctx;
-        gen_branch(
+        jit.gen_branch(
-            jit,
             asm,
             jump_block,
             &ctx,
             Some(next_block),
             Some(&ctx),
             BranchGenFn::BranchIf(Cell::new(BranchShape::Default)),
-        )?;
+        );
     }
 
     Some(EndBlock)
@@ -4626,15 +4633,14 @@ fn gen_branchunless(
 
         // Generate the branch instructions
         let ctx = asm.ctx;
-        gen_branch(
+        jit.gen_branch(
-            jit,
             asm,
             jump_block,
             &ctx,
             Some(next_block),
             Some(&ctx),
             BranchGenFn::BranchUnless(Cell::new(BranchShape::Default)),
-        )?;
+        );
     }
 
     Some(EndBlock)
@@ -4677,15 +4683,14 @@ fn gen_branchnil(
         asm.cmp(val_opnd, Opnd::UImm(Qnil.into()));
         // Generate the branch instructions
         let ctx = asm.ctx;
-        gen_branch(
+        jit.gen_branch(
-            jit,
             asm,
             jump_block,
             &ctx,
             Some(next_block),
             Some(&ctx),
             BranchGenFn::BranchNil(Cell::new(BranchShape::Default)),
-        )?;
+        );
     }
 
     Some(EndBlock)
@@ -8004,19 +8009,14 @@ fn gen_send_iseq(
     return_asm.ctx.set_as_return_landing();
 
     // Write the JIT return address on the callee frame
-    if gen_branch(
+    jit.gen_branch(
-        jit,
         asm,
         return_block,
         &return_asm.ctx,
         None,
         None,
         BranchGenFn::JITReturn,
-    ).is_none() {
+    );
-        // Returning None here would have send_dynamic() code following incomplete
-        // send code. Abandon the block instead.
-        jit.block_abandoned = true;
-    }
 
     // ec->cfp is updated after cfp->jit_return for rb_profile_frames() safety
     asm_comment!(asm, "switch to new CFP");
@@ -8711,8 +8711,7 @@ fn gen_send_general(
 
     // Defer compilation so we can specialize on class of receiver
     if !jit.at_compile_target() {
-        defer_compilation(jit, asm);
+        return jit.defer_compilation(asm);
-        return Some(EndBlock);
     }
 
     let ci_flags = unsafe { vm_ci_flag(ci) };
@@ -9275,8 +9274,7 @@ fn gen_invokeblock_specialized(
     cd: *const rb_call_data,
 ) -> Option<CodegenStatus> {
     if !jit.at_compile_target() {
-        defer_compilation(jit, asm);
+        return jit.defer_compilation(asm);
-        return Some(EndBlock);
     }
 
     // Fallback to dynamic dispatch if this callsite is megamorphic
@@ -9438,8 +9436,7 @@ fn gen_invokesuper_specialized(
 ) -> Option<CodegenStatus> {
     // Defer compilation so we can specialize on class of receiver
     if !jit.at_compile_target() {
-        defer_compilation(jit, asm);
+        return jit.defer_compilation(asm);
-        return Some(EndBlock);
     }
 
     // Handle the last two branches of vm_caller_setup_arg_block
@@ -9672,8 +9669,7 @@ fn gen_objtostring(
     asm: &mut Assembler,
 ) -> Option<CodegenStatus> {
     if !jit.at_compile_target() {
-        defer_compilation(jit, asm);
+        return jit.defer_compilation(asm);
-        return Some(EndBlock);
     }
 
     let recv = asm.stack_opnd(0);
@@ -10014,8 +10010,7 @@ fn gen_getblockparamproxy(
     asm: &mut Assembler,
 ) -> Option<CodegenStatus> {
     if !jit.at_compile_target() {
-        defer_compilation(jit, asm);
+        return jit.defer_compilation(asm);
-        return Some(EndBlock);
     }
 
     // EP level

yjit/src/core.rs

@@ -3916,10 +3916,7 @@ pub fn gen_direct_jump(jit: &mut JITState, ctx: &Context, target0: BlockId, asm:
 }
 
 /// Create a stub to force the code up to this point to be executed
-pub fn defer_compilation(
+pub fn defer_compilation(jit: &mut JITState, asm: &mut Assembler) -> Result<(), ()> {
-    jit: &mut JITState,
-    asm: &mut Assembler,
-) {
     if asm.ctx.is_deferred() {
         panic!("Double defer!");
     }
@@ -3936,7 +3933,7 @@ pub fn defer_compilation(
     };
 
     // Likely a stub since the context is marked as deferred().
-    let target0_address = branch.set_target(0, blockid, &next_ctx, jit);
+    let dst_addr = branch.set_target(0, blockid, &next_ctx, jit).ok_or(())?;
 
     // Pad the block if it has the potential to be invalidated. This must be
     // done before gen_fn() in case the jump is overwritten by a fallthrough.
@@ -3947,9 +3944,7 @@ pub fn defer_compilation(
     // Call the branch generation function
     asm_comment!(asm, "defer_compilation");
     asm.mark_branch_start(&branch);
-    if let Some(dst_addr) = target0_address {
     branch.gen_fn.call(asm, Target::CodePtr(dst_addr), None);
-    }
     asm.mark_branch_end(&branch);
 
     // If the block we're deferring from is empty
@@ -3958,6 +3953,8 @@ pub fn defer_compilation(
     }
 
     incr_counter!(defer_count);
+
+    Ok(())
 }
 
 /// Remove a block from the live control flow graph.