Implement memsize functions for wrapped structures.

Author: aardvark179

lib/truffle/objspace.rb

@@ -152,7 +152,11 @@ def dump_all(output: :file)
   module_function :dump_all
 
   def memsize_of(object)
-    Truffle::ObjSpace.memsize_of(object)
+    unless Truffle::ObjSpace.has_sizer?(object)
+      Truffle::ObjSpace.memsize_of(object)
+    else
+      Truffle::ObjSpace.memsize_of(object) + Truffle::ObjSpace.sizer(object).call
+    end
   end
   module_function :memsize_of
 

lib/truffle/truffle/cext.rb

@@ -1415,13 +1415,14 @@ def rb_data_object_wrap(ruby_class, data, mark, free)
     object
   end
 
-  def rb_data_typed_object_wrap(ruby_class, data, data_type, mark, free)
+  def rb_data_typed_object_wrap(ruby_class, data, data_type, mark, free, size)
     ruby_class = Object if Truffle::Interop.null?(ruby_class)
     object = BASIC_OBJECT_ALLOCATE.bind(ruby_class).call
     data_holder = DataHolder.new(data)
     hidden_variable_set object, :data_type, data_type
     hidden_variable_set object, :data_holder, data_holder
     ObjectSpace.define_finalizer object, data_finalizer(free, data_holder) unless free.nil?
+    Truffle::ObjSpace.define_sizer object, data_sizer(size, data_holder) unless size.nil?
     define_marker object, data_marker(mark, data_holder) unless mark.nil?
     object
   end
@@ -1444,6 +1445,14 @@ def data_marker(mark, data_holder)
     }
   end
 
+  def data_sizer(sizer, data_holder)
+    # In a separate method to avoid capturing the object
+    raise unless sizer.respond_to?(:call)
+    proc {
+      Truffle.invoke_primitive(:interop_call_c_with_mutex, sizer, [data_holder.data])
+    }
+  end
+
   def rb_ruby_verbose_ptr
     $VERBOSE
   end

src/main/c/cext/ruby.c

@@ -2838,7 +2838,7 @@ VALUE rb_data_object_alloc_managed(VALUE klass, size_t size, RUBY_DATA_FUNC dmar
 
 VALUE rb_data_typed_object_wrap(VALUE ruby_class, void *data, const rb_data_type_t *data_type) {
   return rb_tr_wrap(polyglot_invoke(RUBY_CEXT, "rb_data_typed_object_wrap",
-                                            rb_tr_unwrap(ruby_class), data, data_type, data_type->function.dmark, rb_tr_free_function(data_type->function.dfree)));
+                                    rb_tr_unwrap(ruby_class), data, data_type, data_type->function.dmark, rb_tr_free_function(data_type->function.dfree), data_type->function.dsize));
 }
 
 VALUE rb_data_typed_object_zalloc(VALUE ruby_class, size_t size, const rb_data_type_t *data_type) {

src/main/java/org/truffleruby/Layouts.java

@@ -104,6 +104,7 @@ public abstract class Layouts {
     public static final HiddenKey ASSOCIATED_IDENTIFIER = new HiddenKey("associated");
     public static final HiddenKey FINALIZER_REF_IDENTIFIER = new HiddenKey("finalizerRef");
     public static final HiddenKey MARKED_OBJECTS_IDENTIFIER = new HiddenKey("marked_objects");
+    public static final HiddenKey MEMSIZER_IDENTIFIER = new HiddenKey("memsizer");
     public static final HiddenKey VALUE_WRAPPER_IDENTIFIER = new HiddenKey("value_wrapper");
 
     // Generated layouts

src/main/java/org/truffleruby/stdlib/ObjSpaceNodes.java

@@ -19,6 +19,10 @@
 import org.truffleruby.core.regexp.MatchDataNodes.ValuesNode;
 import org.truffleruby.core.string.StringOperations;
 import org.truffleruby.language.objects.ObjectGraph;
+import org.truffleruby.language.objects.ReadObjectFieldNode;
+import org.truffleruby.language.objects.ReadObjectFieldNodeGen;
+import org.truffleruby.language.objects.WriteObjectFieldNode;
+import org.truffleruby.language.objects.WriteObjectFieldNodeGen;
 
 import java.util.Set;
 
@@ -65,7 +69,40 @@ public int memsizeOfObject(DynamicObject object) {
         public int memsize(Object object) {
             return 0;
         }
+    }
+
+    @CoreMethod(names = "has_sizer?", isModuleFunction = true, required = 1)
+    public abstract static class HasMemSizer extends CoreMethodArrayArgumentsNode {
+
+        @Child private ReadObjectFieldNode setSizerNode = ReadObjectFieldNodeGen.create(Layouts.MEMSIZER_IDENTIFIER, nil());
+
+        @Specialization
+        public boolean hasSizer(DynamicObject object) {
+            return setSizerNode.execute(object) != nil();
+        }
+    }
+
+    @CoreMethod(names = "sizer", isModuleFunction = true, required = 1)
+    public abstract static class GetMemSizer extends CoreMethodArrayArgumentsNode {
+
+        @Child private ReadObjectFieldNode setSizerNode = ReadObjectFieldNodeGen.create(Layouts.MEMSIZER_IDENTIFIER, nil());
 
+        @Specialization
+        public Object hasSizer(DynamicObject object) {
+            return setSizerNode.execute(object);
+        }
+    }
+
+    @CoreMethod(names = "define_sizer", isModuleFunction = true, required = 2)
+    public abstract static class DefineMemSizer extends CoreMethodArrayArgumentsNode {
+
+        @Child private WriteObjectFieldNode setSizerNode = WriteObjectFieldNodeGen.create(Layouts.MEMSIZER_IDENTIFIER);
+
+        @Specialization
+        public DynamicObject defineSizer(DynamicObject object, DynamicObject sizer) {
+            setSizerNode.write(object, sizer);
+            return nil();
+        }
     }
 
     @CoreMethod(names = "adjacent_objects", isModuleFunction = true, required = 1)