add opt_str_new

Author: tenderlove

compile.c

@@ -9656,6 +9656,27 @@ 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(..., capacity: <n>): emit opt_string_new so the
+    // allocation can be pre-sized before initialize runs. 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.
+    int string_new_capa_loc = -1;
+    if (inline_new && keywords != NULL &&
+            !(flag & (VM_CALL_ARGS_SPLAT | VM_CALL_KW_SPLAT | VM_CALL_FORWARDING))) {
+        const NODE *recv_node = get_nd_recv(node);
+        if (recv_node && nd_type_p(recv_node, NODE_CONST) &&
+                RNODE_CONST(recv_node)->nd_vid == rb_intern("String")) {
+            VALUE capacity_sym = ID2SYM(rb_intern("capacity"));
+            for (int i = 0; i < keywords->keyword_len; i++) {
+                if (keywords->keywords[i] == capacity_sym) {
+                    string_new_capa_loc = keywords->keyword_len - 1 - i;
+                    break;
+                }
+            }
+        }
+    }
+
     if (inline_new) {
         // Jump unless the receiver uses the "basic" implementation of "new"
         VALUE ci;
@@ -9665,7 +9686,12 @@ 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);
         }
-        ADD_INSN2(ret, node, opt_new, ci, not_basic_new);
+        if (string_new_capa_loc >= 0) {
+            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);
+        }
         LABEL_REF(not_basic_new);
 
         // optimized path

insns.def

@@ -941,6 +941,37 @@ opt_new
     }
 }
 
+/* Specialized String.new: if "new" is still the builtin and the capacity
+ * argument is a fixnum, allocate a right-sized String and use it as the
+ * receiver for the following initialize. Otherwise jump to the generic new.
+ * capa_loc is the TOPN index of the capacity argument on the stack. */
+DEFINE_INSN
+opt_string_new
+(CALL_DATA cd, OFFSET dst, rb_num_t capa_loc)
+()
+()
+// attr bool leaf = false;
+{
+    VALUE argc = vm_ci_argc(cd->ci);
+    VALUE recv = TOPN(argc);
+
+    // The bookkeeping slot should be empty.
+    RUBY_ASSERT(TOPN(argc + 1) == Qnil);
+
+    VALUE capa = TOPN(capa_loc);
+
+    if (recv == rb_cString &&
+            FIXNUM_P(capa) &&
+            vm_method_cfunc_is(GET_CFP(), cd, recv, rb_str_s_new)) {
+        VALUE str = rb_str_new_capa_for_init(FIX2LONG(capa));
+        TOPN(argc) = str;
+        TOPN(argc + 1) = str;
+    }
+    else {
+        JUMP(dst);
+    }
+}
+
 /* Convert object to string using to_s or equivalent. */
 DEFINE_INSN
 objtostring

internal/string.h

@@ -131,6 +131,11 @@ VALUE rb_id_quote_unprintable(ID);
 VALUE rb_sym_proc_call(ID mid, int argc, const VALUE *argv, int kw_splat, VALUE passed_proc);
 VALUE rb_enc_literal_str(const char *ptr, long len, rb_encoding *enc);
 
+/* opt_string_new support: builtin String.new and the right-sized allocation
+ * the instruction hands to String#initialize. */
+VALUE rb_str_s_new(int argc, VALUE *argv, VALUE klass);
+VALUE rb_str_new_capa_for_init(long capa);
+
 struct rb_execution_context_struct;
 VALUE rb_ec_str_resurrect(struct rb_execution_context_struct *ec, VALUE str, bool chilled);
 

prism_compile.c

@@ -3877,6 +3877,28 @@ 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(..., capacity: <n>): emit opt_string_new so the
+    // allocation can be pre-sized before initialize runs. 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.
+    int string_new_capa_loc = -1;
+    if (inline_new && kw_arg != NULL &&
+            !(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")) {
+            VALUE capacity_sym = ID2SYM(rb_intern("capacity"));
+            for (int i = 0; i < kw_arg->keyword_len; i++) {
+                if (kw_arg->keywords[i] == capacity_sym) {
+                    string_new_capa_loc = kw_arg->keyword_len - 1 - i;
+                    break;
+                }
+            }
+        }
+    }
+
     if (inline_new) {
         if (LAST_ELEMENT(ret) == opt_new_prelude) {
             PUSH_INSN(ret, location, putnil);
@@ -3896,7 +3918,12 @@ 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);
         }
 
-        PUSH_INSN2(ret, location, opt_new, ci, not_basic_new);
+        if (string_new_capa_loc >= 0) {
+            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);
+        }
         LABEL_REF(not_basic_new);
         // optimized path
         PUSH_SEND_R(ret, location, rb_intern("initialize"), INT2FIX(orig_argc), block_iseq, INT2FIX(flags | VM_CALL_FCALL), kw_arg);

