ZJIT: Compile ISEQ with optional arguments (ruby#14653)

Author: k0kubun

test/ruby/test_zjit.rb

@@ -308,6 +308,39 @@ def test_setlocal_on_eval
     }
   end
 
+  def test_optional_arguments
+    assert_compiles '[[1, 2, 3], [10, 20, 3], [100, 200, 300]]', %q{
+      def test(a, b = 2, c = 3)
+        [a, b, c]
+      end
+      [test(1), test(10, 20), test(100, 200, 300)]
+    }
+  end
+
+  def test_optional_arguments_setlocal
+    assert_compiles '[[2, 2], [1, nil]]', %q{
+      def test(a = (b = 2))
+        [a, b]
+      end
+      [test, test(1)]
+    }
+  end
+
+  def test_optional_arguments_cyclic
+    assert_compiles '[nil, 1]', %q{
+      test = proc { |a=a| a }
+      [test.call, test.call(1)]
+    }
+  end
+
+  def test_optional_arguments_side_exit
+    # This leads to FailedOptionalArguments, so not using assert_compiles
+    assert_runs '[:foo, nil, 1]', %q{
+      def test(a = (def foo = nil)) = a
+      [test, (undef :foo), test(1)]
+    }
+  end
+
   def test_getblockparamproxy
     assert_compiles '1', %q{
       def test(&block)

zjit/src/codegen.rs

@@ -16,7 +16,7 @@ use crate::stats::{exit_counter_for_compile_error, incr_counter, incr_counter_by
 use crate::stats::{counter_ptr, with_time_stat, Counter, send_fallback_counter, Counter::{compile_time_ns, exit_compile_error}};
 use crate::{asm::CodeBlock, cruby::*, options::debug, virtualmem::CodePtr};
 use crate::backend::lir::{self, asm_comment, asm_ccall, Assembler, Opnd, Target, CFP, C_ARG_OPNDS, C_RET_OPND, EC, NATIVE_STACK_PTR, NATIVE_BASE_PTR, SCRATCH_OPND, SP};
-use crate::hir::{iseq_to_hir, Block, BlockId, BranchEdge, Invariant, RangeType, SideExitReason, SideExitReason::*, MethodType, SpecialObjectType, SpecialBackrefSymbol, SELF_PARAM_IDX};
+use crate::hir::{iseq_to_hir, Block, BlockId, BranchEdge, Invariant, MethodType, RangeType, SideExitReason::{self, *}, SpecialBackrefSymbol, SpecialObjectType, SELF_PARAM_IDX};
 use crate::hir::{Const, FrameState, Function, Insn, InsnId};
 use crate::hir_type::{types, Type};
 use crate::options::get_option;
@@ -34,7 +34,7 @@ struct JITState {
     labels: Vec<Option<Target>>,
 
     /// JIT entry point for the `iseq`
-    jit_entry: Option<Rc<RefCell<JITEntry>>>,
+    jit_entries: Vec<Rc<RefCell<JITEntry>>>,
 
     /// ISEQ calls that need to be compiled later
     iseq_calls: Vec<Rc<RefCell<IseqCall>>>,
@@ -50,7 +50,7 @@ impl JITState {
             iseq,
             opnds: vec![None; num_insns],
             labels: vec![None; num_blocks],
-            jit_entry: None,
+            jit_entries: Vec::default(),
             iseq_calls: Vec::default(),
             c_stack_slots,
         }
@@ -228,7 +228,7 @@ fn gen_iseq_body(cb: &mut CodeBlock, iseq: IseqPtr, function: Option<&Function>,
     };
 
     // Compile the High-level IR
-    let (start_ptr, jit_entry_ptr, gc_offsets, iseq_calls) = gen_function(cb, iseq, function)?;
+    let (iseq_code_ptrs, gc_offsets, iseq_calls) = gen_function(cb, iseq, function)?;
 
     // Stub callee ISEQs for JIT-to-JIT calls
     for iseq_call in iseq_calls.iter() {
@@ -238,11 +238,11 @@ fn gen_iseq_body(cb: &mut CodeBlock, iseq: IseqPtr, function: Option<&Function>,
     // Prepare for GC
     payload.iseq_calls.extend(iseq_calls.clone());
     append_gc_offsets(iseq, &gc_offsets);
-    Ok(IseqCodePtrs { start_ptr, jit_entry_ptr })
+    Ok(iseq_code_ptrs)
 }
 
 /// Compile a function
-fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, function: &Function) -> Result<(CodePtr, CodePtr, Vec<CodePtr>, Vec<IseqCallRef>), CompileError> {
+fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, function: &Function) -> Result<(IseqCodePtrs, Vec<CodePtr>, Vec<IseqCallRef>), CompileError> {
     let c_stack_slots = max_num_params(function).saturating_sub(ALLOC_REGS.len());
     let mut jit = JITState::new(iseq, function.num_insns(), function.num_blocks(), c_stack_slots);
     let mut asm = Assembler::new();
@@ -307,8 +307,13 @@ fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, function: &Function) -> Resul
         }
     }
     result.map(|(start_ptr, gc_offsets)| {
-        let jit_entry_ptr = jit.jit_entry.unwrap().borrow().start_addr.get().unwrap();
-        (start_ptr, jit_entry_ptr, gc_offsets, jit.iseq_calls)
+        // Make sure jit_entry_ptrs can be used as a parallel vector to jit_entry_insns()
+        jit.jit_entries.sort_by_key(|jit_entry| jit_entry.borrow().jit_entry_idx);
+
+        let jit_entry_ptrs = jit.jit_entries.iter().map(|jit_entry|
+            jit_entry.borrow().start_addr.get().expect("start_addr should have been set by pos_marker in gen_entry_point")
+        ).collect();
+        (IseqCodePtrs { start_ptr, jit_entry_ptrs }, gc_offsets, jit.iseq_calls)
     })
 }
 
