Keep references of memory-view-exported objects (#3816)

* memory_view.c: remove a reference in view->obj at rb_memory_view_release

* memory_view.c: keep references of memory-view-exported objects

* Update common.mk

* memory_view.c: Use st_update

5 files changed, 193 insertions, 57 deletions

diff --git a/common.mk b/common.mk
index 6120761cf2..994b5c3a45 100644
--- a/common.mk
+++ b/common.mk
@@ -8006,6 +8006,7 @@ math.$(OBJEXT): {$(VPATH)}missing.h
 math.$(OBJEXT): {$(VPATH)}st.h
 math.$(OBJEXT): {$(VPATH)}subst.h
 memory_view.$(OBJEXT): $(hdrdir)/ruby/ruby.h
+memory_view.$(OBJEXT): $(top_srcdir)/internal/hash.h
 memory_view.$(OBJEXT): $(top_srcdir)/internal/util.h
 memory_view.$(OBJEXT): $(top_srcdir)/internal/variable.h
 memory_view.$(OBJEXT): {$(VPATH)}assert.h
diff --git a/ext/-test-/memory_view/memory_view.c b/ext/-test-/memory_view/memory_view.c
index 0ae9f457ac..7f1f007ba4 100644
--- a/ext/-test-/memory_view/memory_view.c
+++ b/ext/-test-/memory_view/memory_view.c
@@ -24,36 +24,11 @@ static VALUE sym_endianness;
 static VALUE sym_little_endian;
 static VALUE sym_big_endian;
 
-static VALUE exported_objects;
-
 static int
 exportable_string_get_memory_view(VALUE obj, rb_memory_view_t *view, int flags)
 {
     VALUE str = rb_ivar_get(obj, id_str);
     rb_memory_view_init_as_byte_array(view, obj, RSTRING_PTR(str), RSTRING_LEN(str), true);
-
-    VALUE count = rb_hash_lookup2(exported_objects, obj, INT2FIX(0));
-    count = rb_funcall(count, '+', 1, INT2FIX(1));
-    rb_hash_aset(exported_objects, obj, count);
-
-    return 1;
-}
-
-static int
-exportable_string_release_memory_view(VALUE obj, rb_memory_view_t *view)
-{
-    VALUE count = rb_hash_lookup2(exported_objects, obj, INT2FIX(0));
-    if (INT2FIX(1) == count) {
-        rb_hash_delete(exported_objects, obj);
-    }
-    else if (INT2FIX(0) == count) {
-        rb_raise(rb_eRuntimeError, "Duplicated releasing of a memory view has been occurred for %"PRIsVALUE, obj);
-    }
-    else {
-        count = rb_funcall(count, '-', 1, INT2FIX(1));
-        rb_hash_aset(exported_objects, obj, count);
-    }
-
     return 1;
 }
 
@@ -65,7 +40,7 @@ exportable_string_memory_view_available_p(VALUE obj)
 
 static const rb_memory_view_entry_t exportable_string_memory_view_entry = {
     exportable_string_get_memory_view,
-    exportable_string_release_memory_view,
+    NULL,
     exportable_string_memory_view_available_p
 };
 
@@ -207,6 +182,54 @@ memory_view_fill_contiguous_strides(VALUE mod, VALUE ndim_v, VALUE item_size_v,
 }
 
 static VALUE
+memory_view_get_ref_count(VALUE obj)
+{
+    extern VALUE rb_memory_view_exported_object_registry;
+    extern const rb_data_type_t rb_memory_view_exported_object_registry_data_type;
+
+    if (rb_memory_view_exported_object_registry == Qundef) {
+        return Qnil;
+    }
+
+    st_table *table;
+    TypedData_Get_Struct(rb_memory_view_exported_object_registry, st_table,
+                         &rb_memory_view_exported_object_registry_data_type,
+                         table);
+
+    st_data_t count;
+    if (st_lookup(table, (st_data_t)obj, &count)) {
+        return ULL2NUM(count);
+    }
+
+    return Qnil;
+}
+
+static VALUE
+memory_view_ref_count_while_exporting_i(VALUE obj, long n)
+{
+    if (n == 0) {
+        return memory_view_get_ref_count(obj);
+    }
+
+    rb_memory_view_t view;
+    if (!rb_memory_view_get(obj, &view, 0)) {
+        return Qnil;
+    }
+
+    VALUE ref_count = memory_view_ref_count_while_exporting_i(obj, n-1);
+    rb_memory_view_release(&view);
+
+    return ref_count;
+}
+
+static VALUE
+memory_view_ref_count_while_exporting(VALUE mod, VALUE obj, VALUE n)
+{
+    Check_Type(n, T_FIXNUM);
+    return memory_view_ref_count_while_exporting_i(obj, FIX2LONG(n));
+}
+
+static VALUE
 expstr_initialize(VALUE obj, VALUE s)
 {
     rb_ivar_set(obj, id_str, s);
@@ -247,28 +270,6 @@ mdview_get_memory_view(VALUE obj, rb_memory_view_t *view, int flags)
     view->shape = shape;
     view->strides = strides;
 
-    VALUE count = rb_hash_lookup2(exported_objects, obj, INT2FIX(0));
-    count = rb_funcall(count, '+', 1, INT2FIX(1));
-    rb_hash_aset(exported_objects, obj, count);
-
-    return 1;
-}
-
-static int
-mdview_release_memory_view(VALUE obj, rb_memory_view_t *view)
-{
-    VALUE count = rb_hash_lookup2(exported_objects, obj, INT2FIX(0));
-    if (INT2FIX(1) == count) {
-        rb_hash_delete(exported_objects, obj);
-    }
-    else if (INT2FIX(0) == count) {
-        rb_raise(rb_eRuntimeError, "Duplicated releasing of a memory view has been occurred for %"PRIsVALUE, obj);
-    }
-    else {
-        count = rb_funcall(count, '-', 1, INT2FIX(1));
-        rb_hash_aset(exported_objects, obj, count);
-    }
-
     return 1;
 }
 
@@ -280,7 +281,7 @@ mdview_memory_view_available_p(VALUE obj)
 
 static const rb_memory_view_entry_t mdview_memory_view_entry = {
     mdview_get_memory_view,
-    mdview_release_memory_view,
+    NULL,
     mdview_memory_view_available_p
 };
 
@@ -340,6 +341,7 @@ Init_memory_view(void)
     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, "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);
 
     VALUE cExportableString = rb_define_class_under(mMemoryViewTestUtils, "ExportableString", rb_cObject);
     rb_define_method(cExportableString, "initialize", expstr_initialize, 1);
@@ -393,7 +395,4 @@ Init_memory_view(void)
     DEF_ALIGNMENT_CONST(double, DOUBLE);
 
 #undef DEF_ALIGNMENT_CONST
-
-    exported_objects = rb_hash_new();
-    rb_gc_register_mark_object(exported_objects);
 }
