YJIT: replace BLOCKID_NULL with Option<BlockId>, more idiomatic (#5858)

* YJIT: replace BLOCKID_NULL with Option<BlockId>, more idiomatic

* Update yjit/src/core.rs

Co-authored-by: Alan Wu <XrXr@users.noreply.github.com>

* Update yjit/src/core.rs

Co-authored-by: Alan Wu <XrXr@users.noreply.github.com>

Co-authored-by: Alan Wu <XrXr@users.noreply.github.com>

2 files changed, 22 insertions, 29 deletions

diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 1f0de86168..17dc9657ac 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -5925,7 +5925,11 @@ mod tests {
     use super::*;
 
     fn setup_codegen() -> (JITState, Context, CodeBlock, OutlinedCb) {
-        let block = Block::new(BLOCKID_NULL, &Context::default());
+        let blockid = BlockId {
+            iseq: ptr::null(),
+            idx: 0,
+        };
+        let block = Block::new(blockid, &Context::default());
 
         return (
             JITState::new(&block),
diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index 5ea3ee8193..b8229167ed 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -219,12 +219,6 @@ pub struct BlockId {
     pub idx: u32,
 }
 
-/// Null block id constant
-pub const BLOCKID_NULL: BlockId = BlockId {
-    iseq: ptr::null(),
-    idx: 0,
-};
-
 /// Branch code shape enumeration
 #[derive(Copy, Clone, PartialEq, Eq, Debug)]
 pub enum BranchShape {
@@ -251,7 +245,7 @@ struct Branch {
     src_ctx: Context,
 
     // Branch target blocks and their contexts
-    targets: [BlockId; 2],
+    targets: [Option<BlockId>; 2],
     target_ctxs: [Context; 2],
     blocks: [Option<BlockRef>; 2],
 
@@ -510,7 +504,9 @@ pub extern "C" fn rb_yjit_iseq_mark(payload: *mut c_void) {
             for branch in &block.outgoing {
                 let branch = branch.borrow();
                 for target in &branch.targets {
-                    unsafe { rb_gc_mark_movable(target.iseq.into()) };
+                    if let Some(target) = target {
+                        unsafe { rb_gc_mark_movable(target.iseq.into()) };
+                    }
                 }
             }
 
@@ -566,7 +562,9 @@ pub extern "C" fn rb_yjit_iseq_update_references(payload: *mut c_void) {
             for branch in &block.outgoing {
                 let mut branch = branch.borrow_mut();
                 for target in &mut branch.targets {
-                    target.iseq = unsafe { rb_gc_location(target.iseq.into()) }.as_iseq();
+                    if let Some(target) = target {
+                        target.iseq = unsafe { rb_gc_location(target.iseq.into()) }.as_iseq();
+                    }
                 }
             }
 
@@ -1319,12 +1317,8 @@ fn gen_block_series_body(
         };
 
         // Get id and context for the new block
-        let requested_id = last_branch.targets[0];
+        let requested_id = last_branch.targets[0].expect("block id must be filled");
         let requested_ctx = &last_branch.target_ctxs[0];
-        assert_ne!(
-            last_branch.targets[0], BLOCKID_NULL,
-            "block id must be filled"
-        );
 
         // Generate new block using context from the last branch.
         let result = gen_single_block(requested_id, requested_ctx, ec, cb, ocb);
@@ -1479,7 +1473,7 @@ fn make_branch_entry(block: BlockRef, src_ctx: &Context, gen_fn: BranchGenFn) ->
         src_ctx: *src_ctx,
 
         // Branch target blocks and their contexts
-        targets: [BLOCKID_NULL, BLOCKID_NULL],
+        targets: [None, None],
         target_ctxs: [Context::default(), Context::default()],
         blocks: [None, None],
 
@@ -1529,7 +1523,7 @@ fn branch_stub_hit_body(branch_ptr: *const c_void, target_idx: u32, ec: EcPtr) -
     let branch_size_on_entry = branch.code_size();
 
     let target_idx: usize = target_idx.as_usize();
-    let target = branch.targets[target_idx];
+    let target = branch.targets[target_idx].unwrap();
     let target_ctx = branch.target_ctxs[target_idx];
 
     let target_branch_shape = match target_idx {
@@ -1720,8 +1714,6 @@ pub fn gen_branch(
     ctx1: Option<&Context>,
     gen_fn: BranchGenFn,
 ) {
-    assert!(target0 != BLOCKID_NULL);
-
     let branchref = make_branch_entry(jit.get_block(), src_ctx, gen_fn);
 
     // Get the branch targets or stubs
@@ -1738,10 +1730,8 @@ pub fn gen_branch(
     branch.dst_addrs[0] = dst_addr0;
     branch.dst_addrs[1] = dst_addr1;
 
-    branch.targets[0] = target0;
-    if target1.is_some() {
-        branch.targets[1] = target1.unwrap();
-    }
+    branch.targets[0] = Some(target0);
+    branch.targets[1] = target1;
     branch.target_ctxs[0] = *ctx0;
     branch.target_ctxs[1] = if ctx1.is_some() {
         *ctx1.unwrap()
@@ -1770,12 +1760,10 @@ fn gen_jump_branch(
 }
 
 pub fn gen_direct_jump(jit: &JITState, ctx: &Context, target0: BlockId, cb: &mut CodeBlock) {
-    assert!(target0 != BLOCKID_NULL);
-
     let branchref = make_branch_entry(jit.get_block(), ctx, gen_jump_branch);
     let mut branch = branchref.borrow_mut();
 
-    branch.targets[0] = target0;
+    branch.targets[0] = Some(target0);
     branch.target_ctxs[0] = *ctx;
 
     let maybe_block = find_block_version(target0, ctx);
@@ -1828,12 +1816,13 @@ pub fn defer_compilation(
     let mut branch = branch_rc.borrow_mut();
     let block = block_rc.borrow();
 
-    branch.target_ctxs[0] = next_ctx;
-    branch.targets[0] = BlockId {
+    let blockid = BlockId {
         iseq: block.blockid.iseq,
         idx: jit.get_insn_idx(),
     };
-    branch.dst_addrs[0] = get_branch_target(branch.targets[0], &next_ctx, &branch_rc, 0, ocb);
+    branch.target_ctxs[0] = next_ctx;
+    branch.targets[0] = Some(blockid);
+    branch.dst_addrs[0] = get_branch_target(blockid, &next_ctx, &branch_rc, 0, ocb);
 
     // Call the branch generation function
     branch.start_addr = Some(cb.get_write_ptr());