RJIT: Optimize String#+@

3 files changed, 43 insertions, 1 deletions

diff --git a/lib/ruby_vm/rjit/insn_compiler.rb b/lib/ruby_vm/rjit/insn_compiler.rb
index 336dd9cea9..62b4f5c58f 100644
--- a/lib/ruby_vm/rjit/insn_compiler.rb
+++ b/lib/ruby_vm/rjit/insn_compiler.rb
@@ -2855,6 +2855,41 @@ module RubyVM::RJIT
     # @param jit [RubyVM::RJIT::JITState]
     # @param ctx [RubyVM::RJIT::Context]
     # @param asm [RubyVM::RJIT::Assembler]
+    def jit_rb_str_uplus(jit, ctx, asm, argc, _known_recv_class)
+      if argc != 0
+        return false
+      end
+
+      # We allocate when we dup the string
+      jit_prepare_routine_call(jit, ctx, asm)
+
+      asm.comment('Unary plus on string')
+      asm.mov(:rax, ctx.stack_pop(1)) # recv_opnd
+      asm.mov(:rcx, [:rax, C.RBasic.offsetof(:flags)]) # flags_opnd
+      asm.test(:rcx, C::RUBY_FL_FREEZE)
+
+      ret_label = asm.new_label('stack_ret')
+
+      # String#+@ can only exist on T_STRING
+      stack_ret = ctx.stack_push
+
+      # If the string isn't frozen, we just return it.
+      asm.mov(stack_ret, :rax) # recv_opnd
+      asm.jz(ret_label)
+
+      # Str is frozen - duplicate it
+      asm.mov(C_ARGS[0], :rax) # recv_opnd
+      asm.call(C.rb_str_dup)
+      asm.mov(stack_ret, C_RET)
+
+      asm.write_label(ret_label)
+
+      true
+    end
+
+    # @param jit [RubyVM::RJIT::JITState]
+    # @param ctx [RubyVM::RJIT::Context]
+    # @param asm [RubyVM::RJIT::Assembler]
     def jit_rb_str_getbyte(jit, ctx, asm, argc, _known_recv_class)
       return false if argc != 1
       asm.comment('rb_str_getbyte')
@@ -2940,7 +2975,7 @@ module RubyVM::RJIT
       register_cfunc_method(String, :to_str, :jit_rb_str_to_s)
       register_cfunc_method(String, :bytesize, :jit_rb_str_bytesize)
       register_cfunc_method(String, :<<, :jit_rb_str_concat)
-      #register_cfunc_method(String, :+@, :jit_rb_str_uplus)
+      register_cfunc_method(String, :+@, :jit_rb_str_uplus)
 
       # rb_ary_empty_p() method in array.c
       #register_cfunc_method(Array, :empty?, :jit_rb_ary_empty_p)
diff --git a/rjit_c.rb b/rjit_c.rb
index fcc510adcb..0c6cfd3944 100644
--- a/rjit_c.rb
+++ b/rjit_c.rb
@@ -371,6 +371,7 @@ module RubyVM::RJIT # :nodoc: all
   C::RUBY_FIXNUM_FLAG = Primitive.cexpr! %q{ SIZET2NUM(RUBY_FIXNUM_FLAG) }
   C::RUBY_FLONUM_FLAG = Primitive.cexpr! %q{ SIZET2NUM(RUBY_FLONUM_FLAG) }
   C::RUBY_FLONUM_MASK = Primitive.cexpr! %q{ SIZET2NUM(RUBY_FLONUM_MASK) }
+  C::RUBY_FL_FREEZE = Primitive.cexpr! %q{ SIZET2NUM(RUBY_FL_FREEZE) }
   C::RUBY_FL_SINGLETON = Primitive.cexpr! %q{ SIZET2NUM(RUBY_FL_SINGLETON) }
   C::RUBY_IMMEDIATE_MASK = Primitive.cexpr! %q{ SIZET2NUM(RUBY_IMMEDIATE_MASK) }
   C::RUBY_SPECIAL_SHIFT = Primitive.cexpr! %q{ SIZET2NUM(RUBY_SPECIAL_SHIFT) }
@@ -616,6 +617,10 @@ module RubyVM::RJIT # :nodoc: all
     Primitive.cexpr! %q{ SIZET2NUM((size_t)rb_str_concat_literals) }
   end
 
+  def C.rb_str_dup
+    Primitive.cexpr! %q{ SIZET2NUM((size_t)rb_str_dup) }
+  end
+
   def C.rb_str_eql_internal
     Primitive.cexpr! %q{ SIZET2NUM((size_t)rb_str_eql_internal) }
   end
diff --git a/tool/rjit/bindgen.rb b/tool/rjit/bindgen.rb
index 4cca493bcf..bbf174d6fa 100755
--- a/tool/rjit/bindgen.rb
+++ b/tool/rjit/bindgen.rb
@@ -476,6 +476,7 @@ generator = BindingGenerator.new(
       VM_METHOD_TYPE_ZSUPER
       VM_SPECIAL_OBJECT_VMCORE
       RUBY_ENCODING_MASK
+      RUBY_FL_FREEZE
     ],
   },
   values: {
@@ -549,6 +550,7 @@ generator = BindingGenerator.new(
     rb_str_bytesize
     rjit_str_simple_append
     rb_str_buf_append
+    rb_str_dup
   ],
   types: %w[
     CALL_DATA