YJIT: Use a Box for branch targets to save memory

We frequently make branches that only have one target but we used to
always allocate space for two branch targets. This patch moves all the
information a branch target has into a struct and refer to them using
Option<Box<BranchTarget>>, this way when the second branch target is not
present it only takes 8 bytes.

Retained heap size on railsbench went from 16.17 MiB to 14.57 MiB, a
ratio of about 1.1.

1 files changed, 114 insertions, 111 deletions

diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index b45c3f010f..8a5d898c21 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -325,6 +325,15 @@ pub enum BranchShape {
 type BranchGenFn =
     fn(cb: &mut Assembler, target0: CodePtr, target1: Option<CodePtr>, shape: BranchShape) -> ();
 
+/// A place that a branch could jump to
+#[derive(Debug)]
+struct BranchTarget {
+    address: Option<CodePtr>,
+    id: BlockId,
+    ctx: Context,
+    block: Option<BlockRef>,
+}
+
 /// Store info about an outgoing branch in a code segment
 /// Note: care must be taken to minimize the size of branch objects
 struct Branch {
@@ -336,12 +345,7 @@ struct Branch {
     end_addr: Option<CodePtr>, // exclusive
 
     // Branch target blocks and their contexts
-    targets: [Option<BlockId>; 2],
-    target_ctxs: [Context; 2],
-    blocks: [Option<BlockRef>; 2],
-
-    // Jump target addresses
-    dst_addrs: [Option<CodePtr>; 2],
+    targets: [Option<Box<BranchTarget>>; 2],
 
     // Branch code generation function
     gen_fn: BranchGenFn,
@@ -368,6 +372,11 @@ impl Branch {
     fn code_size(&self) -> usize {
         (self.end_addr.unwrap().raw_ptr() as usize) - (self.start_addr.unwrap().raw_ptr() as usize)
     }
+
+    /// Get the address of one of the branch destination
+    fn get_target_address(&self, target_idx: usize) -> Option<CodePtr> {
+        self.targets[target_idx].as_ref().and_then(|target| target.address)
+    }
 }
 
 // In case a block is invalidated, this helps to remove all pointers to the block.
@@ -639,7 +648,7 @@ 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.iter().flatten() {
-                    unsafe { rb_gc_mark_movable(target.iseq.into()) };
+                    unsafe { rb_gc_mark_movable(target.id.iseq.into()) };
                 }
             }
 
@@ -694,7 +703,7 @@ 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 branch.targets.iter_mut().flatten() {
-                    target.iseq = unsafe { rb_gc_location(target.iseq.into()) }.as_iseq();
+                    target.id.iseq = unsafe { rb_gc_location(target.id.iseq.into()) }.as_iseq();
                 }
             }
 
@@ -1472,17 +1481,15 @@ fn gen_block_series_body(
         let mut last_branch = last_branchref.borrow_mut();
 
         // gen_direct_jump() can request a block to be placed immediately after by
-        // leaving `None`s in the `dst_addrs` array.
-        match &last_branch.dst_addrs {
-            [None, None] => (),
-            _ => {
-                break;
-            } // If there is no next block to compile, stop
+        // leaving a single target that has a `None` address.
+        let mut last_target = match &mut last_branch.targets {
+            [Some(last_target), None] if last_target.address.is_none() => last_target,
+            _ => break
         };
 
         // Get id and context for the new block
-        let requested_id = last_branch.targets[0].expect("block id must be filled");
-        let requested_ctx = &last_branch.target_ctxs[0];
+        let requested_id = last_target.id;
+        let requested_ctx = &last_target.ctx;
 
         // Generate new block using context from the last branch.
         let result = gen_single_block(requested_id, requested_ctx, ec, cb, ocb);
@@ -1507,8 +1514,8 @@ fn gen_block_series_body(
         add_block_version(&new_blockref, cb);
 
         // Connect the last branch and the new block
-        last_branch.blocks[0] = Some(new_blockref.clone());
-        last_branch.dst_addrs[0] = new_blockref.borrow().start_addr;
+        last_target.block = Some(new_blockref.clone());
+        last_target.address = new_blockref.borrow().start_addr;
         new_blockref
             .borrow_mut()
             .push_incoming(last_branchref.clone());
@@ -1600,15 +1607,14 @@ fn regenerate_branch(cb: &mut CodeBlock, branch: &mut Branch) {
 
     let mut block = branch.block.borrow_mut();
     let branch_terminates_block = branch.end_addr == block.end_addr;
-    assert!(branch.dst_addrs[0].is_some());
 
     // Generate the branch
     let mut asm = Assembler::new();
     asm.comment("regenerate_branch");
     (branch.gen_fn)(
         &mut asm,
-        branch.dst_addrs[0].unwrap(),
-        branch.dst_addrs[1],
+        branch.get_target_address(0).unwrap(),
+        branch.get_target_address(1),
         branch.shape,
     );
 
@@ -1656,11 +1662,6 @@ fn make_branch_entry(block: &BlockRef, gen_fn: BranchGenFn) -> BranchRef {
 
         // Branch target blocks and their contexts
         targets: [None, None],
-        target_ctxs: [Context::default(), Context::default()],
-        blocks: [None, None],
-
-        // Jump target addresses
-        dst_addrs: [None, None],
 
         // Branch code generation function
         gen_fn: gen_fn,
@@ -1680,7 +1681,7 @@ fn make_branch_entry(block: &BlockRef, gen_fn: BranchGenFn) -> BranchRef {
 
 c_callable! {
     /// Generated code calls this function with the SysV calling convention.
-    /// See [get_branch_target].
+    /// See [set_branch_target].
     fn branch_stub_hit(
         branch_ptr: *const c_void,
         target_idx: u32,
@@ -1713,8 +1714,9 @@ 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].unwrap();
-    let target_ctx = branch.target_ctxs[target_idx];
+    let target = branch.targets[target_idx].as_ref().unwrap();
+    let target_id = target.id;
+    let target_ctx = target.ctx;
 
     let target_branch_shape = match target_idx {
         0 => BranchShape::Next0,
@@ -1727,8 +1729,8 @@ fn branch_stub_hit_body(branch_ptr: *const c_void, target_idx: u32, ec: EcPtr) -
 
     // If this branch has already been patched, return the dst address
     // Note: ractors can cause the same stub to be hit multiple times
-    if let Some(_) = branch.blocks[target_idx] {
-        return branch.dst_addrs[target_idx].unwrap().raw_ptr();
+    if target.block.is_some() {
+        return target.address.unwrap().raw_ptr();
     }
 
     let (cfp, original_interp_sp) = unsafe {
@@ -1736,10 +1738,10 @@ fn branch_stub_hit_body(branch_ptr: *const c_void, target_idx: u32, ec: EcPtr) -
         let original_interp_sp = get_cfp_sp(cfp);
 
         let running_iseq = rb_cfp_get_iseq(cfp);
-        let reconned_pc = rb_iseq_pc_at_idx(running_iseq, target.idx);
+        let reconned_pc = rb_iseq_pc_at_idx(running_iseq, target_id.idx);
         let reconned_sp = original_interp_sp.offset(target_ctx.sp_offset.into());
 
-        assert_eq!(running_iseq, target.iseq as _, "each stub expects a particular iseq");
+        assert_eq!(running_iseq, target_id.iseq as _, "each stub expects a particular iseq");
 
         // Update the PC in the current CFP, because it may be out of sync in JITted code
         rb_set_cfp_pc(cfp, reconned_pc);
@@ -1756,7 +1758,7 @@ fn branch_stub_hit_body(branch_ptr: *const c_void, target_idx: u32, ec: EcPtr) -
     };
 
     // Try to find an existing compiled version of this block
-    let mut block = find_block_version(target, &target_ctx);
+    let mut block = find_block_version(target_id, &target_ctx);
 
     // If this block hasn't yet been compiled
     if block.is_none() {
@@ -1783,7 +1785,7 @@ fn branch_stub_hit_body(branch_ptr: *const c_void, target_idx: u32, ec: EcPtr) -
 
         // Compile the new block version
         drop(branch); // Stop mutable RefCell borrow since GC might borrow branch for marking
-        block = gen_block_series(target, &target_ctx, ec, cb, ocb);
+        block = gen_block_series(target_id, &target_ctx, ec, cb, ocb);
         branch = branch_rc.borrow_mut();
 
         if block.is_none() && branch_modified {
@@ -1806,11 +1808,12 @@ fn branch_stub_hit_body(branch_ptr: *const c_void, target_idx: u32, ec: EcPtr) -
             block.push_incoming(branch_rc.clone());
 
             // Update the branch target address
+            let target = branch.targets[target_idx].as_mut().unwrap();
             let dst_addr = block.start_addr;
-            branch.dst_addrs[target_idx] = dst_addr;
+            target.address = dst_addr;
 
             // Mark this branch target as patched (no longer a stub)
-            branch.blocks[target_idx] = Some(block_rc.clone());
+            target.block = Some(block_rc.clone());
 
             // Rewrite the branch with the new jump target address
             mem::drop(block); // end mut borrow
@@ -1859,14 +1862,15 @@ fn branch_stub_hit_body(branch_ptr: *const c_void, target_idx: u32, ec: EcPtr) -
     dst_addr.raw_ptr()
 }
 
-/// Get a block version or stub corresponding to a branch target
-fn get_branch_target(
+/// Set up a branch target at an index with a block version or a stub
+fn set_branch_target(
+    target_idx: u32,
     target: BlockId,
     ctx: &Context,
     branchref: &BranchRef,
-    target_idx: u32,
+    branch: &mut Branch,
     ocb: &mut OutlinedCb,
-) -> Option<CodePtr> {
+) {
     let maybe_block = find_block_version(target, ctx);
 
     // If the block already exists
@@ -1875,11 +1879,16 @@ fn get_branch_target(
 
         // Add an incoming branch into this block
         block.push_incoming(branchref.clone());
-        let mut branch = branchref.borrow_mut();
-        branch.blocks[target_idx.as_usize()] = Some(blockref.clone());
 
-        // Return a pointer to the compiled code for the block
-        return block.start_addr;
+        // Fill out the target with this block
+        branch.targets[target_idx.as_usize()] = Some(Box::new(BranchTarget {
+            block: Some(blockref.clone()),
+            address: block.start_addr,
+            id: target,
+            ctx: *ctx,
+        }));
+
+        return;
     }
 
     let ocb = ocb.unwrap();
@@ -1912,9 +1921,15 @@ fn get_branch_target(
     asm.compile(ocb);
 
     if ocb.has_dropped_bytes() {
-        None // No space
+        // No space
     } else {
-        Some(stub_addr)
+        // Fill the branch target with a stub
+        branch.targets[target_idx.as_usize()] = Some(Box::new(BranchTarget {
+            block: None, // no block yet
+            address: Some(stub_addr),
+            id: target,
+            ctx: *ctx,
+        }));
     }
 }
 
@@ -1958,37 +1973,21 @@ pub fn gen_branch(
     gen_fn: BranchGenFn,
 ) {
     let branchref = make_branch_entry(&jit.get_block(), gen_fn);
+    let branch = &mut branchref.borrow_mut();
 
     // Get the branch targets or stubs
-    let dst_addr0 = get_branch_target(target0, ctx0, &branchref, 0, ocb);
-    let dst_addr1 = if let Some(ctx) = ctx1 {
-        match get_branch_target(target1.unwrap(), ctx, &branchref, 1, ocb) {
-            Some(dst_addr) => Some(dst_addr),
-            None => return, // avoid unwrap() in gen_fn()
+    set_branch_target(0, target0, ctx0, &branchref, branch, ocb);
+    if let Some(ctx) = ctx1 {
+        set_branch_target(1, target1.unwrap(), ctx, &branchref, branch, ocb);
+        if branch.targets[1].is_none() {
+            return; // avoid unwrap() in gen_fn()
         }
-    } else {
-        None
-    };
-
-    let mut branch = branchref.borrow_mut();
-
-    // Set the branch target adresses
-    branch.dst_addrs[0] = dst_addr0;
-    branch.dst_addrs[1] = dst_addr1;
-
-    branch.targets[0] = Some(target0);
-    branch.targets[1] = target1;
-    branch.target_ctxs[0] = *ctx0;
-    branch.target_ctxs[1] = if let Some(&ctx) = ctx1 {
-        ctx
-    } else {
-        Context::default()
-    };
+    }
 
     // Call the branch generation function
     asm.mark_branch_start(&branchref);
-    if let Some(dst_addr) = branch.dst_addrs[0] {
-        gen_fn(asm, dst_addr, branch.dst_addrs[1], BranchShape::Default);
+    if let Some(dst_addr) = branch.get_target_address(0) {
+        gen_fn(asm, dst_addr, branch.get_target_address(1), BranchShape::Default);
     }
     asm.mark_branch_end(&branchref);
 }
@@ -2012,8 +2011,12 @@ pub fn gen_direct_jump(jit: &JITState, ctx: &Context, target0: BlockId, asm: &mu
     let branchref = make_branch_entry(&jit.get_block(), gen_jump_branch);
     let mut branch = branchref.borrow_mut();
 
-    branch.targets[0] = Some(target0);
-    branch.target_ctxs[0] = *ctx;
+    let mut new_target = BranchTarget {
+        block: None,
+        address: None,
+        ctx: *ctx,
+        id: target0,
+    };
 
     let maybe_block = find_block_version(target0, ctx);
 
@@ -2023,19 +2026,19 @@ pub fn gen_direct_jump(jit: &JITState, ctx: &Context, target0: BlockId, asm: &mu
 
         block.push_incoming(branchref.clone());
 
-        branch.dst_addrs[0] = block.start_addr;
-        branch.blocks[0] = Some(blockref.clone());
+        new_target.address = block.start_addr;
+        new_target.block = Some(blockref.clone());
         branch.shape = BranchShape::Default;
 
         // Call the branch generation function
         asm.comment("gen_direct_jmp: existing block");
         asm.mark_branch_start(&branchref);
-        gen_jump_branch(asm, branch.dst_addrs[0].unwrap(), None, BranchShape::Default);
+        gen_jump_branch(asm, new_target.address.unwrap(), None, BranchShape::Default);
         asm.mark_branch_end(&branchref);
     } else {
         // This None target address signals gen_block_series() to compile the
         // target block right after this one (fallthrough).
-        branch.dst_addrs[0] = None;
+        new_target.address = None;
         branch.shape = BranchShape::Next0;
 
         // The branch is effectively empty (a noop)
@@ -2043,6 +2046,8 @@ pub fn gen_direct_jump(jit: &JITState, ctx: &Context, target0: BlockId, asm: &mu
         asm.mark_branch_start(&branchref);
         asm.mark_branch_end(&branchref);
     }
+
+    branch.targets[0] = Some(Box::new(new_target));
 }
 
 /// Create a stub to force the code up to this point to be executed
@@ -2072,13 +2077,11 @@ pub fn defer_compilation(
         iseq: block.blockid.iseq,
         idx: jit.get_insn_idx(),
     };
-    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);
+    set_branch_target(0, blockid, &next_ctx, &branch_rc, &mut branch, ocb);
 
     // Call the branch generation function
     asm.mark_branch_start(&branch_rc);
-    if let Some(dst_addr) = branch.dst_addrs[0] {
+    if let Some(dst_addr) = branch.get_target_address(0) {
         gen_jump_branch(asm, dst_addr, None, BranchShape::Default);
     }
     asm.mark_branch_end(&branch_rc);
@@ -2098,11 +2101,9 @@ pub fn free_block(blockref: &BlockRef) {
         let mut pred_branch = pred_branchref.borrow_mut();
 
         // If this is us, nullify the target block
-        for pred_succ_ref in &mut pred_branch.blocks {
-            if let Some(pred_succ) = pred_succ_ref {
-                if pred_succ == blockref {
-                    *pred_succ_ref = None;
-                }
+        for pred_succ in pred_branch.targets.iter_mut().flatten() {
+            if pred_succ.block.as_ref() == Some(blockref) {
+                pred_succ.block = None;
             }
         }
     }
@@ -2112,12 +2113,14 @@ pub fn free_block(blockref: &BlockRef) {
         let out_branch = out_branchref.borrow();
 
         // For each successor block
-        for succ in out_branch.blocks.iter().flatten() {
-            // Remove outgoing branch from the successor's incoming list
-            let mut succ_block = succ.borrow_mut();
-            succ_block
-                .incoming
-                .retain(|succ_incoming| !Rc::ptr_eq(succ_incoming, out_branchref));
+        for out_target in out_branch.targets.iter().flatten() {
+            if let Some(succ_blockref) = &out_target.block {
+                // Remove outgoing branch from the successor's incoming list
+                let mut succ_block = succ_blockref.borrow_mut();
+                succ_block
+                    .incoming
+                    .retain(|succ_incoming| !Rc::ptr_eq(succ_incoming, out_branchref));
+            }
         }
     }
 
@@ -2161,7 +2164,7 @@ pub fn invalidate_block_version(blockref: &BlockRef) {
     remove_block_version(blockref);
 
     // Get a pointer to the generated code for this block
-    let code_ptr = block.start_addr;
+    let block_start = block.start_addr;
 
     // Make the the start of the block do an exit. This handles OOM situations
     // and some cases where we can't efficiently patch incoming branches.
@@ -2170,6 +2173,9 @@ pub fn invalidate_block_version(blockref: &BlockRef) {
     // In those situations, there is hopefully no jumps to the start of the block
     // after patching as the start of the block would be in the middle of something
     // generated by branch_t::gen_fn.
+    let block_entry_exit = block
+        .entry_exit
+        .expect("invalidation needs the entry_exit field");
     {
         let block_start = block
             .start_addr
@@ -2177,9 +2183,6 @@ pub fn invalidate_block_version(blockref: &BlockRef) {
         let block_end = block
             .end_addr
             .expect("invalidation needs constructed block");
-        let block_entry_exit = block
-            .entry_exit
-            .expect("invalidation needs the entry_exit field");
 
         if block_start == block_entry_exit {
             // Some blocks exit on entry. Patching a jump to the entry at the
@@ -2213,16 +2216,16 @@ pub fn invalidate_block_version(blockref: &BlockRef) {
     // For each incoming branch
     for branchref in &block.incoming {
         let mut branch = branchref.borrow_mut();
-        let target_idx = if branch.dst_addrs[0] == code_ptr {
+        let target_idx = if branch.get_target_address(0) == block_start {
             0
         } else {
             1
         };
-        assert_eq!(branch.dst_addrs[target_idx], code_ptr);
-        assert_eq!(blockref, branch.blocks[target_idx].as_ref().unwrap());
 
-        // Mark this target as being a stub
-        branch.blocks[target_idx] = None;
+        // Assert that the incoming branch indeed points to the block being invalidated
+        let incoming_target = branch.targets[target_idx].as_ref().unwrap();
+        assert_eq!(block_start, incoming_target.address);
+        assert_eq!(blockref, incoming_target.block.as_ref().unwrap());
 
         // TODO(alan):
         // Don't patch frozen code region
@@ -2230,23 +2233,23 @@ pub fn invalidate_block_version(blockref: &BlockRef) {
         //     continue;
         // }
 
-        // Create a stub for this branch target
-        mem::drop(branch); // end RefCell borrow as get_branch_target() can borrow the branch.
-        let mut branch_target =
-            get_branch_target(block.blockid, &block.ctx, branchref, target_idx as u32, ocb);
+        // Create a stub for this branch target or rewire it to a valid block
+        set_branch_target(target_idx as u32, block.blockid, &block.ctx, branchref, &mut branch, ocb);
 
-        if branch_target.is_none() {
+        if branch.targets[target_idx].is_none() {
             // We were unable to generate a stub (e.g. OOM). Use the block's
             // exit instead of a stub for the block. It's important that we
             // still patch the branch in this situation so stubs are unique
             // to branches. Think about what could go wrong if we run out of
             // memory in the middle of this loop.
-            branch_target = block.entry_exit;
+            branch.targets[target_idx] = Some(Box::new(BranchTarget {
+                block: None,
+                address: block.entry_exit,
+                id: block.blockid,
+                ctx: block.ctx,
+            }));
         }
 
-        branch = branchref.borrow_mut();
-        branch.dst_addrs[target_idx] = branch_target;
-
         // Check if the invalidated block immediately follows
         let target_next = block.start_addr == branch.end_addr;