[pull] master from ruby:master

test/ruby/test_zjit.rb

@@ -18,6 +18,15 @@ def test_enabled
     RUBY
   end
 
+  def test_stats_enabled
+    assert_runs 'false', <<~RUBY, stats: false
+      RubyVM::ZJIT.stats_enabled?
+    RUBY
+    assert_runs 'true', <<~RUBY, stats: true
+      RubyVM::ZJIT.stats_enabled?
+    RUBY
+  end
+
   def test_enable_through_env
     child_env = {'RUBY_YJIT_ENABLE' => nil, 'RUBY_ZJIT_ENABLE' => '1'}
     assert_in_out_err([child_env, '-v'], '') do |stdout, stderr|
@@ -1486,10 +1495,13 @@ def test_require_rubygems_with_auto_compact
   end
 
   def test_stats
-    assert_runs 'true', %q{
+    assert_runs '[true, true]', %q{
       def test = 1
       test
-      RubyVM::ZJIT.stats[:zjit_insns_count] > 0
+      [
+        RubyVM::ZJIT.stats[:zjit_insns_count] > 0,
+        RubyVM::ZJIT.stats(:zjit_insns_count) > 0,
+      ]
     }, stats: true
   end
 
@@ -1890,6 +1902,17 @@ def make_range_then_exit(v)
     }, call_threshold: 2
   end
 
+  def test_raise_in_second_argument
+    assert_compiles '{ok: true}', %q{
+      def write(hash, key)
+        hash[key] = raise rescue true
+        hash
+      end
+
+      write({}, :ok)
+    }
+  end
+
   private
 
   # Assert that every method call in `test_script` can be compiled by ZJIT

zjit.c

@@ -358,7 +358,7 @@ rb_zjit_singleton_class_p(VALUE klass)
 
 // Primitives used by zjit.rb. Don't put other functions below, which wouldn't use them.
 VALUE rb_zjit_assert_compiles(rb_execution_context_t *ec, VALUE self);
-VALUE rb_zjit_stats(rb_execution_context_t *ec, VALUE self);
+VALUE rb_zjit_stats(rb_execution_context_t *ec, VALUE self, VALUE target_key);
 VALUE rb_zjit_stats_enabled_p(rb_execution_context_t *ec, VALUE self);
 
 // Preprocessed zjit.rb generated during build

zjit.rb

@@ -19,10 +19,15 @@ def enabled?
     Primitive.cexpr! 'RBOOL(rb_zjit_enabled_p)'
   end
 
+  # Check if `--zjit-stats` is used
+  def stats_enabled?
+    Primitive.rb_zjit_stats_enabled_p
+  end
+
   # Return ZJIT statistics as a Hash
-  def stats
-    stats = Primitive.rb_zjit_stats
-    return nil if stats.nil?
+  def stats(key = nil)
+    stats = Primitive.rb_zjit_stats(key)
+    return stats if stats.nil? || !key.nil?
 
     if stats.key?(:vm_insns_count) && stats.key?(:zjit_insns_count)
       stats[:total_insns_count] = stats[:vm_insns_count] + stats[:zjit_insns_count]

zjit/src/codegen.rs

@@ -353,10 +353,10 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::Jump(branch) => no_output!(gen_jump(jit, asm, branch)),
         Insn::IfTrue { val, target } => no_output!(gen_if_true(jit, asm, opnd!(val), target)),
         Insn::IfFalse { val, target } => no_output!(gen_if_false(jit, asm, opnd!(val), target)),
-        Insn::SendWithoutBlock { cd, state, self_val, args, .. } => gen_send_without_block(jit, asm, *cd, &function.frame_state(*state), opnd!(self_val), opnds!(args)),
+        Insn::SendWithoutBlock { cd, state, .. } => gen_send_without_block(jit, asm, *cd, &function.frame_state(*state)),
         // Give up SendWithoutBlockDirect for 6+ args since asm.ccall() doesn't support it.
-        Insn::SendWithoutBlockDirect { cd, state, self_val, args, .. } if args.len() + 1 > C_ARG_OPNDS.len() => // +1 for self
-            gen_send_without_block(jit, asm, *cd, &function.frame_state(*state), opnd!(self_val), opnds!(args)),
+        Insn::SendWithoutBlockDirect { cd, state, args, .. } if args.len() + 1 > C_ARG_OPNDS.len() => // +1 for self
+            gen_send_without_block(jit, asm, *cd, &function.frame_state(*state)),
         Insn::SendWithoutBlockDirect { cme, iseq, self_val, args, state, .. } => gen_send_without_block_direct(cb, jit, asm, *cme, *iseq, opnd!(self_val), opnds!(args), &function.frame_state(*state)),
         // Ensure we have enough room fit ec, self, and arguments
         // TODO remove this check when we have stack args (we can use Time.new to test it)
