YJIT: Introduce Opnd::Stack (#7352)

5 files changed, 80 insertions, 50 deletions

diff --git a/yjit/src/backend/arm64/mod.rs b/yjit/src/backend/arm64/mod.rs
index d06f05dcb1..fc26b9d893 100644
--- a/yjit/src/backend/arm64/mod.rs
+++ b/yjit/src/backend/arm64/mod.rs
@@ -54,6 +54,7 @@ impl From<Opnd> for A64Opnd {
             },
             Opnd::InsnOut { .. } => panic!("attempted to lower an Opnd::InsnOut"),
             Opnd::Value(_) => panic!("attempted to lower an Opnd::Value"),
+            Opnd::Stack { .. } => panic!("attempted to lower an Opnd::Stack"),
             Opnd::None => panic!(
                 "Attempted to lower an Opnd::None. This often happens when an out operand was not allocated for an instruction because the output of the instruction was not used. Please ensure you are using the output."
             ),
@@ -252,7 +253,7 @@ impl Assembler
         /// do follow that encoding, and if they don't then we load them first.
         fn split_bitmask_immediate(asm: &mut Assembler, opnd: Opnd, dest_num_bits: u8) -> Opnd {
             match opnd {
-                Opnd::Reg(_) | Opnd::InsnOut { .. } => opnd,
+                Opnd::Reg(_) | Opnd::InsnOut { .. } | Opnd::Stack { .. } => opnd,
                 Opnd::Mem(_) => split_load_operand(asm, opnd),
                 Opnd::Imm(imm) => {
                     if imm == 0 {
@@ -295,7 +296,7 @@ impl Assembler
                         asm.load(opnd)
                     }
                 },
-                Opnd::None | Opnd::Value(_) => unreachable!()
+                Opnd::None | Opnd::Value(_) | Opnd::Stack { .. } => unreachable!()
             }
         }
 
@@ -863,6 +864,9 @@ impl Assembler
                             let ptr_offset: u32 = (cb.get_write_pos() as u32) - (SIZEOF_VALUE as u32);
                             insn_gc_offsets.push(ptr_offset);
                         },
+                        Opnd::Stack { .. } => {
+                            unreachable!("Stack operand was not lowered before arm64_emit");
+                        }
                         Opnd::None => {
                             unreachable!("Attempted to load from None operand");
                         }
