Extend ObjectSpace.each_object to include classes' singletons.

See rails/rails#22376.
jruby · Nov 23, 2015 · 2497978 · 2497978
1 parent 4ddccfe
commit 2497978
Showing 2 changed files with 27 additions and 0 deletions.
diff --git a/core/src/main/java/org/jruby/RubyObjectSpace.java b/core/src/main/java/org/jruby/RubyObjectSpace.java
@@ -153,6 +153,16 @@ public Object apply(IRubyObject arg1) {
             for (IRubyObject arg : modules) {
                 block.yield(context, arg);
             }
+        } else if (args[0].getClass() == MetaClass.class) {
+            // each_object(Cls.singleton_class) is basically a walk of Cls and all descendants of Cls.
+            // In other words, this is walking all instances of Cls's singleton class and its subclasses.
+            IRubyObject attached = ((MetaClass)args[0]).getAttached();
+            block.yield(context, attached);
+            if (attached instanceof RubyClass) {
+                for (RubyClass child : ((RubyClass)attached).subclasses(true)) {
+                    block.yield(context, child);
+                }
+            }
         } else {
             if (!runtime.isObjectSpaceEnabled()) {
                 throw runtime.newRuntimeError("ObjectSpace is disabled; each_object will only work with Class, pass -X+O to enable");

diff --git a/test/test_objectspace.rb b/test/test_objectspace.rb
@@ -88,4 +88,21 @@ def test_finalization
       assert_equal ["finalizing #{obj1_id}"], results
     end
   end
+
+  # See rails/rails#22376.
+  def test_each_object_singleton_class
+    # disable objectspace; we want this to always work
+    old_objectspace = JRuby.objectspace
+    JRuby.objectspace = false
+
+    a = Class.new
+    b = Class.new(a)
+    c = Class.new(a)
+    d = Class.new(b)
+
+    classes = ObjectSpace.each_object(a.singleton_class).to_a
+    assert_equal(classes.sort_by(&:name), [a, b, c, d])
+  ensure
+    JRuby.objectspace = old_objectspace
+  end
 end