Explicit handling of frozen strings in IO::Buffer#for. (#5892)

Merged-By: ioquatix <samuel@codeotaku.com>

io_buffer.c

@@ -208,9 +208,11 @@ io_buffer_free(struct rb_io_buffer *data)
             io_buffer_unmap(data->base, data->size);
         }
 
-        if (RB_TYPE_P(data->source, T_STRING)) {
-            rb_str_unlocktmp(data->source);
-        }
+        // Previously we had this, but we found out due to the way GC works, we
+        // can't refer to any other Ruby objects here.
+        // if (RB_TYPE_P(data->source, T_STRING)) {
+        //     rb_str_unlocktmp(data->source);
+        // }
 
         data->base = NULL;
 
@@ -282,45 +284,14 @@ rb_io_buffer_type_allocate(VALUE self)
     return instance;
 }
 
-/*
- *  call-seq: IO::Buffer.for(string) -> io_buffer
- *
- *  Creates a IO::Buffer from the given string's memory. The buffer remains
- *  associated with the string, and writing to a buffer will update the string's
- *  contents.
- *
- *  Until #free is invoked on the buffer, either explicitly or via the garbage
- *  collector, the source string will be locked and cannot be modified.
- *
- *  If the string is frozen, it will create a read-only buffer which cannot be
- *  modified.
- *
- *    string = 'test'
- *    buffer = IO::Buffer.for(str)
- *    buffer.external? #=> true
- *
- *    buffer.get_string(0, 1)
- *    # => "t"
- *    string
- *    # => "best"
- *
- *    buffer.resize(100)
- *    # in `resize': Cannot resize external buffer! (IO::Buffer::AccessError)
- */
-VALUE
-rb_io_buffer_type_for(VALUE klass, VALUE string)
+static VALUE
+io_buffer_for_make_instance(VALUE klass, VALUE string)
 {
-    io_buffer_experimental();
-
-    StringValue(string);
-
     VALUE instance = rb_io_buffer_type_allocate(klass);
 
     struct rb_io_buffer *data = NULL;
     TypedData_Get_Struct(instance, struct rb_io_buffer, &rb_io_buffer_type, data);
 
-    rb_str_locktmp(string);
-
     enum rb_io_buffer_flags flags = RB_IO_BUFFER_EXTERNAL;
 
     if (RB_OBJ_FROZEN(string))
@@ -331,6 +302,94 @@ rb_io_buffer_type_for(VALUE klass, VALUE string)
     return instance;
 }
 
