From facd83afca498a6d8598da1878e54aa99e60a7b0 Mon Sep 17 00:00:00 2001
From: Charles Oliver Nutter <headius@headius.com>
Date: Mon, 25 Nov 2024 15:19:50 -0600
Subject: [PATCH] Optimize Class#subclasses

There were several issues with the old implementation:

* Iterating by creating a Set and an Iterator adds significant
  allocation overhead. Better to use Map#forEach.
* ConcurrentWeakHashMap is a very old implementation that has
  several inefficiencies as well as no custom implementation of
  Map#forEach (it falls back on the default impl which creates a
  Set and Iterator). Instead we use a synchronized WeakHashMap.
* The initial array was created with a default size of 4 elements.
  For many classes this will be insufficient, leading to excessive
  array allocation and copying as it grows. Instead we use the
  current class's "subclasses" size as an initial estimate.
* When no subclasses appear to be present, bail out early with
  an empty array that does minimal allocation.

Performance is significantly better for zero, small, and large
subclass lists.

1M * Object.subclasses, with 83 elements:

Before:

  1.680000   0.010000   1.690000 (  1.500808)
  1.430000   0.010000   1.440000 (  1.392465)
  1.390000   0.010000   1.400000 (  1.388177)
  1.410000   0.000000   1.410000 (  1.387108)
  1.390000   0.010000   1.400000 (  1.384741)

After:

  1.170000   0.020000   1.190000 (  1.017587)
  0.940000   0.000000   0.940000 (  0.927016)
  0.930000   0.010000   0.940000 (  0.918018)
  1.090000   0.010000   1.100000 (  0.931537)
  0.930000   0.000000   0.930000 (  0.925984)

10M * Numeric.subclasses, with 4 elements:

Before:

  0.720000   0.020000   0.740000 (  0.627838)
  0.570000   0.010000   0.580000 (  0.560524)
  0.550000   0.010000   0.560000 (  0.556029)
  0.760000   0.020000   0.780000 (  0.563721)
  0.570000   0.020000   0.590000 (  0.562469)

After:

  0.610000   0.010000   0.620000 (  0.495835)
  0.620000   0.010000   0.630000 (  0.426894)
  0.450000   0.010000   0.460000 (  0.432614)
  0.440000   0.000000   0.440000 (  0.426404)
  0.420000   0.010000   0.430000 (  0.418811)
---
 core/src/main/java/org/jruby/RubyClass.java | 31 +++++++++++++--------
 1 file changed, 19 insertions(+), 12 deletions(-)

diff --git a/core/src/main/java/org/jruby/RubyClass.java b/core/src/main/java/org/jruby/RubyClass.java
index 6c687fbf779..33e2354c1e5 100644
--- a/core/src/main/java/org/jruby/RubyClass.java
+++ b/core/src/main/java/org/jruby/RubyClass.java
@@ -93,7 +93,7 @@
 import org.jruby.util.Loader;
 import org.jruby.util.OneShotClassLoader;
 import org.jruby.util.StringSupport;
-import org.jruby.util.collections.ConcurrentWeakHashMap;
+import org.jruby.util.WeakIdentityHashMap;
 import org.jruby.util.log.Logger;
 import org.jruby.util.log.LoggerFactory;
 import org.objectweb.asm.AnnotationVisitor;
@@ -1019,7 +1019,14 @@ protected void setModuleSuperClass(RubyClass superClass) {
 
     @JRubyMethod
     public IRubyObject subclasses(ThreadContext context) {
-        RubyArray<RubyClass> subs = RubyArray.newArray(context.runtime);
+        Map<RubyClass, Object> subclasses = this.subclasses;
+        int sizeEstimate = subclasses == null ? 0 : subclasses.size();
+
+        if (sizeEstimate == 0) {
+            return RubyArray.newEmptyArray(context.runtime);
+        }
+
+        RubyArray<RubyClass> subs = RubyArray.newArray(context.runtime, sizeEstimate);
 
         concreteSubclasses(subs);
 
@@ -1081,18 +1088,18 @@ private void subclassesInner(Collection<RubyClass> mine, boolean includeDescenda
         }
     }
 
-    private void concreteSubclasses(Collection<RubyClass> subs) {
+    private void concreteSubclasses(RubyArray<RubyClass> subs) {
         Map<RubyClass, Object> subclasses = this.subclasses;
         if (subclasses != null) {
-            Set<RubyClass> keys = subclasses.keySet();
-            for (RubyClass klass: keys) {
-                if (klass.isSingleton()) continue;
-                if (klass.isIncluded() || klass.isPrepended()) {
-                    klass.concreteSubclasses(subs);
-                    continue;
+            subclasses.forEach((klass, $) -> {
+                if (!klass.isSingleton()) {
+                    if (klass.isIncluded() || klass.isPrepended()) {
+                        klass.concreteSubclasses(subs);
+                    } else {
+                        subs.append(klass);
+                    }
                 }
-                subs.add(klass);
-            }
+            });
         }
     }
 
@@ -1112,7 +1119,7 @@ public void addSubclass(RubyClass subclass) {
             synchronized (this) {
                 subclasses = this.subclasses;
                 if (subclasses == null) {
-                    this.subclasses = subclasses = new ConcurrentWeakHashMap<>(4, 0.75f, 1);
+                    this.subclasses = subclasses = Collections.synchronizedMap(new WeakIdentityHashMap(4));
                 }
             }
         }