string.c

@@ -1736,6 +1736,19 @@ rb_str_buf_new(long capa)
     return str;
 }
 
+/* Allocate a String suitable for becoming the receiver of String#initialize
+ * with a `capacity:` of +capa+. Floors to STR_BUF_MIN_SIZE so that the buffer
+ * rb_str_init computes for the same capacity matches and is reused as-is
+ * (see the "reuse it as-is" branch in rb_str_init). Used by opt_string_new. */
+VALUE
+rb_str_new_capa_for_init(long capa)
+{
+    if (capa < STR_BUF_MIN_SIZE) {
+        capa = STR_BUF_MIN_SIZE;
+    }
+    return rb_str_buf_new(capa);
+}
+
 VALUE
 rb_str_buf_new_cstr(const char *ptr)
 {
@@ -2097,8 +2110,13 @@ rb_str_init(rb_execution_context_t *ec, VALUE str, VALUE orig, VALUE no_str, VAL
                 if (orig == str) n = 0;
             }
             str_modifiable(str);
-            if (STR_EMBED_P(str) || FL_TEST(str, STR_SHARED|STR_NOFREE)) {
-                /* make noembed always */
+            if (!FL_TEST(str, STR_SHARED|STR_NOFREE) && str_capacity(str, termlen) >= (size_t)capa) {
+                /* The receiver already owns a big-enough buffer (e.g. opt_string_new
+                 * pre-sized this object for us). Reuse it as-is, whether it is an
+                 * embedded slot or an owned heap buffer -- do not reallocate. */
+            }
+            else if (STR_EMBED_P(str) || FL_TEST(str, STR_SHARED|STR_NOFREE)) {
+                /* Move to an owned heap buffer of the requested size. */
                 const size_t size = (size_t)capa + termlen;
                 const char *const old_ptr = RSTRING_PTR(str);
                 const size_t osize = RSTRING_LEN(str) + TERM_LEN(str);
@@ -2107,19 +2125,21 @@ rb_str_init(rb_execution_context_t *ec, VALUE str, VALUE orig, VALUE no_str, VAL
                 memcpy(new_ptr, old_ptr, osize < size ? osize : size);
                 FL_UNSET_RAW(str, STR_SHARED|STR_NOFREE);
                 RSTRING(str)->as.heap.ptr = new_ptr;
+                FL_SET(str, STR_NOEMBED);
+                RSTRING(str)->as.heap.aux.capa = capa;
             }
-            else if (STR_HEAP_SIZE(str) != (size_t)capa + termlen) {
+            else {
+                /* Owned heap buffer that is too small: grow it. */
                 SIZED_REALLOC_N(RSTRING(str)->as.heap.ptr, char,
                         (size_t)capa + termlen, STR_HEAP_SIZE(str));
+                RSTRING(str)->as.heap.aux.capa = capa;
             }
             STR_SET_LEN(str, len);
-            TERM_FILL(&RSTRING(str)->as.heap.ptr[len], termlen);
+            TERM_FILL(RSTRING_PTR(str) + len, termlen);
             if (n == 1) {
-                memcpy(RSTRING(str)->as.heap.ptr, RSTRING_PTR(orig), len);
+                memcpy(RSTRING_PTR(str), RSTRING_PTR(orig), len);
                 rb_enc_cr_str_exact_copy(str, orig);
             }
-            FL_SET(str, STR_NOEMBED);
-            RSTRING(str)->as.heap.aux.capa = capa;
         }
         else if (n == 1) {
             rb_str_replace(str, orig);
@@ -2136,7 +2156,7 @@ rb_str_init(rb_execution_context_t *ec, VALUE str, VALUE orig, VALUE no_str, VAL
 }
 
 /* :nodoc: */
-static VALUE
+VALUE
 rb_str_s_new(int argc, VALUE *argv, VALUE klass)
 {
     if (klass != rb_cString) {

string.rb

@@ -23,6 +23,11 @@ def ascii_only?
     Primitive.cexpr! 'rb_str_is_ascii_only_p(self)'
   end
 
+  #  call-seq:
+  #    String.new(string = ''.encode(Encoding::ASCII_8BIT), **options) -> new_string
+  #
+  #  :include: doc/string/new.rdoc
+  #
   def initialize(orig = (no_str = true; nil),
                 encoding: (no_encoding = true; nil),
                 capacity: (no_capacity = true; nil))

vm.c

@@ -27,6 +27,7 @@
 #include "internal/re.h"
 #include "internal/ruby_parser.h"
 #include "internal/st.h"
+#include "internal/string.h"
 #include "internal/symbol.h"
 #include "internal/thread.h"
 #include "internal/transcode.h"