zjit implementation

Author: tenderlove

test/ruby/test_zjit.rb

@@ -187,6 +187,44 @@ def test
     }, insns: [:opt_new], call_threshold: 2
   end
 
+  def test_opt_string_new
+    assert_compiles '["", String, false]', %q{
+      def test = String.new
+      s = test; test
+      [s, s.class, s.frozen?]
+    }, insns: [:opt_string_new], call_threshold: 2
+  end
+
+  def test_opt_string_new_with_capacity
+    assert_compiles '["", "", ""]', %q{
+      def test(cap) = String.new(capacity: cap)
+      [test(100), test(-5), test(0)].tap { test(100) }
+    }, insns: [:opt_string_new], call_threshold: 2
+  end
+
+  def test_opt_string_new_allocates_mutable_string
+    assert_compiles '"hello world"', %q{
+      def test = String.new
+      test; test
+      s = test
+      s << "hello" << " " << "world"
+      s
+    }, insns: [:opt_string_new], call_threshold: 2
+  end
+
+  def test_opt_string_new_falls_back_when_initialize_redefined
+    assert_compiles '[true, true]', %q{
+      def test = String.new
+      test; test
+      class String
+        def initialize(*) = (super; @marked = true)
+      end
+      a = test
+      b = String.new
+      [a.instance_variable_get(:@marked), b.instance_variable_get(:@marked)]
+    }, insns: [:opt_string_new], call_threshold: 2
+  end
+
   def test_uncached_getconstant_path
     assert_compiles RUBY_COPYRIGHT.dump, %q{
       def test = RUBY_COPYRIGHT

yjit/bindgen/src/main.rs

@@ -67,6 +67,8 @@ fn main() {
 
         .allowlist_type("ruby_special_consts")
         .allowlist_function("rb_utf8_str_new")
+        .allowlist_function("rb_str_buf_new")
+        .allowlist_function("rb_str_new_capa_for_init")
         .allowlist_function("rb_str_buf_append")
         .allowlist_function("rb_str_dup")
         .allowlist_type("ruby_preserved_encindex")

yjit/src/codegen.rs

@@ -5004,11 +5004,6 @@ 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)) {
 
-        extern "C" {
-            fn rb_str_buf_new(capa: std::os::raw::c_long) -> VALUE;
-            fn rb_str_new_capa_for_init(capa: std::os::raw::c_long) -> VALUE;
-        }
-
         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.

yjit/src/cruby_bindings.inc.rs

@@ -74,6 +74,8 @@ fn main() {
 
         .allowlist_type("ruby_special_consts")
         .allowlist_function("rb_utf8_str_new")
+        .allowlist_function("rb_str_buf_new")
+        .allowlist_function("rb_str_new_capa_for_init")
         .allowlist_function("rb_str_buf_append")
         .allowlist_function("rb_str_dup")
         .allowlist_function("rb_str_getbyte")

zjit/bindgen/src/main.rs

@@ -625,6 +625,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::ArrayLength { array } => gen_array_length(asm, opnd!(array)),
         Insn::ObjectAlloc { val, state } => gen_object_alloc(jit, asm, opnd!(val), &function.frame_state(*state)),
         &Insn::ObjectAllocClass { class, state } => gen_object_alloc_class(asm, class, &function.frame_state(state)),
+        &Insn::StringNew { capa, cfunc, state } => gen_string_new(asm, opnd!(capa), cfunc, &function.frame_state(state)),
         Insn::StringCopy { val, chilled, state } => gen_string_copy(asm, opnd!(val), *chilled, &function.frame_state(*state)),
         Insn::StringConcat { strings, state } => gen_string_concat(jit, asm, opnds!(strings), &function.frame_state(*state)),
         &Insn::StringGetbyte { string, index } => gen_string_getbyte(asm, opnd!(string), opnd!(index)),
@@ -2311,6 +2312,12 @@ fn gen_object_alloc(jit: &JITState, asm: &mut Assembler, val: lir::Opnd, state:
     asm_ccall!(asm, rb_obj_alloc, val)
 }
 
+fn gen_string_new(asm: &mut Assembler, capa: lir::Opnd, cfunc: *const u8, state: &FrameState) -> lir::Opnd {
+    // rb_str_buf_new / rb_str_new_capa_for_init are leaf allocators that may trigger GC.
+    gen_prepare_leaf_call_with_gc(asm, state);
+    asm.ccall(cfunc, vec![capa])
+}
+
 fn gen_object_alloc_class(asm: &mut Assembler, class: VALUE, state: &FrameState) -> lir::Opnd {
     // Allocating an object for a known class with default allocator is leaf; see doc for
     // `ObjectAllocClass`.

zjit/src/codegen.rs

@@ -1634,6 +1634,7 @@ pub(crate) mod ids {
     def_ids! {
         name: NULL               content: b""
         name: respond_to_missing content: b"respond_to_missing?"
+        name: initialize
         name: eq                 content: b"=="
         name: string_eq          content: b"String#=="
         name: include_p          content: b"include?"