diff --git a/include/ruby/memory_view.h b/include/ruby/memory_view.h
index e3897f830e..e2c5cd9a03 100644
--- a/include/ruby/memory_view.h
+++ b/include/ruby/memory_view.h
@@ -136,6 +136,10 @@ int rb_memory_view_available_p(VALUE obj);
 int rb_memory_view_get(VALUE obj, rb_memory_view_t* memory_view, int flags);
 int rb_memory_view_release(rb_memory_view_t* memory_view);
 
+/* for testing */
+RUBY_EXTERN VALUE rb_memory_view_exported_object_registry;
+RUBY_EXTERN const rb_data_type_t rb_memory_view_exported_object_registry_data_type;
+
 RBIMPL_SYMBOL_EXPORT_END()
 
 RBIMPL_ATTR_PURE()
diff --git a/memory_view.c b/memory_view.c
index b2d81cb957..aade3a4aaf 100644
--- a/memory_view.c
+++ b/memory_view.c
@@ -7,6 +7,7 @@
 **********************************************************************/
 
 #include "internal.h"
+#include "internal/hash.h"
 #include "internal/variable.h"
 #include "internal/util.h"
 #include "ruby/memory_view.h"
@@ -15,10 +16,119 @@
     (result) = RUBY_ALIGNOF(T); \
 } while(0)
 
