1 files changed, 123 insertions, 102 deletions
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 31582c26bc..e9c17cb537 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -12,6 +12,7 @@ use crate::utils::*;
 use CodegenStatus::*;
 use YARVOpnd::*;
 
+use std::cell::Cell;
 use std::cmp;
 use std::collections::HashMap;
 use std::ffi::CStr;
@@ -38,63 +39,89 @@ type InsnGenFn = fn(
     ocb: &mut OutlinedCb,
 ) -> CodegenStatus;
 
-/// Code generation state
-/// This struct only lives while code is being generated
+/// Ephemeral code generation state.
+/// Represents a [core::Block] while we build it.
 pub struct JITState {
-    // Block version being compiled
-    block: BlockRef,
-
-    // Instruction sequence this is associated with
+    /// Instruction sequence for the compiling block
     iseq: IseqPtr,
 
-    // Index of the current instruction being compiled
-    insn_idx: u16,
+    /// The iseq index of the first instruction in the block
+    starting_insn_idx: IseqIdx,
+
+    /// The [Context] entering into the first instruction of the block
+    starting_ctx: Context,
+
+    /// The placement for the machine code of the [Block]
+    output_ptr: CodePtr,
 
-    // Opcode for the instruction being compiled
+    /// Index of the current instruction being compiled
+    insn_idx: IseqIdx,
+
+    /// Opcode for the instruction being compiled
     opcode: usize,
 
-    // PC of the instruction being compiled
+    /// PC of the instruction being compiled
     pc: *mut VALUE,
 
-    // Side exit to the instruction being compiled. See :side-exit:.
+    /// Side exit to the instruction being compiled. See :side-exit:.
     side_exit_for_pc: Option<CodePtr>,
 
-    // Execution context when compilation started
-    // This allows us to peek at run-time values
+    /// Execution context when compilation started
+    /// This allows us to peek at run-time values
     ec: EcPtr,
 
-    // Whether we need to record the code address at
-    // the end of this bytecode instruction for global invalidation
-    record_boundary_patch_point: bool,
+    /// The outgoing branches the block will have
+    pub pending_outgoing: Vec<PendingBranchRef>,
+
+    // --- Fields for block invalidation and invariants tracking below:
+    // Public mostly so into_block defined in the sibling module core
+    // can partially move out of Self.
 
-    // The block's outgoing branches
-    outgoing: Vec<BranchRef>,
+    /// Whether we need to record the code address at
+    /// the end of this bytecode instruction for global invalidation
+    pub record_boundary_patch_point: bool,
 
-    // The block's CME dependencies
-    cme_dependencies: Vec<CmePtr>,
+    /// Code for immediately exiting upon entry to the block.
+    /// Required for invalidation.
+    pub block_entry_exit: Option<CodePtr>,
+
+    /// A list of callable method entries that must be valid for the block to be valid.
+    pub method_lookup_assumptions: Vec<CmePtr>,
+
+    /// A list of basic operators that not be redefined for the block to be valid.
+    pub bop_assumptions: Vec<(RedefinitionFlag, ruby_basic_operators)>,
+
+    /// A list of constant expression path segments that must have
+    /// not been written to for the block to be valid.
+    pub stable_constant_names_assumption: Option<*const ID>,
+
+    /// When true, the block is valid only when there is a total of one ractor running
+    pub block_assumes_single_ractor: bool,
 }
 
 impl JITState {
-    pub fn new(blockref: &BlockRef, ec: EcPtr) -> Self {
+    pub fn new(blockid: BlockId, starting_ctx: Context, output_ptr: CodePtr, ec: EcPtr) -> Self {
         JITState {
-            block: blockref.clone(),
-            iseq: ptr::null(), // TODO: initialize this from the blockid
+            iseq: blockid.iseq,
+            starting_insn_idx: blockid.idx,
+            starting_ctx,
+            output_ptr,
             insn_idx: 0,
             opcode: 0,
             pc: ptr::null_mut::<VALUE>(),
             side_exit_for_pc: None,
+            pending_outgoing: vec![],
             ec,
             record_boundary_patch_point: false,
-            outgoing: Vec::new(),
-            cme_dependencies: Vec::new(),
+            block_entry_exit: None,
+            method_lookup_assumptions: vec![],
+            bop_assumptions: vec![],
+            stable_constant_names_assumption: None,
+            block_assumes_single_ractor: false,
         }
     }
 
-    pub fn get_block(&self) -> BlockRef {
-        self.block.clone()
-    }
-
-    pub fn get_insn_idx(&self) -> u16 {
+    pub fn get_insn_idx(&self) -> IseqIdx {
         self.insn_idx
     }
 
@@ -110,6 +137,18 @@ impl JITState {
         self.pc
     }
 
+    pub fn get_starting_insn_idx(&self) -> IseqIdx {
+        self.starting_insn_idx
+    }
+
+    pub fn get_block_entry_exit(&self) -> Option<CodePtr> {
+        self.block_entry_exit
+    }
+
+    pub fn get_starting_ctx(&self) -> Context {
+        self.starting_ctx.clone()
+    }
+
     pub fn get_arg(&self, arg_idx: isize) -> VALUE {
         // insn_len require non-test config
         #[cfg(not(test))]
@@ -175,18 +214,22 @@ impl JITState {
         }
     }
 
+    pub fn assume_method_lookup_stable(&mut self, ocb: &mut OutlinedCb, cme: CmePtr) {
+        jit_ensure_block_entry_exit(self, ocb);
+        self.method_lookup_assumptions.push(cme);
+    }
+
     fn get_cfp(&self) -> *mut rb_control_frame_struct {
         unsafe { get_ec_cfp(self.ec) }
     }
 
-    // Push an outgoing branch ref
-    pub fn push_outgoing(&mut self, branch: BranchRef) {
-        self.outgoing.push(branch);
+    pub fn assume_stable_constant_names(&mut self, ocb: &mut OutlinedCb, id: *const ID) {
+        jit_ensure_block_entry_exit(self, ocb);
+        self.stable_constant_names_assumption = Some(id);
     }
 
-    // Push a CME dependency
-    pub fn push_cme_dependency(&mut self, cme: CmePtr) {
-        self.cme_dependencies.push(cme);
+    pub fn queue_outgoing_branch(&mut self, branch: PendingBranchRef) {
+        self.pending_outgoing.push(branch)
     }
 }
 
@@ -498,22 +541,20 @@ fn get_side_exit(jit: &mut JITState, ocb: &mut OutlinedCb, ctx: &Context) -> Tar
 // Ensure that there is an exit for the start of the block being compiled.
 // Block invalidation uses this exit.
 pub fn jit_ensure_block_entry_exit(jit: &mut JITState, ocb: &mut OutlinedCb) {
-    let blockref = jit.block.clone();
-    let mut block = blockref.borrow_mut();
-    let block_ctx = block.get_ctx();
-    let blockid = block.get_blockid();
-
-    if block.entry_exit.is_some() {
+    if jit.block_entry_exit.is_some() {
         return;
     }
 
+    let block_starting_context = &jit.get_starting_ctx();
+
     // If we're compiling the first instruction in the block.
-    if jit.insn_idx == blockid.idx {
+    if jit.insn_idx == jit.starting_insn_idx {
         // Generate the exit with the cache in jitstate.
-        block.entry_exit = Some(get_side_exit(jit, ocb, &block_ctx).unwrap_code_ptr());
+        let entry_exit = get_side_exit(jit, ocb, block_starting_context).unwrap_code_ptr();
+        jit.block_entry_exit = Some(entry_exit);
     } else {
-        let block_entry_pc = unsafe { rb_iseq_pc_at_idx(blockid.iseq, blockid.idx.into()) };
-        block.entry_exit = Some(gen_outlined_exit(block_entry_pc, &block_ctx, ocb));
+        let block_entry_pc = unsafe { rb_iseq_pc_at_idx(jit.iseq, jit.starting_insn_idx.into()) };
+        jit.block_entry_exit = Some(gen_outlined_exit(block_entry_pc, block_starting_context, ocb));
     }
 }
 
@@ -725,18 +766,18 @@ pub fn gen_single_block(
     verify_blockid(blockid);
     assert!(!(blockid.idx == 0 && ctx.get_stack_size() > 0));
 
+    // Save machine code placement of the block. `cb` might page switch when we
+    // generate code in `ocb`.
+    let block_start_addr = cb.get_write_ptr();
+
     // Instruction sequence to compile
     let iseq = blockid.iseq;
     let iseq_size = unsafe { get_iseq_encoded_size(iseq) };
     let iseq_size: u16 = iseq_size.try_into().unwrap();
     let mut insn_idx: u16 = blockid.idx;
-    let starting_insn_idx = insn_idx;
-
-    // Allocate the new block
-    let blockref = Block::new(blockid, &ctx, cb.get_write_ptr());
 
     // Initialize a JIT state object
-    let mut jit = JITState::new(&blockref, ec);
+    let mut jit = JITState::new(blockid, ctx.clone(), cb.get_write_ptr(), ec);
     jit.iseq = blockid.iseq;
 
     // Create a backend assembler instance
@@ -762,7 +803,7 @@ pub fn gen_single_block(
         // We need opt_getconstant_path to be in a block all on its own. Cut the block short
         // if we run into it. This is necessary because we want to invalidate based on the
         // instruction's index.
-        if opcode == YARVINSN_opt_getconstant_path.as_usize() && insn_idx > starting_insn_idx {
+        if opcode == YARVINSN_opt_getconstant_path.as_usize() && insn_idx > jit.starting_insn_idx {
             jump_to_next_insn(&mut jit, &ctx, &mut asm, ocb);
             break;
         }
@@ -814,17 +855,15 @@ pub fn gen_single_block(
                 println!("can't compile {}", insn_name(opcode));
             }
 
-            let mut block = jit.block.borrow_mut();
-
             // TODO: if the codegen function makes changes to ctx and then return YJIT_CANT_COMPILE,
             // the exit this generates would be wrong. We could save a copy of the entry context
             // and assert that ctx is the same here.
             gen_exit(jit.pc, &ctx, &mut asm);
 
-            // If this is the first instruction in the block, then we can use
-            // the exit for block->entry_exit.
-            if insn_idx == block.get_blockid().idx {
-                block.entry_exit = Some(block.get_start_addr());
+            // If this is the first instruction in the block, then
+            // the entry address is the address for block_entry_exit
+            if insn_idx == jit.starting_insn_idx {
+                jit.block_entry_exit = Some(jit.output_ptr);
             }
 
             break;
@@ -842,45 +881,28 @@ pub fn gen_single_block(
             break;
         }
     }
+    let end_insn_idx = insn_idx;
 
-    // Finish filling out the block
-    {
-        let mut block = jit.block.borrow_mut();
-        if block.entry_exit.is_some() {
-            asm.pad_inval_patch();
-        }
-
-        // Compile code into the code block
-        let gc_offsets = asm.compile(cb);
-
-        // Add the GC offsets to the block
-        block.set_gc_obj_offsets(gc_offsets);
-
-        // Set CME dependencies to the block
-        block.set_cme_dependencies(jit.cme_dependencies);
-
-        // Set outgoing branches to the block
-        block.set_outgoing(jit.outgoing);
-
-        // Mark the end position of the block
-        block.set_end_addr(cb.get_write_ptr());
+    // We currently can't handle cases where the request is for a block that
+    // doesn't go to the next instruction in the same iseq.
+    assert!(!jit.record_boundary_patch_point);
 
-        // Store the index of the last instruction in the block
-        block.set_end_idx(insn_idx);
+    // Pad the block if it has the potential to be invalidated
+    if jit.block_entry_exit.is_some() {
+        asm.pad_inval_patch();
     }
 
-    // We currently can't handle cases where the request is for a block that
-    // doesn't go to the next instruction.
-    assert!(!jit.record_boundary_patch_point);
+    // Compile code into the code block
+    let gc_offsets = asm.compile(cb);
+    let end_addr = cb.get_write_ptr();
 
     // If code for the block doesn't fit, fail
     if cb.has_dropped_bytes() || ocb.unwrap().has_dropped_bytes() {
-        free_block(&blockref);
         return Err(());
     }
 
     // Block compiled successfully
-    Ok(blockref)
+    Ok(jit.into_block(end_insn_idx, block_start_addr, end_addr, gc_offsets))
 }
 
 fn gen_nop(
@@ -3595,7 +3617,7 @@ fn gen_branchif(
             ctx,
             Some(next_block),
             Some(ctx),
-            BranchGenFn::BranchIf(BranchShape::Default),
+            BranchGenFn::BranchIf(Cell::new(BranchShape::Default)),
         );
     }
 
@@ -3652,7 +3674,7 @@ fn gen_branchunless(
             ctx,
             Some(next_block),
             Some(ctx),
-            BranchGenFn::BranchUnless(BranchShape::Default),
+            BranchGenFn::BranchUnless(Cell::new(BranchShape::Default)),
         );
     }
 
@@ -3706,7 +3728,7 @@ fn gen_branchnil(
             ctx,
             Some(next_block),
             Some(ctx),
-            BranchGenFn::BranchNil(BranchShape::Default),
+            BranchGenFn::BranchNil(Cell::new(BranchShape::Default)),
         );
     }
 
@@ -4533,7 +4555,7 @@ fn jit_obj_respond_to(
     // Invalidate this block if method lookup changes for the method being queried. This works
     // both for the case where a method does or does not exist, as for the latter we asked for a
     // "negative CME" earlier.
-    assume_method_lookup_stable(jit, ocb, target_cme);
+    jit.assume_method_lookup_stable(ocb, target_cme);
 
     // Generate a side exit
     let side_exit = get_side_exit(jit, ocb, ctx);
@@ -6308,7 +6330,7 @@ fn gen_send_general(
 
     // Register block for invalidation
     //assert!(cme->called_id == mid);
-    assume_method_lookup_stable(jit, ocb, cme);
+    jit.assume_method_lookup_stable(ocb, cme);
 
     // To handle the aliased method case (VM_METHOD_TYPE_ALIAS)
     loop {
@@ -6478,7 +6500,7 @@ fn gen_send_general(
 
                         flags |= VM_CALL_FCALL | VM_CALL_OPT_SEND;
 
-                        assume_method_lookup_stable(jit, ocb, cme);
+                        jit.assume_method_lookup_stable(ocb, cme);
 
                         let (known_class, type_mismatch_exit) = {
                             if compile_time_name.string_p() {
@@ -6973,8 +6995,8 @@ fn gen_invokesuper(
 
     // We need to assume that both our current method entry and the super
     // method entry we invoke remain stable
-    assume_method_lookup_stable(jit, ocb, me);
-    assume_method_lookup_stable(jit, ocb, cme);
+    jit.assume_method_lookup_stable(ocb, me);
+    jit.assume_method_lookup_stable(ocb, cme);
 
     // Method calls may corrupt types
     ctx.clear_local_types();
@@ -7428,14 +7450,13 @@ fn gen_opt_getconstant_path(
         asm.store(stack_top, ic_entry_val);
     } else {
         // Optimize for single ractor mode.
-        // FIXME: This leaks when st_insert raises NoMemoryError
         if !assume_single_ractor_mode(jit, ocb) {
             return CantCompile;
         }
 
         // Invalidate output code on any constant writes associated with
         // constants referenced within the current block.
-        assume_stable_constant_names(jit, ocb, idlist);
+        jit.assume_stable_constant_names(ocb, idlist);
 
         jit_putobject(jit, ctx, asm, unsafe { (*ice).value });
     }
@@ -8112,15 +8133,15 @@ mod tests {
     use super::*;
 
     fn setup_codegen() -> (JITState, Context, Assembler, CodeBlock, OutlinedCb) {
-        let blockid = BlockId {
-            iseq: ptr::null(),
-            idx: 0,
-        };
         let cb = CodeBlock::new_dummy(256 * 1024);
-        let block = Block::new(blockid, &Context::default(), cb.get_write_ptr());
 
         return (
-            JITState::new(&block, ptr::null()), // No execution context in tests. No peeking!
+            JITState::new(
+                BlockId { iseq: std::ptr::null(), idx: 0 },
+                Context::default(),
+                cb.get_write_ptr(),
+                ptr::null(), // No execution context in tests. No peeking!
+            ),
             Context::default(),
             Assembler::new(),
             cb,