RJIT: Support rest args

lib/ruby_vm/rjit/compiler.rb

@@ -27,7 +27,7 @@ module RubyVM::RJIT
   CFP = :r15
   SP  = :rbx
 
-  # Scratch registers: rax, rcx
+  # Scratch registers: rax, rcx, rdx
 
   # Mark objects in this Array during GC
   GC_REFS = []

lib/ruby_vm/rjit/insn_compiler.rb

@@ -4287,8 +4287,89 @@ module RubyVM::RJIT
       end
 
       if iseq_has_rest
-        asm.incr_counter(:send_iseq_has_rest)
+        # We are going to allocate so setting pc and sp.
-        return CantCompile
+        jit_save_pc(jit, asm) # clobbers rax
+        jit_save_sp(ctx, asm)
+
+        if flags & C::VM_CALL_ARGS_SPLAT != 0
+          non_rest_arg_count = argc - 1
+          # We start by dupping the array because someone else might have
+          # a reference to it.
+          array = ctx.stack_pop(1)
+          asm.mov(C_ARGS[0], array)
+          asm.call(C.rb_ary_dup)
+          array = C_RET
+          if non_rest_arg_count > required_num
+            # If we have more arguments than required, we need to prepend
+            # the items from the stack onto the array.
+            diff = (non_rest_arg_count - required_num)
+
+            # diff is >0 so no need to worry about null pointer
+            asm.comment('load pointer to array elements')
+            offset_magnitude = C.VALUE.size * diff
+            values_opnd = ctx.sp_opnd(-offset_magnitude)
+            values_ptr = :rcx
+            asm.lea(values_ptr, values_opnd)
+
+            asm.comment('prepend stack values to rest array')
+            asm.mov(C_ARGS[0], diff)
+            asm.mov(C_ARGS[1], values_ptr)
+            asm.mov(C_ARGS[2], array)
+            asm.call(C.rb_yjit_rb_ary_unshift_m)
+            ctx.stack_pop(diff)
+
+            stack_ret = ctx.stack_push
+            asm.mov(stack_ret, C_RET)
+            # We now should have the required arguments
+            # and an array of all the rest arguments
+            argc = required_num + 1
+          elsif non_rest_arg_count < required_num
+            # If we have fewer arguments than required, we need to take some
+            # from the array and move them to the stack.
+            diff = (required_num - non_rest_arg_count)
+            # This moves the arguments onto the stack. But it doesn't modify the array.
+            move_rest_args_to_stack(array, diff, jit, ctx, asm)
+
+            # We will now slice the array to give us a new array of the correct size
+            asm.mov(C_ARGS[0], array)
+            asm.mov(C_ARGS[1], diff)
+            asm.call(C.rjit_rb_ary_subseq_length)
+            stack_ret = ctx.stack_push
+            asm.mov(stack_ret, C_RET)
+
+            # We now should have the required arguments
+            # and an array of all the rest arguments
+            argc = required_num + 1
+          else
+            # The arguments are equal so we can just push to the stack
+            assert_equal(non_rest_arg_count, required_num)
+            stack_ret = ctx.stack_push
+            asm.mov(stack_ret, array)
+          end
+        else
+          assert_equal(true, argc >= required_num)
+          n = (argc - required_num)
+          argc = required_num + 1
+          # If n is 0, then elts is never going to be read, so we can just pass null
+          if n == 0
+            values_ptr = 0
+          else
+            asm.comment('load pointer to array elements')
+            offset_magnitude = C.VALUE.size * n
+            values_opnd = ctx.sp_opnd(-offset_magnitude)
+            values_ptr = :rcx
+            asm.lea(values_ptr, values_opnd)
+          end
+
+          asm.mov(C_ARGS[0], EC)
+          asm.mov(C_ARGS[1], n)
+          asm.mov(C_ARGS[2], values_ptr)
+          asm.call(C.rb_ec_ary_new_from_values)
+
+          ctx.stack_pop(n)
+          stack_ret = ctx.stack_push
+          asm.mov(stack_ret, C_RET)
+        end
       end
 
       if doing_kw_call
@@ -5017,6 +5098,49 @@ module RubyVM::RJIT
       end
     end
 