@@ -859,26 +859,19 @@ fn gen_send_without_block(
     asm: &mut Assembler,
     cd: *const rb_call_data,
     state: &FrameState,
-    self_val: Opnd,
-    args: Vec<Opnd>,
 ) -> lir::Opnd {
-    gen_spill_locals(jit, asm, state);
-    // Spill the receiver and the arguments onto the stack.
-    // They need to be on the interpreter stack to let the interpreter access them.
-    // TODO: Avoid spilling operands that have been spilled before.
-    // TODO: Despite https://github.com/ruby/ruby/pull/13468, Kokubun thinks this should
-    // spill the whole stack in case it raises an exception. The HIR might need to change
-    // for opt_aref_with, which pushes to the stack in the middle of the instruction.
-    asm_comment!(asm, "spill receiver and arguments");
-    for (idx, &val) in [self_val].iter().chain(args.iter()).enumerate() {
-        // Currently, we don't move the SP register. So it's equal to the base pointer.
-        let stack_opnd = Opnd::mem(64, SP, idx as i32 * SIZEOF_VALUE_I32);
-        asm.mov(stack_opnd, val);
-    }
+    // Note that it's incorrect to use this frame state to side exit because
+    // the state might not be on the boundary of an interpreter instruction.
+    // For example, `opt_aref_with` pushes to the stack and then sends.
+    asm_comment!(asm, "spill frame state");
 
     // Save PC and SP
     gen_save_pc(asm, state);
-    gen_save_sp(asm, 1 + args.len()); // +1 for receiver
+    gen_save_sp(asm, state.stack().len());
+
+    // Spill locals and stack
+    gen_spill_locals(jit, asm, state);
+    gen_spill_stack(jit, asm, state);
 
     asm_comment!(asm, "call #{} with dynamic dispatch", ruby_call_method_name(cd));
     unsafe extern "C" {

zjit/src/hir.rs

@@ -818,7 +818,7 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
                 }
                 Ok(())
             },
-            Insn::Snapshot { state } => write!(f, "Snapshot {}", state),
+            Insn::Snapshot { state } => write!(f, "Snapshot {}", state.print(self.ptr_map)),
             Insn::Defined { op_type, v, .. } => {
                 // op_type (enum defined_type) printing logic from iseq.c.
                 // Not sure why rb_iseq_defined_string() isn't exhaustive.
@@ -2513,11 +2513,10 @@ pub struct FrameState {
     locals: Vec<InsnId>,
 }
 
-impl FrameState {
-    /// Get the opcode for the current instruction
-    pub fn get_opcode(&self) -> i32 {
-        unsafe { rb_iseq_opcode_at_pc(self.iseq, self.pc) }
-    }
+/// Print adaptor for [`FrameState`]. See [`PtrPrintMap`].
+pub struct FrameStatePrinter<'a> {
+    inner: &'a FrameState,
+    ptr_map: &'a PtrPrintMap,
 }
 
 /// Compute the index of a local variable from its slot index
@@ -2624,14 +2623,24 @@ impl FrameState {
         args.extend(self.locals.iter().chain(self.stack.iter()).map(|op| *op));
         args
     }
+
+    /// Get the opcode for the current instruction
+    pub fn get_opcode(&self) -> i32 {
+        unsafe { rb_iseq_opcode_at_pc(self.iseq, self.pc) }
+    }
+
+    pub fn print<'a>(&'a self, ptr_map: &'a PtrPrintMap) -> FrameStatePrinter<'a> {
+        FrameStatePrinter { inner: self, ptr_map }
+    }
 }
 
-impl std::fmt::Display for FrameState {
+impl Display for FrameStatePrinter<'_> {
     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
-        write!(f, "FrameState {{ pc: {:?}, stack: ", self.pc)?;
-        write_vec(f, &self.stack)?;
+        let inner = self.inner;
+        write!(f, "FrameState {{ pc: {:?}, stack: ", self.ptr_map.map_ptr(inner.pc))?;
+        write_vec(f, &inner.stack)?;
         write!(f, ", locals: ")?;
-        write_vec(f, &self.locals)?;
+        write_vec(f, &inner.locals)?;
         write!(f, " }}")
     }
 }
@@ -3195,9 +3204,11 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                     let aref_arg = fun.push_insn(block, Insn::Const { val: Const::Value(get_arg(pc, 0)) });
                     let args = vec![aref_arg];
 