+struct io_buffer_for_yield_instance_arguments {
+  VALUE klass;
+  VALUE string;
+  VALUE instance;
+};
+
+static VALUE
+io_buffer_for_yield_instance(VALUE _arguments) {
+    struct io_buffer_for_yield_instance_arguments *arguments = (struct io_buffer_for_yield_instance_arguments *)_arguments;
+
+    rb_str_locktmp(arguments->string);
+
+    arguments->instance = io_buffer_for_make_instance(arguments->klass, arguments->string);
+
+    return rb_yield(arguments->instance);
+}
+
+static VALUE
+io_buffer_for_yield_instance_ensure(VALUE _arguments)
+{
+    struct io_buffer_for_yield_instance_arguments *arguments = (struct io_buffer_for_yield_instance_arguments *)_arguments;
+
+    if (arguments->instance != Qnil) {
+        rb_io_buffer_free(arguments->instance);
+    }
+
+    rb_str_unlocktmp(arguments->string);
+
+    return Qnil;
+}
+
+/*
+ *  call-seq:
+ *    IO::Buffer.for(string) -> readonly io_buffer
+ *    IO::Buffer.for(string) {|io_buffer| ... read/write io_buffer ...}
+ *
+ *  Creates a IO::Buffer from the given string's memory. Without a block a
+ *  frozen internal copy of the string is created efficiently and used as the
+ *  buffer source. When a block is provided, the buffer is associated directly
+ *  with the string's internal data and updating the buffer will update the
+ *  string.
+ *
+ *  Until #free is invoked on the buffer, either explicitly or via the garbage
+ *  collector, the source string will be locked and cannot be modified.
+ *
+ *  If the string is frozen, it will create a read-only buffer which cannot be
+ *  modified.
+ *
+ *    string = 'test'
+ *    buffer = IO::Buffer.for(string)
+ *    buffer.external? #=> true
+ *
+ *    buffer.get_string(0, 1)
+ *    # => "t"
+ *    string
+ *    # => "best"
+ *
+ *    buffer.resize(100)
+ *    # in `resize': Cannot resize external buffer! (IO::Buffer::AccessError)
+ *
+ *    IO::Buffer.for(string) do |buffer|
+ *      buffer.set_string("T")
+ *      string
+ *      # => "Test"
+ *    end
+ */
+VALUE
+rb_io_buffer_type_for(VALUE klass, VALUE string)
+{
+    StringValue(string);
+
+    // If the string is frozen, both code paths are okay.
+    // If the string is not frozen, if a block is not given, it must be frozen.
+    if (rb_block_given_p()) {
+      struct io_buffer_for_yield_instance_arguments arguments = {
+          .klass = klass,
+          .string = string,
+          .instance = Qnil,
+      };
+
+      return rb_ensure(io_buffer_for_yield_instance, (VALUE)&arguments, io_buffer_for_yield_instance_ensure, (VALUE)&arguments);
+    } else {
+      // This internally returns the source string if it's already frozen.
+      string = rb_str_tmp_frozen_acquire(string);
+      return io_buffer_for_make_instance(klass, string);
+    }
+}
+
 VALUE
 rb_io_buffer_new(void *base, size_t size, enum rb_io_buffer_flags flags)
 {
@@ -2079,7 +2138,20 @@ io_buffer_pwrite(VALUE self, VALUE io, VALUE length, VALUE offset)
  *  C mechanisms like `memcpy`.
  *
  *  The class is meant to be an utility for implementing more high-level mechanisms
- *  like Fiber::SchedulerInterface#io_read and Fiber::SchedulerInterface#io_write.
+ *  like Fiber::SchedulerInterface#io_read and Fiber::Sc            io_buffer_unmap(data->base, data->size);
+        }
+
+        if (RB_TYPE_P(data->source, T_STRING)) {
+            rb_str_unlocktmp(data->source);
+        }
+        // Previously we had this, but we found out due to the way GC works, we
+        // can't refer to any other Ruby objects here.
+        // if (RB_TYPE_P(data->source, T_STRING)) {
+        //     rb_str_unlocktmp(data->source);
+        // }
+
+        data->base = NULL;
+hedulerInterface#io_write.
  *
  *  <b>Examples of usage:</b>
  *

test/ruby/test_io_buffer.rb

@@ -88,30 +88,34 @@ class TestIOBuffer < Test::Unit::TestCase
   def test_string_mapped
     string = "Hello World"
     buffer = IO::Buffer.for(string)
-    refute buffer.readonly?
-
-    # Cannot modify string as it's locked by the buffer:
-    assert_raise RuntimeError do
-      string[0] = "h"
-    end
-
-    buffer.set_value(:U8, 0, "h".ord)
-
-    # Buffer releases it's ownership of the string:
-    buffer.free
-
-    assert_equal "hello World", string
-    string[0] = "H"
-    assert_equal "Hello World", string
+    assert buffer.readonly?
   end
 
   def test_string_mapped_frozen
     string = "Hello World".freeze
     buffer = IO::Buffer.for(string)
-
     assert buffer.readonly?
   end
 
+  def test_string_mapped_mutable
+    string = "Hello World"
+    IO::Buffer.for(string) do |buffer|
+      refute buffer.readonly?
+
+      # Cannot modify string as it's locked by the buffer:
+      assert_raise RuntimeError do
+        string[0] = "h"
+      end
+
+      buffer.set_value(:U8, 0, "h".ord)
+
+      # Buffer releases it's ownership of the string:
+      buffer.free
+
+      assert_equal "hello World", string
+    end
+  end
+
   def test_non_string
     not_string = Object.new