Implement rb_get_alloc_func and spec related functions

lib/truffle/truffle/cext.rb

@@ -14,6 +14,7 @@ module Truffle::CExt
   DATA_HOLDER = Object.new
   DATA_MEMSIZER = Object.new
   RB_TYPE = Object.new
+  ALLOCATOR_FUNC = Object.new
 
   extend self
 
@@ -1278,17 +1279,40 @@ def rb_eval_string(str)
     eval(str)
   end
 
+  BASIC_ALLOC = ::BasicObject.singleton_class.instance_method(:__allocate__)
+
+  def rb_newobj_of(ruby_class)
+    # we need to bypass __allocate__ on ruby_class
+    BASIC_ALLOC.bind(ruby_class).call
+  end
+
   def rb_define_alloc_func(ruby_class, function)
     ruby_class.singleton_class.define_method(:__allocate__) do
       TrufflePrimitive.cext_unwrap(TrufflePrimitive.call_with_c_mutex(function, [TrufflePrimitive.cext_wrap(self)]))
     end
     class << ruby_class
       private :__allocate__
     end
+    hidden_variable_set(ruby_class.singleton_class, ALLOCATOR_FUNC, function)
+  end
+
+  def rb_get_alloc_func(ruby_class)
+    return nil unless Class === ruby_class
+    begin
+      allocate_method = ruby_class.method(:__allocate__).owner
+    rescue NameError
+      nil
+    else
+      hidden_variable_get(allocate_method, ALLOCATOR_FUNC)
+    end
   end
 
   def rb_undef_alloc_func(ruby_class)
     ruby_class.singleton_class.send(:undef_method, :__allocate__)
+  rescue NameError
+    # it's fine to call this on a class that doesn't have an allocator
+  else
+    hidden_variable_set(ruby_class.singleton_class, ALLOCATOR_FUNC, nil)
   end
 
   def rb_alias(mod, new_name, old_name)

spec/ruby/optional/capi/ext/object_spec.c

@@ -345,6 +345,31 @@ static VALUE object_spec_rb_class_inherited_p(VALUE self, VALUE mod, VALUE arg)
   return rb_class_inherited_p(mod, arg);
 }
 
+static VALUE speced_allocator(VALUE klass) {
+  VALUE instance = rb_newobj_of(klass, 0);
+  rb_funcall(instance, rb_intern("instance_variable_set"), 2, ID2SYM(rb_intern("@from_custom_allocator")), Qtrue);
+  return instance;
+}
+
+static VALUE define_alloc_func(VALUE self, VALUE klass) {
+  rb_define_alloc_func(klass, speced_allocator);
+  return Qnil;
+}
+
+static VALUE undef_alloc_func(VALUE self, VALUE klass) {
+  rb_undef_alloc_func(klass);
+  return Qnil;
+}
+
+static VALUE speced_allocator_p(VALUE self, VALUE klass) {
+  rb_alloc_func_t allocator = rb_get_alloc_func(klass);
+  return (allocator == speced_allocator) ? Qtrue : Qfalse;
+}
+
+static VALUE allocator_nil_p(VALUE self, VALUE klass) {
+  rb_alloc_func_t allocator = rb_get_alloc_func(klass);
+  return allocator ? Qfalse : Qtrue;
+}
 
 void Init_object_spec(void) {
   VALUE cls = rb_define_class("CApiObjectSpecs", rb_cObject);
@@ -412,6 +437,10 @@ void Init_object_spec(void) {
   rb_define_method(cls, "rb_ivar_defined", object_spec_rb_ivar_defined, 2);
   rb_define_method(cls, "rb_copy_generic_ivar", object_spec_rb_copy_generic_ivar, 2);
   rb_define_method(cls, "rb_free_generic_ivar", object_spec_rb_free_generic_ivar, 1);
+  rb_define_method(cls, "rb_define_alloc_func", define_alloc_func, 1);
+  rb_define_method(cls, "rb_undef_alloc_func", undef_alloc_func, 1);
+  rb_define_method(cls, "speced_allocator?", speced_allocator_p, 1);
+  rb_define_method(cls, "allocator_nil?", allocator_nil_p, 1);
 }
 
 #ifdef __cplusplus

spec/ruby/optional/capi/object_spec.rb

@@ -886,5 +886,55 @@ def reach
         o.instance_variables.should == []
       end
     end
+
+    describe "allocator accessors" do
+      describe "rb_define_alloc_func" do
+        it "sets up the allocator" do
+          klass = Class.new
+          @o.rb_define_alloc_func(klass)
+          klass.allocate.should have_instance_variable(:@from_custom_allocator)
+        end
+      end
+
+      describe "rb_get_alloc_func" do
+        it "gets the allocator that is defined directly on a class" do
+          klass = Class.new
+          @o.rb_define_alloc_func(klass)
+          @o.speced_allocator?(Object).should be_false
+          @o.speced_allocator?(klass).should be_true
+        end
+
+        it "gets the allocator that is inherited" do
+          parent = Class.new
+          @o.rb_define_alloc_func(parent)
+          klass = Class.new(parent)
+          @o.speced_allocator?(Object).should be_false
+          @o.speced_allocator?(klass).should be_true
+        end
+      end
+
+      describe "rb_undef_alloc_func" do
+        it "does nothing when called on a class without a custom allocator" do
+          -> { @o.allocator_nil?(Class.new) }.should_not raise_error
+        end
+
+        it "undefs the allocator for the class" do
+          klass = Class.new
+          @o.rb_define_alloc_func(klass)
+          @o.speced_allocator?(klass).should be_true
+          @o.rb_undef_alloc_func(klass)
+          @o.allocator_nil?(klass).should be_true
+        end
+
+        it "undefs the allocator for a class that inherits a allocator" do
+          parent = Class.new
+          @o.rb_define_alloc_func(parent)
+          klass = Class.new(parent)
+          @o.speced_allocator?(klass).should be_true
+          @o.rb_undef_alloc_func(klass)
+          @o.allocator_nil?(klass).should be_true
+        end
+      end
+    end
   end
 end

src/main/c/cext/ruby.c

@@ -3685,7 +3685,7 @@ void rb_remove_method_id(VALUE klass, ID mid) {
 }
 
 rb_alloc_func_t rb_get_alloc_func(VALUE klass) {
-  rb_tr_error("rb_get_alloc_func not implemented");
+  return RUBY_CEXT_INVOKE_NO_WRAP("rb_get_alloc_func", klass);
 }
 
 void rb_clear_constant_cache(void) {
@@ -4667,7 +4667,8 @@ VALUE rb_newobj(void) {
 }
 
 VALUE rb_newobj_of(VALUE klass, VALUE flags) {
-  rb_tr_error("rb_newobj_of not implemented");
+  // ignore flags for now
+  return RUBY_CEXT_INVOKE("rb_newobj_of", klass);
 }
 
 VALUE rb_obj_setup(VALUE obj, VALUE klass, VALUE type) {