Make IO::Buffer MemoryView ready (#17965)

`null?` buffer isn't MemoryView available.

Not `null?` buffer is MemoryView available. It can be exported as a
1-dimensional contiguous array of bytes.

Users can't change the underlying memory of the exporting IO::Buffer
while exported MemoryView is using.

ext/-test-/memory_view/memory_view.c

@@ -114,51 +114,85 @@ memory_view_parse_item_format(VALUE mod, VALUE format)
 }
 
 static VALUE
-memory_view_get_memory_view_info(VALUE mod, VALUE obj)
+memory_view_release_ensure(VALUE data)
 {
-    rb_memory_view_t view;
+    rb_memory_view_t *view = (rb_memory_view_t *)data;
+    rb_memory_view_release(view);
+    return Qnil;
+}
 
-    if (!rb_memory_view_get(obj, &view, 0)) {
-        return Qnil;
-    }
+static VALUE
+memory_view_get_memory_view_info_body(VALUE args)
+{
+    rb_memory_view_t *view = (rb_memory_view_t *)args;
 
     VALUE hash = rb_hash_new();
-    rb_hash_aset(hash, sym_obj, view.obj);
-    rb_hash_aset(hash, sym_byte_size, SSIZET2NUM(view.byte_size));
-    rb_hash_aset(hash, sym_readonly, view.readonly ? Qtrue : Qfalse);
-    rb_hash_aset(hash, sym_format, view.format ? rb_str_new_cstr(view.format) : Qnil);
-    rb_hash_aset(hash, sym_item_size, SSIZET2NUM(view.item_size));
-    rb_hash_aset(hash, sym_ndim, SSIZET2NUM(view.ndim));
+    rb_hash_aset(hash, sym_obj, view->obj);
+    rb_hash_aset(hash, sym_byte_size, SSIZET2NUM(view->byte_size));
+    rb_hash_aset(hash, sym_readonly, view->readonly ? Qtrue : Qfalse);
+    rb_hash_aset(hash, sym_format, view->format ? rb_str_new_cstr(view->format) : Qnil);
+    rb_hash_aset(hash, sym_item_size, SSIZET2NUM(view->item_size));
+    rb_hash_aset(hash, sym_ndim, SSIZET2NUM(view->ndim));
 
-    if (view.shape) {
-        VALUE shape = rb_ary_new_capa(view.ndim);
+    if (view->shape) {
+        VALUE shape = rb_ary_new_capa(view->ndim);
+        ssize_t i;
+        for (i = 0; i < view->ndim; i++) {
+            rb_ary_push(shape, SSIZET2NUM(view->shape[i]));
+        }
         rb_hash_aset(hash, sym_shape, shape);
     }
     else {
         rb_hash_aset(hash, sym_shape, Qnil);
     }
 
-    if (view.strides) {
-        VALUE strides = rb_ary_new_capa(view.ndim);
+    if (view->strides) {
+        VALUE strides = rb_ary_new_capa(view->ndim);
+        ssize_t i;
+        for (i = 0; i < view->ndim; i++) {
+            rb_ary_push(strides, SSIZET2NUM(view->strides[i]));
+        }
         rb_hash_aset(hash, sym_strides, strides);
     }
     else {
         rb_hash_aset(hash, sym_strides, Qnil);
     }
 
-    if (view.sub_offsets) {
-        VALUE sub_offsets = rb_ary_new_capa(view.ndim);
+    if (view->sub_offsets) {
+        VALUE sub_offsets = rb_ary_new_capa(view->ndim);
         rb_hash_aset(hash, sym_sub_offsets, sub_offsets);
     }
     else {
         rb_hash_aset(hash, sym_sub_offsets, Qnil);
     }
 
-    rb_memory_view_release(&view);
-
     return hash;
 }
 
+static VALUE
+memory_view_get_memory_view_info(int argc, VALUE *args, VALUE mod)
+{
+    rb_memory_view_t view;
+    VALUE obj;
+    VALUE flags;
+    enum ruby_memory_view_flags view_flags;
+
+    rb_scan_args(argc, args, "11", &obj, &flags);
+    if (NIL_P(flags)) {
+        view_flags = RUBY_MEMORY_VIEW_SIMPLE;
+    }
+    else {
+        view_flags = NUM2UINT(flags);
+    }
+
+    if (!rb_memory_view_get(obj, &view, view_flags)) {
+        return Qnil;
+    }
+
+    return rb_ensure(memory_view_get_memory_view_info_body, (VALUE)&view,
+                     memory_view_release_ensure, (VALUE)&view);
+}
+
 static VALUE
 memory_view_is_row_major_contiguous(VALUE mod, VALUE obj)
 {
@@ -281,6 +315,108 @@ memory_view_extract_item_members(VALUE mod, VALUE str, VALUE format)
     return item;
 }
 
+typedef struct {
+    VALUE location;
+    rb_memory_view_t view;
+} memory_view_get_data_args;
+
+static VALUE
+memory_view_get_data_body(VALUE _args)
+{
+    memory_view_get_data_args *args = (memory_view_get_data_args *)_args;
+    long beg, len;
+
+    if (RTEST(rb_range_beg_len(args->location,
+                               &beg,
+                               &len,
+                               args->view.byte_size,
+                               0))) {
+        const char *content = ((const char *)(args->view.data)) + beg;
+        return rb_str_new(content, len);
+    }
+    else {
+        const char *content =
+            ((const char *)(args->view.data)) + NUM2LONG(args->location);
+        return rb_str_new(content, 1);
+    }
+}
+
+static VALUE
+memory_view_get_data(VALUE mod, VALUE obj, VALUE location)
+{
+    memory_view_get_data_args args;
+    args.location = location;
+
+    if (!rb_memory_view_get(obj, &(args.view), 0)) {
+        return Qnil;
+    }
+
+    return rb_ensure(memory_view_get_data_body, (VALUE)&args,
+                     memory_view_release_ensure, (VALUE)&(args.view));
+}
+
+typedef struct {
+    VALUE offset;
+    VALUE data;
+    rb_memory_view_t view;
+} memory_view_set_data_args;
+
+static VALUE
+memory_view_set_data_body(VALUE _args)
+{
+    memory_view_set_data_args *args = (memory_view_set_data_args *)_args;
+
+    if (FIXNUM_P(args->data)) {
+        ((char *)(args->view.data))[NUM2LONG(args->offset)] =
+            NUM2LONG(args->data);
+    }
+    else {
+        StringValue(args->data);
+        memcpy(((char *)(args->view.data)) + NUM2LONG(args->offset),
+               RSTRING_PTR(args->data),
+               RSTRING_LEN(args->data));
+    }
+
+    return Qtrue;
+}
+
+static VALUE
+memory_view_set_data(VALUE mod, VALUE obj, VALUE offset, VALUE data)
+{
+    memory_view_set_data_args args;
+    args.offset = offset;
+    args.data = data;
+
+    if (!rb_memory_view_get(obj, &(args.view), RUBY_MEMORY_VIEW_WRITABLE)) {
+        return Qnil;
+    }
+
+    return rb_ensure(memory_view_set_data_body, (VALUE)&args,
+                     memory_view_release_ensure, (VALUE)&(args.view));
+}
+
+static VALUE
+memory_view_get_body(VALUE data)
+{
+    return rb_yield_values(0);
+}
+
+static VALUE
+memory_view_get(VALUE mod, VALUE obj, VALUE flags)
+{
+    rb_memory_view_t view;
+
+    if (!rb_memory_view_get(obj, &view, NUM2UINT(flags))) {
+        rb_raise(rb_eArgError,
+                 "Unable to get memory view: "
+                 "object=%+" PRIsVALUE " flags=%+" PRIsVALUE,
+                 obj, flags);
+    }
+
+    return rb_ensure(memory_view_get_body, Qnil,
+                     memory_view_release_ensure, (VALUE)&view);
+}
+
 static VALUE
 expstr_initialize(VALUE obj, VALUE s)
 {
@@ -416,16 +552,38 @@ Init_memory_view(void)
 #ifdef HAVE_RUBY_MEMORY_VIEW_H
     VALUE mMemoryViewTestUtils = rb_define_module("MemoryViewTestUtils");
 
+    rb_define_const(mMemoryViewTestUtils, "SIMPLE",
+                    UINT2NUM(RUBY_MEMORY_VIEW_SIMPLE));
+    rb_define_const(mMemoryViewTestUtils, "WRITABLE",
+                    UINT2NUM(RUBY_MEMORY_VIEW_WRITABLE));
+    rb_define_const(mMemoryViewTestUtils, "FORMAT",
+                    UINT2NUM(RUBY_MEMORY_VIEW_FORMAT));
+    rb_define_const(mMemoryViewTestUtils, "MULTI_DIMENSIONAL",
+                    UINT2NUM(RUBY_MEMORY_VIEW_MULTI_DIMENSIONAL));
+    rb_define_const(mMemoryViewTestUtils, "STRIDES",
+                    UINT2NUM(RUBY_MEMORY_VIEW_STRIDES));
+    rb_define_const(mMemoryViewTestUtils, "ROW_MAJOR",
+                    UINT2NUM(RUBY_MEMORY_VIEW_ROW_MAJOR));
+    rb_define_const(mMemoryViewTestUtils, "COLUMN_MAJOR",
+                    UINT2NUM(RUBY_MEMORY_VIEW_COLUMN_MAJOR));
+    rb_define_const(mMemoryViewTestUtils, "ANY_CONTIGUOUS",
+                    UINT2NUM(RUBY_MEMORY_VIEW_ANY_CONTIGUOUS));
+    rb_define_const(mMemoryViewTestUtils, "INDIRECT",
+                    UINT2NUM(RUBY_MEMORY_VIEW_INDIRECT));
+
     rb_define_module_function(mMemoryViewTestUtils, "available?", memory_view_available_p, 1);
     rb_define_module_function(mMemoryViewTestUtils, "register", memory_view_register, 1);
     rb_define_module_function(mMemoryViewTestUtils, "item_size_from_format", memory_view_item_size_from_format, 1);
     rb_define_module_function(mMemoryViewTestUtils, "parse_item_format", memory_view_parse_item_format, 1);
-    rb_define_module_function(mMemoryViewTestUtils, "get_memory_view_info", memory_view_get_memory_view_info, 1);
+    rb_define_module_function(mMemoryViewTestUtils, "get_memory_view_info", memory_view_get_memory_view_info, -1);
     rb_define_module_function(mMemoryViewTestUtils, "is_row_major_contiguous", memory_view_is_row_major_contiguous, 1);
     rb_define_module_function(mMemoryViewTestUtils, "is_column_major_contiguous", memory_view_is_column_major_contiguous, 1);
     rb_define_module_function(mMemoryViewTestUtils, "fill_contiguous_strides", memory_view_fill_contiguous_strides, 4);
     rb_define_module_function(mMemoryViewTestUtils, "ref_count_while_exporting", memory_view_ref_count_while_exporting, 2);
     rb_define_module_function(mMemoryViewTestUtils, "extract_item_members", memory_view_extract_item_members, 2);
+    rb_define_module_function(mMemoryViewTestUtils, "get_data", memory_view_get_data, 2);
+    rb_define_module_function(mMemoryViewTestUtils, "set_data", memory_view_set_data, 3);
+    rb_define_module_function(mMemoryViewTestUtils, "get", memory_view_get, 2);
 
     VALUE cExportableString = rb_define_class_under(mMemoryViewTestUtils, "ExportableString", rb_cObject);
     rb_define_method(cExportableString, "initialize", expstr_initialize, 1);

inits.c

@@ -47,6 +47,7 @@ rb_call_inits(void)
     CALL(marshal);
     CALL(Range);
     CALL(IO);
+    CALL(MemoryView); /* Must precede IO_Buffer */
     CALL(IO_Buffer)
     CALL(Dir);
     CALL(Time);
@@ -70,7 +71,6 @@ rb_call_inits(void)
     CALL(process);
     CALL(Rational);
     CALL(Complex);
-    CALL(MemoryView);
     CALL(pathname);
     CALL(version);
     CALL(vm_trace);

io_buffer.c

@@ -8,6 +8,7 @@
 
 #include "ruby/io/buffer.h"
 #include "ruby/fiber/scheduler.h"
+#include "ruby/memory_view.h"
 
 // For `rb_nogvl`.
 #include "ruby/thread.h"
@@ -4252,6 +4253,81 @@ io_buffer_bit_count(int argc, VALUE *argv, VALUE self)
     return SIZET2NUM(count);
 }
 
+static bool
+io_buffer_memory_view_get(VALUE self, rb_memory_view_t *view, int flags)
+{
+    struct rb_io_buffer *buffer = get_io_buffer(self);
+
+    if (buffer->base == NULL || !io_buffer_validate(buffer)) {
+        return false;
+    }
+
+    bool readonly = true;
+    if (flags & RUBY_MEMORY_VIEW_WRITABLE) {
+        if (io_buffer_readonly_p(buffer)) {
+            return false;
+        } else {
+            readonly = false;
+        }
+    }
+    rb_memory_view_init_as_byte_array(view, self, buffer->base, buffer->size, readonly);
+    if (flags & RUBY_MEMORY_VIEW_FORMAT) {
+        view->format = "C";
+    }
+    bool request_multi_dimensional = flags & RUBY_MEMORY_VIEW_MULTI_DIMENSIONAL;
+    bool request_strides =
+        (flags & RUBY_MEMORY_VIEW_STRIDES) == RUBY_MEMORY_VIEW_STRIDES;
+    if (request_multi_dimensional || request_strides) {
+        size_t n_metadata = 0;
+        if (request_multi_dimensional)
+            n_metadata++;
+        if (request_strides)
+            n_metadata++;
+        ssize_t *metadata_buffer = ALLOC_N(ssize_t, n_metadata);
+        size_t i = 0;
+        if (request_multi_dimensional) {
+            ssize_t *shape = &metadata_buffer[i];
+            shape[0] = buffer->size;
+            view->shape = shape;
+            i++;
+        }
+        if (request_strides) {
+            ssize_t *strides = &metadata_buffer[i];
+            strides[0] = 1;
+            view->strides = strides;
+            i++;
+        }
+        view->private_data = metadata_buffer;
+    }
+    io_buffer_lock(buffer);
+
+    return true;
+}
+
+static bool
+io_buffer_memory_view_release(VALUE self, rb_memory_view_t *view)
+{
+    rb_io_buffer_unlock(self);
+    if (view->private_data) {
+        xfree(view->private_data);
+    }
+    return true;
+}
+
+static bool
+io_buffer_memory_view_available_p(VALUE self)
+{
+    struct rb_io_buffer *buffer = get_io_buffer(self);
+
+    return buffer->base != NULL && io_buffer_validate(buffer);
+}
+
+static const rb_memory_view_entry_t io_buffer_memory_view_entry = {
+    .get_func = io_buffer_memory_view_get,
+    .release_func = io_buffer_memory_view_release,
+    .available_p_func = io_buffer_memory_view_available_p,
+};
+
 /*
  *  Document-class: IO::Buffer
  *
@@ -4276,6 +4352,16 @@ io_buffer_bit_count(int argc, VALUE *argv, VALUE self)
  *  like Fiber::Scheduler#io_read and Fiber::Scheduler#io_write and parsing binary
  *  protocols.
  *
+ *  == MemoryView Support
+ *
+ *  IO::Buffer supports the C-level MemoryView protocol, so C
+ *  extensions can use +rb_memory_view_get()+ to access the buffer's
+ *  memory directly (zero-copy) as a 1-dimensional contiguous array of
+ *  bytes. The memory view is writable if the buffer is not
+ *  #readonly? and +RUBY_MEMORY_VIEW_WRITABLE+ is specified.
+ *
+ *  While a MemoryView is exported, the buffer is locked.
+ *
  *  == Examples of Usage
  *
  *  Empty buffer:
@@ -4516,4 +4602,7 @@ Init_IO_Buffer(void)
     rb_define_method(rb_cIOBuffer, "pread", io_buffer_pread, -1);
     rb_define_method(rb_cIOBuffer, "write", io_buffer_write, -1);
     rb_define_method(rb_cIOBuffer, "pwrite", io_buffer_pwrite, -1);
+
+    // MemoryView:
+    rb_memory_view_register(rb_cIOBuffer, &io_buffer_memory_view_entry);
 }

test/ruby/test_io_buffer.rb

@@ -4,6 +4,7 @@ require 'tempfile'
 require 'rbconfig/sizeof'
 require 'io/nonblock'
 require '-test-/io_buffer'
+require "-test-/memory_view"
 
 class TestIOBuffer < Test::Unit::TestCase
   experimental = Warning[:experimental]
@@ -1703,4 +1704,249 @@ class TestIOBuffer < Test::Unit::TestCase
       buffer.hexdump(0, 1, 0)
     end
   end
+
+  def test_memory_view_null
+    buffer = IO::Buffer.new(0)
+    assert_false(MemoryViewTestUtils.available?(buffer))
+  end
+
+  def test_memory_view_available
+    buffer = IO::Buffer.new(8)
+    assert_true(MemoryViewTestUtils.available?(buffer))
+  end
+
+  def test_memory_view_get_simple
+    data = +"\x00\x01\x02\x03"
+    IO::Buffer.for(data) do |buffer|
+      info = MemoryViewTestUtils.get_memory_view_info(buffer)
+      assert_equal({
+                     obj: buffer,
+                     byte_size: data.bytesize,
+                     readonly: true,
+                     format: nil,
+                     item_size: 1,
+                     ndim: 1,
+                     shape: nil,
+                     strides: nil,
+                     sub_offsets: nil,
+                   },
+                   info)
+    end
+  end
+
+  def test_memory_view_get_writable
+    data = +"\x00\x01\x02\x03"
+    IO::Buffer.for(data) do |buffer|
+      flags = MemoryViewTestUtils::WRITABLE
+      info = MemoryViewTestUtils.get_memory_view_info(buffer, flags)
+      assert_equal({
+                     obj: buffer,
+                     byte_size: data.bytesize,
+                     readonly: false,
+                     format: nil,
+                     item_size: 1,
+                     ndim: 1,
+                     shape: nil,
+                     strides: nil,
+                     sub_offsets: nil,
+                   },
+                   info)
+    end
+  end
+
+  def test_memory_view_get_format
+    data = "\x00\x01\x02\x03".freeze
+    IO::Buffer.for(data) do |buffer|
+      flags = MemoryViewTestUtils::FORMAT
+      info = MemoryViewTestUtils.get_memory_view_info(buffer, flags)
+      assert_equal({
+                     obj: buffer,
+                     byte_size: data.bytesize,
+                     readonly: true,
+                     format: "C",
+                     item_size: 1,
+                     ndim: 1,
+                     shape: nil,
+                     strides: nil,
+                     sub_offsets: nil,
+                   },
+                   info)
+    end
+  end
+
+  def test_memory_view_get_multi_dimensional
+    data = "\x00\x01\x02\x03".freeze
+    IO::Buffer.for(data) do |buffer|
+      flags = MemoryViewTestUtils::MULTI_DIMENSIONAL
+      info = MemoryViewTestUtils.get_memory_view_info(buffer, flags)
+      assert_equal({
+                     obj: buffer,
+                     byte_size: data.bytesize,
+                     readonly: true,
+                     format: nil,
+                     item_size: 1,
+                     ndim: 1,
+                     shape: [data.bytesize],
+                     strides: nil,
+                     sub_offsets: nil,
+                   },
+                   info)
+    end
+  end
+
+  def test_memory_view_get_strides
+    data = "\x00\x01\x02\x03".freeze
+    IO::Buffer.for(data) do |buffer|
+      flags = MemoryViewTestUtils::STRIDES
+      info = MemoryViewTestUtils.get_memory_view_info(buffer, flags)
+      assert_equal({
+                     obj: buffer,
+                     byte_size: data.bytesize,
+                     readonly: true,
+                     format: nil,
+                     item_size: 1,
+                     ndim: 1,
+                     shape: [data.bytesize],
+                     strides: [1],
+                     sub_offsets: nil,
+                   },
+                   info)
+    end
+  end
+
+  def test_memory_view_get_row_major
+    data = "\x00\x01\x02\x03".freeze
+    IO::Buffer.for(data) do |buffer|
+      flags = MemoryViewTestUtils::ROW_MAJOR
+      info = MemoryViewTestUtils.get_memory_view_info(buffer, flags)
+      assert_equal({
+                     obj: buffer,
+                     byte_size: data.bytesize,
+                     readonly: true,
+                     format: nil,
+                     item_size: 1,
+                     ndim: 1,
+                     shape: [data.bytesize],
+                     strides: [1],
+                     sub_offsets: nil,
+                   },
+                   info)
+    end
+  end
+
+  def test_memory_view_get_column_major
+    data = "\x00\x01\x02\x03".freeze
+    IO::Buffer.for(data) do |buffer|
+      flags = MemoryViewTestUtils::COLUMN_MAJOR
+      info = MemoryViewTestUtils.get_memory_view_info(buffer, flags)
+      assert_equal({
+                     obj: buffer,
+                     byte_size: data.bytesize,
+                     readonly: true,
+                     format: nil,
+                     item_size: 1,
+                     ndim: 1,
+                     shape: [data.bytesize],
+                     strides: [1],
+                     sub_offsets: nil,
+                   },
+                   info)
+    end
+  end
+
+  def test_memory_view_get_any_contiguous
+    data = "\x00\x01\x02\x03".freeze
+    IO::Buffer.for(data) do |buffer|
+      flags = MemoryViewTestUtils::ANY_CONTIGUOUS
+      info = MemoryViewTestUtils.get_memory_view_info(buffer, flags)
+      assert_equal({
+                     obj: buffer,
+                     byte_size: data.bytesize,
+                     readonly: true,
+                     format: nil,
+                     item_size: 1,
+                     ndim: 1,
+                     shape: [data.bytesize],
+                     strides: [1],
+                     sub_offsets: nil,
+                   },
+                   info)
+    end
+  end
+
+  def test_memory_view_get_indirect
+    data = "\x00\x01\x02\x03".freeze
+    IO::Buffer.for(data) do |buffer|
+      flags = MemoryViewTestUtils::INDIRECT
+      info = MemoryViewTestUtils.get_memory_view_info(buffer, flags)
+      assert_equal({
+                     obj: buffer,
+                     byte_size: data.bytesize,
+                     readonly: true,
+                     format: nil,
+                     item_size: 1,
+                     ndim: 1,
+                     shape: [data.bytesize],
+                     strides: [1],
+                     sub_offsets: nil,
+                   },
+                   info)
+    end
+  end
+
+  def test_memory_view_readonly
+    data = "\x00\x01\x02\x03".freeze
+    buffer = IO::Buffer.for(data)
+    # rb_memory_view_get(RUBY_MEMORY_VIEW_WRITABLE) is failed with
+    # readonly IO::Buffer.
+    assert_nil(MemoryViewTestUtils.set_data(buffer, 1, 0x11))
+    assert_equal(data, MemoryViewTestUtils.get_data(buffer, 0..(data.bytesize)))
+  end
+
+  def test_memory_view_writable
+    IO::Buffer.for(+"\x00\x01\x02\x03") do |buffer|
+      assert_true(MemoryViewTestUtils.set_data(buffer, 1, 0x11))
+      assert_equal(0x11, buffer.get_value(:U8, 1))
+    end
+  end
+
+  def test_memory_view_locked
+    IO::Buffer.for("\x00\x01\x02\x03".freeze) do |buffer|
+      flags = MemoryViewTestUtils::SIMPLE
+      MemoryViewTestUtils.get(buffer, flags) do
+        assert_predicate buffer, :locked?
+      end
+      refute_predicate buffer, :locked?
+    end
+  end
+
+  def test_memory_view_nested_locked
+    IO::Buffer.for("\x00\x01\x02\x03".freeze) do |buffer|
+      flags = MemoryViewTestUtils::SIMPLE
+      MemoryViewTestUtils.get(buffer, flags) do
+        MemoryViewTestUtils.get(buffer, flags) do
+          assert_predicate buffer, :locked?
+        end
+        assert_predicate buffer, :locked?
+      end
+      refute_predicate buffer, :locked?
+    end
+  end
+
+  def test_memory_view_sliced_nested_locked
+    IO::Buffer.for("\x00\x01\x02\x03".freeze) do |buffer|
+      sliced = buffer.slice(1, 2)
+      flags = MemoryViewTestUtils::SIMPLE
+      MemoryViewTestUtils.get(sliced, flags) do
+        MemoryViewTestUtils.get(sliced, flags) do
+          assert_predicate sliced, :locked?
+          assert_predicate buffer, :locked?
+        end
+        assert_predicate sliced, :locked?
+        assert_predicate buffer, :locked?
+      end
+      refute_predicate sliced, :locked?
+      refute_predicate buffer, :locked?
+    end
+  end
 end