@@ -1072,7 +1076,7 @@ impl Assembler
     /// Optimize and compile the stored instructions
     pub fn compile_with_regs(self, cb: &mut CodeBlock, regs: Vec<Reg>) -> Vec<u32>
     {
-        let mut asm = self.arm64_split().alloc_regs(regs);
+        let mut asm = self.lower_stack().arm64_split().alloc_regs(regs);
 
         // Create label instances in the code block
         for (idx, name) in asm.label_names.iter().enumerate() {
diff --git a/yjit/src/backend/ir.rs b/yjit/src/backend/ir.rs
index 1dea189f24..dd0390f39c 100644
--- a/yjit/src/backend/ir.rs
+++ b/yjit/src/backend/ir.rs
@@ -7,7 +7,7 @@ use std::fmt;
 use std::convert::From;
 use std::io::Write;
 use std::mem::take;
-use crate::cruby::{VALUE};
+use crate::cruby::{VALUE, SIZEOF_VALUE_I32};
 use crate::virtualmem::{CodePtr};
 use crate::asm::{CodeBlock, uimm_num_bits, imm_num_bits};
 use crate::core::{Context, Type, TempMapping};
@@ -72,6 +72,9 @@ pub enum Opnd
     // Output of a preceding instruction in this block
     InsnOut{ idx: usize, num_bits: u8 },
 
+    // Pointer to a slot on the VM stack
+    Stack { idx: i32, sp_offset: i16, num_bits: u8 },
+
     // Low-level operands, for lowering
     Imm(i64),           // Raw signed immediate
     UImm(u64),          // Raw unsigned immediate
@@ -85,6 +88,7 @@ impl fmt::Debug for Opnd {
         match self {
             Self::None => write!(fmt, "None"),
             Value(val) => write!(fmt, "Value({val:?})"),
+            Stack { idx, sp_offset, .. } => write!(fmt, "SP[{}]", *sp_offset as i32 - idx - 1),
             InsnOut { idx, num_bits } => write!(fmt, "Out{num_bits}({idx})"),
             Imm(signed) => write!(fmt, "{signed:x}_i64"),
             UImm(unsigned) => write!(fmt, "{unsigned:x}_u64"),
@@ -158,6 +162,7 @@ impl Opnd
             Opnd::Reg(reg) => Some(Opnd::Reg(reg.with_num_bits(num_bits))),
             Opnd::Mem(Mem { base, disp, .. }) => Some(Opnd::Mem(Mem { base, disp, num_bits })),
             Opnd::InsnOut { idx, .. } => Some(Opnd::InsnOut { idx, num_bits }),
+            Opnd::Stack { idx, sp_offset, .. } => Some(Opnd::Stack { idx, sp_offset, num_bits }),
             _ => None,
         }
     }
@@ -914,6 +919,25 @@ impl Assembler
         Target::Label(label_idx)
     }
 
+    /// Convert Stack operands to memory operands
+    pub fn lower_stack(mut self) -> Assembler
+    {
+        let mut asm = Assembler::new_with_label_names(take(&mut self.label_names));
+        let mut iterator = self.into_draining_iter();
+
+        while let Some((index, mut insn)) = iterator.next_unmapped() {
+            let mut opnd_iter = insn.opnd_iter_mut();
+            while let Some(opnd) = opnd_iter.next() {
+                if let Opnd::Stack { idx, sp_offset, num_bits } = *opnd {
+                    *opnd = Opnd::mem(num_bits, SP, (sp_offset as i32 - idx - 1) * SIZEOF_VALUE_I32);
+                }
+            }
+            asm.push_insn(insn);
+        }
+
+        asm
+    }
+
     /// Sets the out field on the various instructions that require allocated
     /// registers because their output is used as the operand on a subsequent
     /// instruction. This is our implementation of the linear scan algorithm.
diff --git a/yjit/src/backend/x86_64/mod.rs b/yjit/src/backend/x86_64/mod.rs
index 297a0fd852..0ea943b75e 100644
--- a/yjit/src/backend/x86_64/mod.rs
+++ b/yjit/src/backend/x86_64/mod.rs
@@ -701,7 +701,7 @@ impl Assembler
     /// Optimize and compile the stored instructions
     pub fn compile_with_regs(self, cb: &mut CodeBlock, regs: Vec<Reg>) -> Vec<u32>
     {
-        let mut asm = self.x86_split().alloc_regs(regs);
+        let mut asm = self.lower_stack().x86_split().alloc_regs(regs);
 
         // Create label instances in the code block
         for (idx, name) in asm.label_names.iter().enumerate() {
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 564a39d42c..5adeddf6df 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -880,9 +880,8 @@ fn gen_dup(
     asm: &mut Assembler,
     _ocb: &mut OutlinedCb,
 ) -> CodegenStatus {
-
-    let dup_val = ctx.stack_pop(0);
-    let (mapping, tmp_type) = ctx.get_opnd_mapping(StackOpnd(0));
+    let dup_val = ctx.stack_opnd(0);
+    let (mapping, tmp_type) = ctx.get_opnd_mapping(dup_val.into());
 
     let loc0 = ctx.stack_push_mapping((mapping, tmp_type));
     asm.mov(loc0, dup_val);
@@ -907,8 +906,8 @@ fn gen_dupn(
     let opnd1: Opnd = ctx.stack_opnd(1);
     let opnd0: Opnd = ctx.stack_opnd(0);
 
-    let mapping1 = ctx.get_opnd_mapping(StackOpnd(1));
-    let mapping0 = ctx.get_opnd_mapping(StackOpnd(0));
+    let mapping1 = ctx.get_opnd_mapping(opnd1.into());
+    let mapping0 = ctx.get_opnd_mapping(opnd0.into());
 
     let dst1: Opnd = ctx.stack_push_mapping(mapping1);
     asm.mov(dst1, opnd1);
@@ -940,16 +939,16 @@ fn stack_swap(
     let stack0_mem = ctx.stack_opnd(offset0 as i32);
     let stack1_mem = ctx.stack_opnd(offset1 as i32);
 
-    let mapping0 = ctx.get_opnd_mapping(StackOpnd(offset0));
-    let mapping1 = ctx.get_opnd_mapping(StackOpnd(offset1));
+    let mapping0 = ctx.get_opnd_mapping(stack0_mem.into());
+    let mapping1 = ctx.get_opnd_mapping(stack1_mem.into());
 
     let stack0_reg = asm.load(stack0_mem);
     let stack1_reg = asm.load(stack1_mem);
     asm.mov(stack0_mem, stack1_reg);
     asm.mov(stack1_mem, stack0_reg);
 
-    ctx.set_opnd_mapping(StackOpnd(offset0), mapping1);
-    ctx.set_opnd_mapping(StackOpnd(offset1), mapping0);
+    ctx.set_opnd_mapping(stack0_mem.into(), mapping1);
+    ctx.set_opnd_mapping(stack1_mem.into(), mapping0);
 }
 
 fn gen_putnil(
@@ -1043,15 +1042,15 @@ fn gen_setn(
 ) -> CodegenStatus {
     let n = jit.get_arg(0).as_usize();
 
-    let top_val = ctx.stack_pop(0);
+    let top_val = ctx.stack_opnd(0);
     let dst_opnd = ctx.stack_opnd(n.try_into().unwrap());
     asm.mov(
         dst_opnd,
         top_val
     );
 
-    let mapping = ctx.get_opnd_mapping(StackOpnd(0));
-    ctx.set_opnd_mapping(StackOpnd(n.try_into().unwrap()), mapping);
+    let mapping = ctx.get_opnd_mapping(top_val.into());
+    ctx.set_opnd_mapping(dst_opnd.into(), mapping);
 
     KeepCompiling
 }
@@ -1066,7 +1065,7 @@ fn gen_topn(
     let n = jit.get_arg(0).as_usize();
 
     let top_n_val = ctx.stack_opnd(n.try_into().unwrap());
-    let mapping = ctx.get_opnd_mapping(StackOpnd(n.try_into().unwrap()));
+    let mapping = ctx.get_opnd_mapping(top_n_val.into());
     let loc0 = ctx.stack_push_mapping(mapping);
     asm.mov(loc0, top_n_val);
 
@@ -2481,9 +2480,13 @@ fn guard_two_fixnums(
     ocb: &mut OutlinedCb,
     side_exit: Target
 ) {
+    // Get stack operands without popping them
+    let arg1 = ctx.stack_opnd(0);
+    let arg0 = ctx.stack_opnd(1);
+
     // Get the stack operand types
-    let arg1_type = ctx.get_opnd_type(StackOpnd(0));
-    let arg0_type = ctx.get_opnd_type(StackOpnd(1));
+    let arg1_type = ctx.get_opnd_type(arg1.into());
+    let arg0_type = ctx.get_opnd_type(arg0.into());
 
     if arg0_type.is_heap() || arg1_type.is_heap() {
         asm.comment("arg is heap object");
@@ -2508,10 +2511,6 @@ fn guard_two_fixnums(
     assert!(arg0_type == Type::Fixnum || arg0_type.is_unknown());
     assert!(arg1_type == Type::Fixnum || arg1_type.is_unknown());
 
-    // Get stack operands without popping them
-    let arg1 = ctx.stack_opnd(0);
-    let arg0 = ctx.stack_opnd(1);
-
     // If not fixnums at run-time, fall back
     if arg0_type != Type::Fixnum {
         asm.comment("guard arg0 fixnum");
@@ -2543,8 +2542,8 @@ fn guard_two_fixnums(
     }
 
     // Set stack types in context
-    ctx.upgrade_opnd_type(StackOpnd(0), Type::Fixnum);
-    ctx.upgrade_opnd_type(StackOpnd(1), Type::Fixnum);
+    ctx.upgrade_opnd_type(arg1.into(), Type::Fixnum);
+    ctx.upgrade_opnd_type(arg0.into(), Type::Fixnum);
 }
 
 // Conditional move operation used by comparison operators
@@ -2697,7 +2696,7 @@ fn gen_equality_specialized(
             ocb,
             unsafe { rb_cString },
             a_opnd,
-            StackOpnd(1),
+            a_opnd.into(),
             comptime_a,
             SEND_MAX_DEPTH,
             side_exit,
@@ -2711,7 +2710,7 @@ fn gen_equality_specialized(
         asm.je(equal);
 
         // Otherwise guard that b is a T_STRING (from type info) or String (from runtime guard)
-        let btype = ctx.get_opnd_type(StackOpnd(0));
+        let btype = ctx.get_opnd_type(b_opnd.into());
         if btype.known_value_type() != Some(RUBY_T_STRING) {
             // Note: any T_STRING is valid here, but we check for a ::String for simplicity
             // To pass a mutable static variable (rb_cString) requires an unsafe block
@@ -2722,7 +2721,7 @@ fn gen_equality_specialized(
                 ocb,
                 unsafe { rb_cString },
                 b_opnd,
-                StackOpnd(0),
+                b_opnd.into(),
                 comptime_b,
                 SEND_MAX_DEPTH,
                 side_exit,
@@ -2833,7 +2832,7 @@ fn gen_opt_aref(
             ocb,
             unsafe { rb_cArray },
             recv_opnd,
-            StackOpnd(1),
+            recv_opnd.into(),
             comptime_recv,
             OPT_AREF_MAX_CHAIN_DEPTH,
             side_exit,
@@ -2876,7 +2875,7 @@ fn gen_opt_aref(
             ocb,
             unsafe { rb_cHash },
             recv_opnd,
-            StackOpnd(1),
+            recv_opnd.into(),
             comptime_recv,
             OPT_AREF_MAX_CHAIN_DEPTH,
             side_exit,
@@ -2937,7 +2936,7 @@ fn gen_opt_aset(
             ocb,
             unsafe { rb_cArray },
             recv,
-            StackOpnd(2),
+            recv.into(),
             comptime_recv,
             SEND_MAX_DEPTH,
             side_exit,
@@ -2951,7 +2950,7 @@ fn gen_opt_aset(
             ocb,
             unsafe { rb_cInteger },
             key,
-            StackOpnd(1),
+            key.into(),
             comptime_key,
             SEND_MAX_DEPTH,
             side_exit,
@@ -2989,7 +2988,7 @@ fn gen_opt_aset(
             ocb,
             unsafe { rb_cHash },
             recv,
-            StackOpnd(2),
+            recv.into(),
             comptime_recv,
             SEND_MAX_DEPTH,
             side_exit,
@@ -5791,8 +5790,8 @@ fn gen_send_iseq(
     // side exits, so you still need to allow side exits here if block_arg0_splat is true.
     // Note that you can't have side exits after this arg0 splat.
     if block_arg0_splat {
-        let arg0_type = ctx.get_opnd_type(StackOpnd(0));
         let arg0_opnd = ctx.stack_opnd(0);
+        let arg0_type = ctx.get_opnd_type(arg0_opnd.into());
 
         // Only handle the case that you don't need to_ary conversion
         let not_array_exit = counted_exit!(ocb, side_exit, invokeblock_iseq_arg0_not_array);
@@ -6123,7 +6122,7 @@ fn gen_send_general(
 
     // Points to the receiver operand on the stack
     let recv = ctx.stack_opnd(recv_idx);
-    let recv_opnd = StackOpnd(recv_idx.try_into().unwrap());
+    let recv_opnd: YARVOpnd = recv.into();
 
     // Log the name of the method we're calling to
     #[cfg(feature = "disasm")]
@@ -6387,14 +6386,15 @@ fn gen_send_general(
                             }
                         };
 
+                        let name_opnd = ctx.stack_opnd(argc);
                         jit_guard_known_klass(
                             jit,
                             ctx,
                             asm,
                             ocb,
                             known_class,
-                            ctx.stack_opnd(argc),
-                            StackOpnd(argc as u16),
+                            name_opnd,
+                            name_opnd.into(),
                             compile_time_name,
                             2, // We have string or symbol, so max depth is 2
                             type_mismatch_exit
@@ -6402,7 +6402,7 @@ fn gen_send_general(
 
                         // Need to do this here so we don't have too many live
                         // values for the register allocator.
-                        let name_opnd = asm.load(ctx.stack_opnd(argc));
+                        let name_opnd = asm.load(name_opnd);
 
                         let symbol_id_opnd = asm.ccall(rb_get_symbol_id as *const u8, vec![name_opnd]);
 
@@ -7016,7 +7016,7 @@ fn gen_objtostring(
             ocb,
             comptime_recv.class_of(),
             recv,
-            StackOpnd(0),
+            recv.into(),
             comptime_recv,
             SEND_MAX_DEPTH,
             side_exit,
diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index 1d64a62964..cb6a328ee3 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -294,6 +294,15 @@ pub enum YARVOpnd {
     StackOpnd(u16),
 }
 
+impl From<Opnd> for YARVOpnd {
+    fn from(value: Opnd) -> Self {
+        match value {
+            Opnd::Stack { idx, .. } => StackOpnd(idx as u16),
+            _ => unreachable!("{:?} cannot be converted to YARVOpnd", value)
+        }
+    }
+}
+
 /// Code generation context
 /// Contains information we can use to specialize/optimize code
 /// There are a lot of context objects so we try to keep the size small.
@@ -1176,9 +1185,7 @@ impl Context {
         self.stack_size += 1;
         self.sp_offset += 1;
 
-        // SP points just above the topmost value
-        let offset = ((self.sp_offset as i32) - 1) * (SIZEOF_VALUE as i32);
-        return Opnd::mem(64, SP, offset);
+        return self.stack_opnd(0);
     }
 
     /// Push one new value on the temp stack
@@ -1206,9 +1213,7 @@ impl Context {
     pub fn stack_pop(&mut self, n: usize) -> Opnd {
         assert!(n <= self.stack_size.into());
 
-        // SP points just above the topmost value
-        let offset = ((self.sp_offset as i32) - 1) * (SIZEOF_VALUE as i32);
-        let top = Opnd::mem(64, SP, offset);
+        let top = self.stack_opnd(0);
 
         // Clear the types of the popped values
         for i in 0..n {
@@ -1243,10 +1248,7 @@ impl Context {
 
     /// Get an operand pointing to a slot on the temp stack
     pub fn stack_opnd(&self, idx: i32) -> Opnd {
-        // SP points just above the topmost value
-        let offset = ((self.sp_offset as i32) - 1 - idx) * (SIZEOF_VALUE as i32);
-        let opnd = Opnd::mem(64, SP, offset);
-        return opnd;
+        Opnd::Stack { idx, sp_offset: self.sp_offset, num_bits: 64 }
     }
 
     /// Get the type of an instruction operand