Add `IO::Buffer.string` for efficient string creation. (#7364)

2 files changed, 44 insertions, 0 deletions

diff --git a/io_buffer.c b/io_buffer.c
index 6a6bc25628..513eddc954 100644
--- a/io_buffer.c
+++ b/io_buffer.c
@@ -407,6 +407,35 @@ rb_io_buffer_type_for(VALUE klass, VALUE string)
     }
 }
 
+/*
+ * call-seq:
+ *   IO::Buffer.string(length) {|io_buffer| ... read/write io_buffer ...} -> string
+ *
+ * Creates a new string of the given length and yields a IO::Buffer instance
+ * to the block which uses the string as a source. The block is expected to
+ * write to the buffer and the string will be returned.
+ *
+ *    IO::Buffer.string(4) do |buffer|
+ *      buffer.set_string("Ruby")
+ *    end
+ *    # => "Ruby"
+ */
+VALUE
+rb_io_buffer_type_string(VALUE klass, VALUE length)
+{
+    VALUE string = rb_str_new(NULL, NUM2SIZET(length));
+
+    struct io_buffer_for_yield_instance_arguments arguments = {
+        .klass = klass,
+        .string = string,
+        .instance = Qnil,
+    };
+
+    rb_ensure(io_buffer_for_yield_instance, (VALUE)&arguments, io_buffer_for_yield_instance_ensure, (VALUE)&arguments);
+
+    return string;
+}
+
 VALUE
 rb_io_buffer_new(void *base, size_t size, enum rb_io_buffer_flags flags)
 {
@@ -3236,6 +3265,7 @@ Init_IO_Buffer(void)
 
     rb_define_alloc_func(rb_cIOBuffer, rb_io_buffer_type_allocate);
     rb_define_singleton_method(rb_cIOBuffer, "for", rb_io_buffer_type_for, 1);
+    rb_define_singleton_method(rb_cIOBuffer, "string", rb_io_buffer_type_string, 1);
 
 #ifdef _WIN32
     SYSTEM_INFO info;
diff --git a/test/ruby/test_io_buffer.rb b/test/ruby/test_io_buffer.rb
index d9fa22fd2a..55e347471d 100644
--- a/test/ruby/test_io_buffer.rb
+++ b/test/ruby/test_io_buffer.rb
@@ -124,6 +124,20 @@ class TestIOBuffer < Test::Unit::TestCase
     end
   end
 
+  def test_string
+    result = IO::Buffer.string(12) do |buffer|
+      buffer.set_string("Hello World!")
+    end
+
+    assert_equal "Hello World!", result
+  end
+
+  def test_string_negative
+    assert_raise ArgumentError do
+      IO::Buffer.string(-1)
+    end
+  end
+
   def test_resize_mapped
     buffer = IO::Buffer.new