+    # Pushes arguments from an array to the stack. Differs from push splat because
+    # the array can have items left over.
+    # @param jit [RubyVM::RJIT::JITState]
+    # @param ctx [RubyVM::RJIT::Context]
+    # @param asm [RubyVM::RJIT::Assembler]
+    def move_rest_args_to_stack(array, num_args, jit, ctx, asm)
+      side_exit = side_exit(jit, ctx)
+
+      asm.comment('move_rest_args_to_stack')
+
+      # array is :rax
+      array_len_opnd = :rcx
+      jit_array_len(asm, array, array_len_opnd)
+
+      asm.comment('Side exit if length is less than required')
+      asm.cmp(array_len_opnd, num_args)
+      asm.jl(counted_exit(side_exit, :send_iseq_has_rest_and_splat_not_equal))
+
+      asm.comment('Push arguments from array')
+
+      # Load the address of the embedded array
+      # (struct RArray *)(obj)->as.ary
+      array_reg = array
+
+      # Conditionally load the address of the heap array
+      # (struct RArray *)(obj)->as.heap.ptr
+      flags_opnd = [array_reg, C.RBasic.offsetof(:flags)]
+      asm.test(flags_opnd, C::RARRAY_EMBED_FLAG)
+      heap_ptr_opnd = [array_reg, C.RArray.offsetof(:as, :heap, :ptr)]
+      # Load the address of the embedded array
+      # (struct RArray *)(obj)->as.ary
+      ary_opnd = :rdx # NOTE: array :rax is used after move_rest_args_to_stack too
+      asm.lea(:rcx, [array_reg, C.RArray.offsetof(:as, :ary)])
+      asm.mov(ary_opnd, heap_ptr_opnd)
+      asm.cmovnz(ary_opnd, :rcx)
+
+      num_args.times do |i|
+        top = ctx.stack_push
+        asm.mov(:rcx, [ary_opnd, i * C.VALUE.size])
+        asm.mov(top, :rcx)
+      end
+    end
+
     # vm_caller_setup_arg_splat (+ CALLER_SETUP_ARG):
     # Pushes arguments from an array to the stack that are passed with a splat (i.e. *args).
     # It optimistically compiles to a static size that is the exact number of arguments needed for the function.

rjit_c.c

@@ -176,6 +176,13 @@ rjit_str_simple_append(VALUE str1, VALUE str2)
     return rb_str_cat(str1, RSTRING_PTR(str2), RSTRING_LEN(str2));
 }
 
+VALUE
+rjit_rb_ary_subseq_length(VALUE ary, long beg)
+{
+    long len = RARRAY_LEN(ary);
+    return rb_ary_subseq(ary, beg, len);
+}
+
 // The code we generate in gen_send_cfunc() doesn't fire the c_return TracePoint event
 // like the interpreter. When tracing for c_return is enabled, we patch the code after
 // the C method return to call into this to fire the event.
@@ -507,6 +514,7 @@ extern VALUE rb_str_bytesize(VALUE str);
 extern const rb_callable_method_entry_t *rb_callable_method_entry_or_negative(VALUE klass, ID mid);
 extern VALUE rb_vm_yield_with_cfunc(rb_execution_context_t *ec, const struct rb_captured_block *captured, int argc, const VALUE *argv);
 extern VALUE rb_vm_set_ivar_id(VALUE obj, ID id, VALUE val);
+extern VALUE rb_yjit_rb_ary_unshift_m(int argc, VALUE *argv, VALUE ary);
 
 #include "rjit_c.rbinc"
 

rjit_c.rb

@@ -519,6 +519,10 @@ module RubyVM::RJIT # :nodoc: all
     Primitive.cexpr! %q{ SIZET2NUM((size_t)rb_ary_clear) }
   end
 
+  def C.rb_ary_dup
+    Primitive.cexpr! %q{ SIZET2NUM((size_t)rb_ary_dup) }
+  end
+
   def C.rb_ary_entry_internal
     Primitive.cexpr! %q{ SIZET2NUM((size_t)rb_ary_entry_internal) }
   end
@@ -727,6 +731,10 @@ module RubyVM::RJIT # :nodoc: all
     Primitive.cexpr! %q{ SIZET2NUM((size_t)rb_vm_yield_with_cfunc) }
   end
 
+  def C.rb_yjit_rb_ary_unshift_m
+    Primitive.cexpr! %q{ SIZET2NUM((size_t)rb_yjit_rb_ary_unshift_m) }
+  end
+
   def C.rjit_full_cfunc_return
     Primitive.cexpr! %q{ SIZET2NUM((size_t)rjit_full_cfunc_return) }
   end
@@ -735,6 +743,10 @@ module RubyVM::RJIT # :nodoc: all
     Primitive.cexpr! %q{ SIZET2NUM((size_t)rjit_optimized_call) }
   end
 
+  def C.rjit_rb_ary_subseq_length
+    Primitive.cexpr! %q{ SIZET2NUM((size_t)rjit_rb_ary_subseq_length) }
+  end
+
   def C.rjit_record_exit_stack
     Primitive.cexpr! %q{ SIZET2NUM((size_t)rjit_record_exit_stack) }
   end

tool/rjit/bindgen.rb

@@ -563,6 +563,9 @@ generator = BindingGenerator.new(
     rb_str_dup
     rb_vm_yield_with_cfunc
     rb_vm_set_ivar_id
+    rb_ary_dup
+    rjit_rb_ary_subseq_length
+    rb_yjit_rb_ary_unshift_m
   ],
   types: %w[
     CALL_DATA