+                    let mut send_state = state.clone();
+                    send_state.stack_push(aref_arg);
+                    let send_state = fun.push_insn(block, Insn::Snapshot { state: send_state });
                     let recv = state.stack_pop()?;
-                    let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
-                    let send = fun.push_insn(block, Insn::SendWithoutBlock { self_val: recv, cd, args, state: exit_id });
+                    let send = fun.push_insn(block, Insn::SendWithoutBlock { self_val: recv, cd, args, state: send_state });
                     state.stack_push(send);
                 }
                 YARVINSN_opt_neq => {
@@ -3920,6 +3931,12 @@ mod tests {
         expected_hir.assert_eq(&actual_hir);
     }
 
+    #[track_caller]
+    pub fn assert_function_hir_with_frame_state(function: Function, expected_hir: Expect) {
+        let actual_hir = format!("{}", FunctionPrinter::with_snapshot(&function));
+        expected_hir.assert_eq(&actual_hir);
+    }
+
     #[track_caller]
     fn assert_compile_fails(method: &str, reason: ParseError) {
         let iseq = crate::cruby::with_rubyvm(|| get_method_iseq("self", method));
@@ -5154,11 +5171,18 @@ mod tests {
         eval("
             def test(a) = a['string lit triggers aref_with']
         ");
-        assert_method_hir("test",  expect![[r#"
+
+        let iseq = crate::cruby::with_rubyvm(|| get_method_iseq("self", "test"));
+        assert!(iseq_contains_opcode(iseq, YARVINSN_opt_aref_with));
+        let function = iseq_to_hir(iseq).unwrap();
+        assert_function_hir_with_frame_state(function, expect![[r#"
             fn test@<compiled>:2:
             bb0(v0:BasicObject, v1:BasicObject):
-              v3:StringExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
+              v2:Any = Snapshot FrameState { pc: 0x1000, stack: [], locals: [v1] }
+              v3:StringExact[VALUE(0x1008)] = Const Value(VALUE(0x1008))
+              v4:Any = Snapshot FrameState { pc: 0x1010, stack: [v1, v3], locals: [v1] }
               v5:BasicObject = SendWithoutBlock v1, :[], v3
+              v6:Any = Snapshot FrameState { pc: 0x1018, stack: [v5], locals: [v1] }
               CheckInterrupts
               Return v5
         "#]]);

zjit/src/stats.rs

@@ -72,29 +72,41 @@ fn incr_counter(counter: Counter, amount: u64) {
 
 /// Return a Hash object that contains ZJIT statistics
 #[unsafe(no_mangle)]
-pub extern "C" fn rb_zjit_stats(_ec: EcPtr, _self: VALUE) -> VALUE {
+pub extern "C" fn rb_zjit_stats(_ec: EcPtr, _self: VALUE, target_key: VALUE) -> VALUE {
     if !zjit_enabled_p() {
         return Qnil;
     }
 
-    fn set_stat(hash: VALUE, key: &str, value: u64) {
-        unsafe { rb_hash_aset(hash, rust_str_to_sym(key), VALUE::fixnum_from_usize(value as usize)); }
+    macro_rules! set_stat {
+        ($hash:ident, $key:expr, $value:expr) => {
+            let key = rust_str_to_sym($key);
+            // Evaluate $value only when it's needed
+            if key == target_key {
+                return VALUE::fixnum_from_usize($value as usize);
+            } else if $hash != Qnil {
+                #[allow(unused_unsafe)]
+                unsafe { rb_hash_aset($hash, key, VALUE::fixnum_from_usize($value as usize)); }
+            }
+        }
     }
 
-    let hash = unsafe { rb_hash_new() };
+    let hash = if target_key.nil_p() {
+        unsafe { rb_hash_new() }
+    } else {
+        Qnil
+    };
     let counters = ZJITState::get_counters();
 
     for &counter in DEFAULT_COUNTERS {
-        let counter_val = unsafe { *counter_ptr(counter) };
-        set_stat(hash, &counter.name(), counter_val);
+        set_stat!(hash, &counter.name(), unsafe { *counter_ptr(counter) });
     }
 
     // Set counters that are enabled when --zjit-stats is enabled
     if get_option!(stats) {
-        set_stat(hash, "zjit_insns_count", counters.zjit_insns_count);
+        set_stat!(hash, "zjit_insns_count", counters.zjit_insns_count);
 
         if unsafe { rb_vm_insns_count } > 0 {
-            set_stat(hash, "vm_insns_count", unsafe { rb_vm_insns_count });
+            set_stat!(hash, "vm_insns_count", unsafe { rb_vm_insns_count });
         }
     }