@@ -375,7 +380,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         // TODO remove this check when we have stack args (we can use Time.new to test it)
         Insn::InvokeBuiltin { bf, state, .. } if bf.argc + 2 > (C_ARG_OPNDS.len() as i32) => return Err(*state),
         Insn::InvokeBuiltin { bf, args, state, .. } => gen_invokebuiltin(jit, asm, &function.frame_state(*state), bf, opnds!(args)),
-        &Insn::EntryPoint { jit_entry } => no_output!(gen_entry_point(jit, asm, jit_entry)),
+        &Insn::EntryPoint { jit_entry_idx } => no_output!(gen_entry_point(jit, asm, jit_entry_idx)),
         Insn::Return { val } => no_output!(gen_return(asm, opnd!(val))),
         Insn::FixnumAdd { left, right, state } => gen_fixnum_add(jit, asm, opnd!(left), opnd!(right), &function.frame_state(*state)),
         Insn::FixnumSub { left, right, state } => gen_fixnum_sub(jit, asm, opnd!(left), opnd!(right), &function.frame_state(*state)),
@@ -424,7 +429,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         &Insn::DefinedIvar { self_val, id, pushval, .. } => { gen_defined_ivar(asm, opnd!(self_val), id, pushval) },
         &Insn::ArrayExtend { left, right, state } => { no_output!(gen_array_extend(jit, asm, opnd!(left), opnd!(right), &function.frame_state(state))) },
         &Insn::GuardShape { val, shape, state } => gen_guard_shape(jit, asm, opnd!(val), shape, &function.frame_state(state)),
-        Insn::LoadPC => gen_load_pc(),
+        Insn::LoadPC => gen_load_pc(asm),
         &Insn::LoadIvarEmbedded { self_val, id, index } => gen_load_ivar_embedded(asm, opnd!(self_val), id, index),
         &Insn::LoadIvarExtended { self_val, id, index } => gen_load_ivar_extended(asm, opnd!(self_val), id, index),
         &Insn::ArrayMax { state, .. }
@@ -825,8 +830,8 @@ fn gen_guard_shape(jit: &mut JITState, asm: &mut Assembler, val: Opnd, shape: Sh
     val
 }
 
-fn gen_load_pc() -> Opnd {
-    Opnd::mem(64, CFP, RUBY_OFFSET_CFP_PC)
+fn gen_load_pc(asm: &mut Assembler) -> Opnd {
+    asm.load(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_PC))
 }
 
 fn gen_load_ivar_embedded(asm: &mut Assembler, self_val: Opnd, id: ID, index: u16) -> Opnd {
@@ -1313,12 +1318,11 @@ fn gen_object_alloc_class(asm: &mut Assembler, class: VALUE, state: &FrameState)
     }
 }
 
