RJIT: Optimize String#+@

ruby · Mar 19, 2023 · cd5a8d0 · cd5a8d0
1 parent 2121282
commit cd5a8d0
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 1 deletion.
diff --git a/lib/ruby_vm/rjit/insn_compiler.rb b/lib/ruby_vm/rjit/insn_compiler.rb
@@ -2852,6 +2852,41 @@ def jit_rb_str_concat(jit, ctx, asm, argc, known_recv_class)
       true
     end
 
+    # @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]
@@ -2940,7 +2975,7 @@ def register_cfunc_codegen_funcs
       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
@@ -371,6 +371,7 @@ def rb_iseqw_to_iseq(iseqw)
   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 @@ def C.rb_str_concat_literals
     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
@@ -476,6 +476,7 @@ def push_target(target)
       VM_METHOD_TYPE_ZSUPER
       VM_SPECIAL_OBJECT_VMCORE
       RUBY_ENCODING_MASK
+      RUBY_FL_FREEZE
     ],
   },
   values: {
@@ -549,6 +550,7 @@ def push_target(target)
     rb_str_bytesize
     rjit_str_simple_append
     rb_str_buf_append
+    rb_str_dup
   ],
   types: %w[
     CALL_DATA