From 177ead8d16ba989c9c8cf3776205ecf00116e1fd Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Tue, 25 Aug 2026 05:50:39 +0900 Subject: [PATCH] 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 | 198 ++++++++++++++++++--- inits.c | 2 +- io_buffer.c | 89 ++++++++++ test/ruby/test_io_buffer.rb | 246 +++++++++++++++++++++++++++ 4 files changed, 514 insertions(+), 21 deletions(-) diff --git a/ext/-test-/memory_view/memory_view.c b/ext/-test-/memory_view/memory_view.c index 4dee65dd3b9..a61652eefc7 100644 --- a/ext/-test-/memory_view/memory_view.c +++ b/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); diff --git a/inits.c b/inits.c index 1be4916e87b..e4323cc4a84 100644 --- a/inits.c +++ b/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); diff --git a/io_buffer.c b/io_buffer.c index 5a5128aa1dd..8099c23a37e 100644 --- a/io_buffer.c +++ b/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); } diff --git a/test/ruby/test_io_buffer.rb b/test/ruby/test_io_buffer.rb index 500ca5b7eb0..6e0c858fff5 100644 --- a/test/ruby/test_io_buffer.rb +++ b/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