-/// Compile a frame setup. If is_jit_entry is true, remember the address of it as a JIT entry.
-fn gen_entry_point(jit: &mut JITState, asm: &mut Assembler, is_jit_entry: bool) {
-    if is_jit_entry {
-        assert!(jit.jit_entry.is_none(), "only one jit_entry is expected");
-        let jit_entry = JITEntry::new();
-        jit.jit_entry = Some(jit_entry.clone());
+/// Compile a frame setup. If jit_entry_idx is Some, remember the address of it as a JIT entry.
+fn gen_entry_point(jit: &mut JITState, asm: &mut Assembler, jit_entry_idx: Option<usize>) {
+    if let Some(jit_entry_idx) = jit_entry_idx {
+        let jit_entry = JITEntry::new(jit_entry_idx);
+        jit.jit_entries.push(jit_entry.clone());
         asm.pos_marker(move |code_ptr, _| {
             jit_entry.borrow_mut().start_addr.set(Some(code_ptr));
         });
@@ -1776,21 +1780,14 @@ fn compile_iseq(iseq: IseqPtr) -> Result<Function, CompileError> {
     let mut function = match iseq_to_hir(iseq) {
         Ok(function) => function,
         Err(err) => {
-            let name = crate::cruby::iseq_get_location(iseq, 0);
-            debug!("ZJIT: iseq_to_hir: {err:?}: {name}");
+            debug!("ZJIT: iseq_to_hir: {err:?}: {}", iseq_get_location(iseq, 0));
             return Err(CompileError::ParseError(err));
         }
     };
     if !get_option!(disable_hir_opt) {
         function.optimize();
     }
     function.dump_hir();
-    #[cfg(debug_assertions)]
-    if let Err(err) = function.validate() {
-        debug!("ZJIT: compile_iseq: {err:?}");
-        use crate::hir::ParseError;
-        return Err(CompileError::ParseError(ParseError::Validation(err)));
-    }
     Ok(function)
 }
 
@@ -1928,10 +1925,15 @@ c_callable! {
 /// Compile an ISEQ for a function stub
 fn function_stub_hit_body(cb: &mut CodeBlock, iseq_call: &Rc<RefCell<IseqCall>>) -> Result<CodePtr, CompileError> {
     // Compile the stubbed ISEQ
-    let IseqCodePtrs { jit_entry_ptr, .. } = gen_iseq(cb, iseq_call.borrow().iseq, None).inspect_err(|err| {
+    let IseqCodePtrs { jit_entry_ptrs, .. } = gen_iseq(cb, iseq_call.borrow().iseq, None).inspect_err(|err| {
         debug!("{err:?}: gen_iseq failed: {}", iseq_get_location(iseq_call.borrow().iseq, 0));
     })?;
 
+    // We currently don't support JIT-to-JIT calls for ISEQs with optional arguments.
+    // So we only need to use jit_entry_ptrs[0] for now. TODO: Support optional arguments.
+    assert_eq!(1, jit_entry_ptrs.len());
+    let jit_entry_ptr = jit_entry_ptrs[0];
+
     // Update the stub to call the code pointer
     let code_addr = jit_entry_ptr.raw_ptr(cb);
     let iseq = iseq_call.borrow().iseq;
@@ -2131,14 +2133,17 @@ impl Assembler {
 
 /// Store info about a JIT entry point
 pub struct JITEntry {
+    /// Index that corresponds to jit_entry_insns()
+    jit_entry_idx: usize,
     /// Position where the entry point starts
     start_addr: Cell<Option<CodePtr>>,
 }
 
 impl JITEntry {
     /// Allocate a new JITEntry
-    fn new() -> Rc<RefCell<Self>> {
+    fn new(jit_entry_idx: usize) -> Rc<RefCell<Self>> {
         let jit_entry = JITEntry {
+            jit_entry_idx,
             start_addr: Cell::new(None),
         };
         Rc::new(RefCell::new(jit_entry))

zjit/src/cruby.rs

@@ -1105,9 +1105,14 @@ pub mod test_utils {
         ruby_str_to_rust_string(eval(&inspect))
     }
 
-    /// Get the ISeq of a specified method
+    /// Get IseqPtr for a specified method
     pub fn get_method_iseq(recv: &str, name: &str) -> *const rb_iseq_t {
-        let wrapped_iseq = eval(&format!("RubyVM::InstructionSequence.of({}.method(:{}))", recv, name));
+        get_proc_iseq(&format!("{}.method(:{})", recv, name))
+    }
+
+    /// Get IseqPtr for a specified Proc object
+    pub fn get_proc_iseq(obj: &str) -> *const rb_iseq_t {
+        let wrapped_iseq = eval(&format!("RubyVM::InstructionSequence.of({obj})"));
         unsafe { rb_iseqw_to_iseq(wrapped_iseq) }
     }
 

zjit/src/gc.rs

@@ -40,8 +40,8 @@ impl IseqPayload {
 pub struct IseqCodePtrs {
     /// Entry for the interpreter
     pub start_ptr: CodePtr,
-    /// Entry for JIT-to-JIT calls
-    pub jit_entry_ptr: CodePtr,
+    /// Entries for JIT-to-JIT calls
+    pub jit_entry_ptrs: Vec<CodePtr>,
 }
 
 #[derive(Debug, PartialEq)]