+// Exported Object Registry
+
+VALUE rb_memory_view_exported_object_registry = Qundef;
+
+static int
+exported_object_registry_mark_key_i(st_data_t key, st_data_t value, st_data_t data)
+{
+    rb_gc_mark(key);
+    return ST_CONTINUE;
+}
+
+static void
+exported_object_registry_mark(void *ptr)
+{
+    st_table *table = ptr;
+    st_foreach(table, exported_object_registry_mark_key_i, 0);
+}
+
+static void
+exported_object_registry_free(void *ptr)
+{
+    st_table *table = ptr;
+    st_clear(table);
+    st_free_table(table);
+}
+
+const rb_data_type_t rb_memory_view_exported_object_registry_data_type = {
+    "memory_view/exported_object_registry",
+    {
+	exported_object_registry_mark,
+	exported_object_registry_free,
+	0,
+    },
+    0, 0, RUBY_TYPED_FREE_IMMEDIATELY
+};
+
+static void
+init_exported_object_registry(void)
+{
+    if (rb_memory_view_exported_object_registry != Qundef) {
+        return;
+    }
+
+    st_table *table = rb_init_identtable();
+    VALUE obj = TypedData_Wrap_Struct(
+        0, &rb_memory_view_exported_object_registry_data_type, table);
+    rb_gc_register_mark_object(obj);
+    rb_memory_view_exported_object_registry = obj;
+}
+
+static inline st_table *
+get_exported_object_table(void)
+{
+    st_table *table;
+    TypedData_Get_Struct(rb_memory_view_exported_object_registry, st_table,
+                         &rb_memory_view_exported_object_registry_data_type,
+                         table);
+    return table;
+}
+
+static int
+update_exported_object_ref_count(st_data_t *key, st_data_t *val, st_data_t arg, int existing)
+{
+    if (existing) {
+        *val += 1;
+    }
+    else {
+        *val = 1;
+    }
+    return ST_CONTINUE;
+}
+
+static void
+register_exported_object(VALUE obj)
+{
+    if (rb_memory_view_exported_object_registry == Qundef) {
+        init_exported_object_registry();
+    }
+
+    st_table *table = get_exported_object_table();
+
+    st_update(table, (st_data_t)obj, update_exported_object_ref_count, 0);
+}
+
+static void
+unregister_exported_object(VALUE obj)
+{
+    if (rb_memory_view_exported_object_registry == Qundef) {
+        return;
+    }
+
+    st_table *table = get_exported_object_table();
+
+    st_data_t count;
+    if (!st_lookup(table, (st_data_t)obj, &count)) {
+        return;
+    }
+
+    if (--count == 0) {
+        st_data_t key = (st_data_t)obj;
+        st_delete(table, &key, &count);
+    }
+    else {
+        st_insert(table, (st_data_t)obj, count);
+    }
+}
+
+// MemoryView
+
 static ID id_memory_view;
 
 static const rb_data_type_t memory_view_entry_data_type = {
-    "memory_view",
+    "memory_view/entry",
     {
 	0,
 	0,
@@ -481,8 +591,13 @@ rb_memory_view_get(VALUE obj, rb_memory_view_t* view, int flags)
 {
     VALUE klass = CLASS_OF(obj);
     const rb_memory_view_entry_t *entry = lookup_memory_view_entry(klass);
-    if (entry)
-        return (*entry->get_func)(obj, view, flags);
+    if (entry) {
+        int rv = (*entry->get_func)(obj, view, flags);
+        if (rv) {
+            register_exported_object(view->obj);
+        }
+        return rv;
+    }
     else
         return 0;
 }
@@ -493,8 +608,17 @@ rb_memory_view_release(rb_memory_view_t* view)
 {
     VALUE klass = CLASS_OF(view->obj);
     const rb_memory_view_entry_t *entry = lookup_memory_view_entry(klass);
-    if (entry)
-        return (*entry->release_func)(view->obj, view);
+    if (entry) {
+        int rv = 1;
+        if (entry->release_func) {
+            rv = (*entry->release_func)(view->obj, view);
+        }
+        if (rv) {
+            unregister_exported_object(view->obj);
+            view->obj = Qnil;
+        }
+        return rv;
+    }
     else
         return 0;
 }
diff --git a/test/ruby/test_memory_view.rb b/test/ruby/test_memory_view.rb
index 7a5ebb994c..668d738974 100644
--- a/test/ruby/test_memory_view.rb
+++ b/test/ruby/test_memory_view.rb
@@ -197,6 +197,14 @@ class TestMemoryView < Test::Unit::TestCase
     assert_equal(expected_result, members)
   end
 
+  def test_ref_count_with_exported_object
+    es = MemoryViewTestUtils::ExportableString.new("ruby")
+    assert_equal(1, MemoryViewTestUtils.ref_count_while_exporting(es, 1))
+    assert_equal(2, MemoryViewTestUtils.ref_count_while_exporting(es, 2))
+    assert_equal(10, MemoryViewTestUtils.ref_count_while_exporting(es, 10))
+    assert_nil(MemoryViewTestUtils.ref_count_while_exporting(es, 0))
+  end
+
   def test_rb_memory_view_init_as_byte_array
     # ExportableString's memory view is initialized by rb_memory_view_init_as_byte_array
     es = MemoryViewTestUtils::ExportableString.new("ruby")