conditionally call initialize

Author: tenderlove

compile.c

@@ -9656,20 +9656,29 @@ compile_call(rb_iseq_t *iseq, LINK_ANCHOR *const ret, const NODE *const node, co
     LABEL *not_basic_new = NEW_LABEL(nd_line(node));
     LABEL *not_basic_new_finish = NEW_LABEL(nd_line(node));
 
-    // Detect String.new and String.new(..., capacity: <n>): emit opt_string_new
-    // so the allocation can be pre-sized. The instruction guards at runtime that
-    // the receiver is String and "new" is unredefined, so a compile-time name
-    // match is sufficient here. string_new_capa_loc is the TOPN index of the
-    // capacity argument, or -1 when not applicable. When the only arguments are
-    // ones the allocation already satisfies (nothing, or just capacity),
-    // initialize is a no-op and string_new_skip_init lets us omit the call.
+    // Detect String.new with no positional arguments: allocate via the
+    // String-specific opt_string_new instead of the generic opt_new. The
+    // instruction guards at runtime that the receiver is String and "new" is
+    // unredefined, so a compile-time name match is sufficient here.
+    //
+    //   - string_new_capa_loc is the TOPN offset of the capacity: argument, or
+    //     -1 when none is given. It can be deeper than 0 because other keywords
+    //     (e.g. encoding:) may sit above it on the stack.
+    //   - string_new_skip_init means the allocation already satisfies the call
+    //     (nothing, or just capacity), so #initialize is omitted. Otherwise a
+    //     trailing initialize send applies the remaining keywords, like opt_new.
     int string_new_capa_loc = -1;
+    bool string_new = false;
     bool string_new_skip_init = false;
     if (inline_new &&
             !(flag & (VM_CALL_ARGS_SPLAT | VM_CALL_KW_SPLAT | VM_CALL_FORWARDING))) {
         const NODE *recv_node = get_nd_recv(node);
+        // NUM2INT(argc) is the positional count (new_callinfo adds the keywords
+        // separately), so it must be zero for both the bare and capacity forms.
         if (recv_node && nd_type_p(recv_node, NODE_CONST) &&
-                RNODE_CONST(recv_node)->nd_vid == rb_intern("String")) {
+                RNODE_CONST(recv_node)->nd_vid == rb_intern("String") &&
+                NUM2INT(argc) == 0) {
+            string_new = true;
             int keyword_len = keywords ? keywords->keyword_len : 0;
             if (keywords) {
                 VALUE capacity_sym = ID2SYM(rb_intern("capacity"));
@@ -9680,11 +9689,7 @@ compile_call(rb_iseq_t *iseq, LINK_ANCHOR *const ret, const NODE *const node, co
                     }
                 }
             }
-            // NUM2INT(argc) is the positional count (new_callinfo adds the
-            // keywords separately). Skip initialize when there are no
-            // positional args and either no keywords or only capacity.
-            if (NUM2INT(argc) == 0 &&
-                    (keyword_len == 0 || (keyword_len == 1 && string_new_capa_loc >= 0))) {
+            if (keyword_len == 0 || (keyword_len == 1 && string_new_capa_loc >= 0)) {
                 string_new_skip_init = true;
             }
         }
@@ -9699,10 +9704,11 @@ compile_call(rb_iseq_t *iseq, LINK_ANCHOR *const ret, const NODE *const node, co
         else {
             ci = (VALUE)new_callinfo(iseq, mid, NUM2INT(argc), flag, keywords, 0);
         }
-        if (string_new_skip_init) {
-            // opt_string_new allocates the String and skips initialize entirely.
-            // capa_loc is unused for the bare form (argc == 0).
-            ADD_INSN3(ret, node, opt_string_new, ci, not_basic_new, INT2FIX(string_new_capa_loc >= 0 ? string_new_capa_loc : 0));
+        if (string_new) {
+            // opt_string_new allocates the String via the String-specific
+            // allocator (pre-sized when capacity: is present). Whether
+            // #initialize runs is decided by the trailing bytecode below.
+            ADD_INSN3(ret, node, opt_string_new, ci, not_basic_new, INT2FIX(string_new_capa_loc));
         }
         else {
             ADD_INSN2(ret, node, opt_new, ci, not_basic_new);

insns.def

@@ -967,17 +967,20 @@ opt_string_new
             vm_method_cfunc_is(GET_CFP(), cd, recv, rb_class_new_instance_pass_kw) &&
             vm_method_basic_definition_p(recv, idInitialize)) {
 
-        if (argc == 0) {
+        if (capa_loc == (rb_num_t)-1) {
+            // No capacity given: allocate an empty (embedded) buffer. When a
+            // trailing #initialize send follows (e.g. String.new(encoding:)) it
+            // applies the remaining arguments; otherwise this is the result.
             str = rb_str_buf_new(0);
         }
         else {
+            // The capacity is not necessarily the top of the stack: with
+            // String.new(capacity:, encoding:) the encoding sits above it.
             VALUE rbcapa = TOPN(capa_loc);
             if (FIXNUM_P(rbcapa)) {
-                long capa = FIX2LONG(rbcapa);
-                if (capa < 0) {
-                  capa = 0;
-                }
-                str = rb_str_new_capa_for_init(capa);
+                // rb_str_new_capa_for_init floors anything below the minimum
+                // buffer size, including negatives, so no clamping is needed.
+                str = rb_str_new_capa_for_init(FIX2LONG(rbcapa));
             }
         }
 

prism_compile.c

@@ -3877,21 +3877,31 @@ pm_compile_call(rb_iseq_t *iseq, const pm_call_node_t *call_node, LINK_ANCHOR *c
         call_node->block == NULL &&
         (flags & VM_CALL_ARGS_BLOCKARG) == 0;
 
-    // Detect String.new and String.new(..., capacity: <n>): emit opt_string_new
-    // so the allocation can be pre-sized. The instruction guards at runtime that
-    // the receiver is String and "new" is unredefined, so a compile-time name
-    // match is sufficient here. string_new_capa_loc is the TOPN index of the
-    // capacity argument, or -1 when not applicable. When the only arguments are
-    // ones the allocation already satisfies (nothing, or just capacity),
-    // initialize is a no-op and string_new_skip_init lets us omit the call.
+    // Detect String.new with no positional arguments: allocate via the
+    // String-specific opt_string_new instead of the generic opt_new. The
+    // instruction guards at runtime that the receiver is String and "new" is
+    // unredefined, so a compile-time name match is sufficient here.
+    //
+    //   - string_new_capa_loc is the TOPN offset of the capacity: argument, or
+    //     -1 when none is given. It can be deeper than 0 because other keywords
+    //     (e.g. encoding:) may sit above it on the stack.
+    //   - string_new_skip_init means the allocation already satisfies the call
+    //     (nothing, or just capacity), so #initialize is omitted. Otherwise a
+    //     trailing initialize send applies the remaining keywords, exactly like
+    //     opt_new does.
     int string_new_capa_loc = -1;
+    bool string_new = false;
     bool string_new_skip_init = false;
     if (inline_new &&
             !(flags & (VM_CALL_ARGS_SPLAT | VM_CALL_KW_SPLAT | VM_CALL_FORWARDING)) &&
             call_node->receiver != NULL &&
             PM_NODE_TYPE_P(call_node->receiver, PM_CONSTANT_READ_NODE)) {
         const pm_constant_read_node_t *recv = (const pm_constant_read_node_t *) call_node->receiver;
-        if (pm_constant_id_lookup(scope_node, recv->name) == rb_intern("String")) {
+        // orig_argc is the positional count (new_callinfo adds the keywords
+        // separately).
+        if (pm_constant_id_lookup(scope_node, recv->name) == rb_intern("String") &&
+                orig_argc == 0) {
+            string_new = true;
             int keyword_len = kw_arg ? kw_arg->keyword_len : 0;
             if (kw_arg) {
                 VALUE capacity_sym = ID2SYM(rb_intern("capacity"));
@@ -3902,11 +3912,7 @@ pm_compile_call(rb_iseq_t *iseq, const pm_call_node_t *call_node, LINK_ANCHOR *c
                     }
                 }
             }
-            // orig_argc is the positional count (new_callinfo adds the keywords
-            // separately). Skip initialize when there are no positional args and
-            // either no keywords or only capacity.
-            if (orig_argc == 0 &&
-                    (keyword_len == 0 || (keyword_len == 1 && string_new_capa_loc >= 0))) {
+            if (keyword_len == 0 || (keyword_len == 1 && string_new_capa_loc >= 0)) {
                 string_new_skip_init = true;
             }
         }
@@ -3931,10 +3937,11 @@ pm_compile_call(rb_iseq_t *iseq, const pm_call_node_t *call_node, LINK_ANCHOR *c
             ci = (VALUE)new_callinfo(iseq, method_id, orig_argc, flags, kw_arg, 0);
         }
 
-        if (string_new_skip_init) {
-            // opt_string_new allocates the String and skips initialize entirely.
-            // capa_loc is unused for the bare form (orig_argc == 0).
-            PUSH_INSN3(ret, location, opt_string_new, ci, not_basic_new, INT2FIX(string_new_capa_loc >= 0 ? string_new_capa_loc : 0));
+        if (string_new) {
+            // opt_string_new allocates the String via the String-specific
+            // allocator (pre-sized when capacity: is present). Whether
+            // #initialize runs is decided by the trailing bytecode below.
+            PUSH_INSN3(ret, location, opt_string_new, ci, not_basic_new, INT2FIX(string_new_capa_loc));
         }
         else {
             PUSH_INSN2(ret, location, opt_new, ci, not_basic_new);

test/ruby/test_string.rb

@@ -86,6 +86,46 @@ def test_initialize
     end
   end
 
+  def test_initialize_keyword_encoding
+    # These go through the opt_string_new "hybrid" path (allocate via the
+    # String-specific allocator, then run #initialize for the encoding).
+    s = String.new(encoding: "UTF-8")
+    assert_equal(Encoding::UTF_8, s.encoding)
+    assert_equal("", s)
+    assert_not_predicate(s, :frozen?)
+    s << "abc"
+    assert_equal("abc", s)
+
+    s = String.new(capacity: 100, encoding: "UTF-8")
+    assert_equal(Encoding::UTF_8, s.encoding)
+    assert_equal("", s)
+    s << "abc"
+    assert_equal("abc", s)
+
+    # Reversed keyword order (capacity is the top of the stack).
+    s = String.new(encoding: "UTF-8", capacity: 100)
+    assert_equal(Encoding::UTF_8, s.encoding)
+    assert_equal("", s)
+
+    s = String.new(encoding: Encoding::US_ASCII)
+    assert_equal(Encoding::US_ASCII, s.encoding)
+  end
+
+  def test_initialize_keyword_redefined
+    assert_separately([], <<-'RUBY')
+      EnvUtil.suppress_warning do
+        class String
+          def initialize(*args, **kw)
+            @marked = true
+          end
+        end
+      end
+      assert_equal(true, String.new(encoding: "UTF-8").instance_variable_get(:@marked))
+      assert_equal(true, String.new(capacity: 100, encoding: "UTF-8").instance_variable_get(:@marked))
+      assert_equal(true, String.new.instance_variable_get(:@marked))
+    RUBY
+  end
+
   def test_initialize_shared
     S(str = "mystring" * 10).__send__(:initialize, capacity: str.bytesize)
     assert_equal("mystring", str[0, 8])

test/ruby/test_zjit.rb

@@ -225,6 +225,36 @@ def initialize(*) = (super; @marked = true)
     }, insns: [:opt_string_new], call_threshold: 2
   end
 
+  def test_opt_string_new_with_encoding
+    # Hybrid: opt_string_new allocates, then a trailing #initialize applies the encoding.
+    assert_compiles '["UTF-8", true]', %q{
+      def test = String.new(encoding: "UTF-8")
+      test; test
+      s = test
+      [s.encoding.name, s.empty?]
+    }, insns: [:opt_string_new], call_threshold: 2
+  end
+
+  def test_opt_string_new_with_capacity_and_encoding
+    # Hybrid with a pre-sized buffer; capacity is below encoding on the stack.
+    assert_compiles '["UTF-8", "abc"]', %q{
+      def test = String.new(capacity: 100, encoding: "UTF-8")
+      test; test
+      s = test
+      s << "abc"
+      [s.encoding.name, s]
+    }, insns: [:opt_string_new], call_threshold: 2
+  end
+
+  def test_opt_string_new_with_encoding_then_capacity
+    # Keyword order reversed: capacity is the top of the stack (capa_loc 0).
+    assert_compiles '"UTF-8"', %q{
+      def test = String.new(encoding: "UTF-8", capacity: 100)
+      test; test
+      test.encoding.name
+    }, insns: [:opt_string_new], call_threshold: 2
+  end
+
   def test_uncached_getconstant_path
     assert_compiles RUBY_COPYRIGHT.dump, %q{
       def test = RUBY_COPYRIGHT

yjit/src/codegen.rs

@@ -4963,7 +4963,8 @@ fn gen_opt_string_new(
 ) -> Option<CodegenStatus> {
     let cd = jit.get_arg(0).as_ptr();
     let jump_offset = jit.get_arg(1).as_i32();
-    let capa_loc = jit.get_arg(2).as_i32();
+    // TOPN offset of the capacity: argument, or -1 when none was given.
+    let capa_loc = jit.get_arg(2).as_i64();
 
     if !jit.at_compile_target() {
         return jit.defer_compilation(asm);
@@ -5004,22 +5005,26 @@ fn gen_opt_string_new(
     if jit.assume_expected_cfunc(asm, comptime_recv_klass, mid, rb_class_new_instance_pass_kw as _)
         && assume_method_basic_definition(jit, asm, comptime_recv_klass, ID!(initialize)) {
 
-        let str = if argc == 0 {
-            // bare String.new: an empty (embedded) string. rb_str_buf_new skips
-            // the encoding/coderange setup that rb_str_new would do.
+        let str = if capa_loc == -1 {
+            // No capacity given: an empty (embedded) string. rb_str_buf_new skips
+            // the encoding/coderange setup that rb_str_new would do. When a
+            // trailing initialize send follows (e.g. String.new(encoding:)) it
+            // applies the remaining arguments.
             jit_prepare_non_leaf_call(jit, asm);
             asm.ccall(rb_str_buf_new as *const u8, vec![Opnd::Imm(0)])
         } else {
-            // String.new(capacity: n): only fixnum capacities take the fast path.
-            // A non-fixnum side-exits so the interpreter takes the generic path.
-            let capa = asm.stack_opnd(capa_loc);
+            // The capacity is at TOPN(capa_loc) (not necessarily the top of the
+            // stack: with String.new(capacity:, encoding:) the encoding sits
+            // above it). Only fixnum capacities take the fast path; a non-fixnum
+            // side-exits so the interpreter takes the generic path.
+            let capa = asm.stack_opnd(capa_loc as i32);
             guard_object_is_fixnum(jit, asm, capa, capa.into());
 
             jit_prepare_non_leaf_call(jit, asm);
 
             // FIX2LONG; rb_str_new_capa_for_init floors anything below the
             // minimum buffer size (including negatives) so no clamping needed.
-            let capa = asm.load(asm.stack_opnd(capa_loc));
+            let capa = asm.load(asm.stack_opnd(capa_loc as i32));
             let capa = asm.rshift(capa, Opnd::UImm(1));
             asm.ccall(rb_str_new_capa_for_init as *const u8, vec![capa])
         };

zjit/src/hir.rs

@@ -8088,7 +8088,8 @@ fn add_iseq_to_hir(
                     // mirroring the interpreter's opt_string_new fast path.
                     let cd: *const rb_call_data = get_arg(pc, 0).as_ptr();
                     let dst = get_arg(pc, 1).as_i64();
-                    let capa_loc = get_arg(pc, 2).as_u64() as usize;
+                    // TOPN offset of the capacity: argument, or -1 when none was given.
+                    let capa_loc = get_arg(pc, 2).as_i64();
 
                     let argc = crate::profile::num_arguments_on_stack(cd);
                     let ci = unsafe { get_call_data_ci(cd) };
@@ -8140,16 +8141,19 @@ fn add_iseq_to_hir(
                         state: exit_id,
                     });
 
-                    let str = if argc == 0 {
-                        // Bare String.new: rb_str_buf_new(0) yields an embedded empty string,
-                        // skipping the encoding/coderange setup rb_str_new would do.
+                    let str = if capa_loc == -1 {
+                        // No capacity given: rb_str_buf_new(0) yields an embedded empty string,
+                        // skipping the encoding/coderange setup rb_str_new would do. A trailing
+                        // initialize send (e.g. String.new(encoding:)) applies the rest.
                         let capa = fun.push_insn(block, Insn::Const { val: Const::CInt64(0) });
                         fun.push_insn(block, Insn::StringNew { capa, cfunc: rb_str_buf_new as *const u8, state: exit_id })
                     } else {
-                        // String.new(capacity: n): only the Fixnum case is handled here; anything
-                        // else side-exits and the interpreter re-runs opt_string_new (which jumps
-                        // to the generic path for a non-Fixnum capacity).
-                        let capa_val = state.stack_topn(capa_loc)?;
+                        // The capacity is at TOPN(capa_loc), not necessarily the top of the stack
+                        // (with String.new(capacity:, encoding:) the encoding sits above it). Only
+                        // the Fixnum case is handled here; anything else side-exits and the
+                        // interpreter re-runs opt_string_new (which jumps to the generic path for
+                        // a non-Fixnum capacity).
+                        let capa_val = state.stack_topn(capa_loc as usize)?;
                         let capa_fix = fun.push_insn(block, Insn::GuardType { val: capa_val, guard_type: types::Fixnum, state: exit_id, recompile: None });
                         let capa = fun.push_insn(block, Insn::UnboxFixnum { val: capa_fix });
                         fun.push_insn(block, Insn::StringNew { capa, cfunc: rb_str_new_capa_for_init as *const u8, state: exit_id })