Partially restore allModules set

All Class instances in the system are traversible by walking down
from BasicObject, but natural Module instances (i.e. all `module`
objects) do not descend from BasicObject. We partially restore the
allModules set in order to track these objects for each_object
when called with Module.

core/src/main/java/org/jruby/Ruby.java

@@ -1345,15 +1345,32 @@ public int allocModuleId() {
         return moduleLastId.incrementAndGet();
     }
 
-    @Deprecated
+    /**
+     * A collection of all natural Module instances in the system.
+     *
+     * Instances of Module, which are themselves modules, do not have ancestors and can't be traversed by walking down
+     * from BasicObject. We track them separately here for purposes of ObjectSpace.each_object.
+     *
+     * @param module the true module to add to the allModules collection
+     */
     public void addModule(RubyModule module) {
-        // ignored
+        assert module.getMetaClass() == moduleClass;
+
+        allModules.put(module, RubyBasicObject.NEVER);
     }
 
-    @Deprecated
+    /**
+     * Walk all natural Module instances in the system.
+     *
+     * This will only include direct instances of Module, not instances of Class.
+     *
+     * @param func the consumer to call for each module
+     */
     public void eachModule(Consumer<RubyModule> func) {
-        // walk all subclasses starting from BasicObject
-        basicObjectClass.eachDescendant(func);
+        Enumeration<RubyModule> e = allModules.keys();
+        while (e.hasMoreElements()) {
+            func.accept(e.nextElement());
+        }
     }
 
     @Deprecated(since = "10.0")
@@ -3406,6 +3423,7 @@ private void systemTeardown(final ThreadContext context) {
 
         // Clear runtime tables to aid GC
         boundMethods.clear();
+        allModules.clear();
         constantNameInvalidators.clear();
         symbolTable.clear();
         javaSupport = loadJavaSupport();
@@ -5815,6 +5833,11 @@ public void warn(String message) {
     private final AtomicInteger symbolLastId = new AtomicInteger(128);
     private final AtomicInteger moduleLastId = new AtomicInteger(0);
 
+    // Weak map of all natural instances of Module in the system (not including Classes).
+    // a ConcurrentMap<RubyModule, ?> is used to emulate WeakHashSet<RubyModule>
+    // NOTE: module instances are unique and we only addModule from <init> - could use a ConcurrentLinkedQueue
+    private final ConcurrentWeakHashMap<RubyModule, Object> allModules = new ConcurrentWeakHashMap<>(128);
+
     private final Map<String, DateTimeZone> timeZoneCache = new HashMap<>();
 
     /**

core/src/main/java/org/jruby/RubyClass.java

@@ -1136,13 +1136,6 @@ public Collection<RubyClass> subclasses(boolean includeDescendants) {
         return mine;
     }
 
-    public void eachDescendant(Consumer<? super RubyClass> consumer) {
-        getSubclassesForRead().forEachClass((k) -> {
-            consumer.accept(k);
-            k.eachDescendant(consumer);
-        });
-    }
-
     private SubclassArray newConcreteSubclassesArray(ThreadContext context) {
         return new SubclassArray(context.runtime, this.concreteSubclassesEstimate);
     }
@@ -1166,7 +1159,7 @@ void addAllSubclasses(SubclassList mine) {
         getSubclassesForRead().forEachClass(mine);
     }
 
-    private RubyClassSet getSubclassesForRead() {
+    public RubyClassSet getSubclassesForRead() {
         RubyClassSet result = this.subclasses;
 
         if (result != null) {

core/src/main/java/org/jruby/RubyModule.java

@@ -424,7 +424,11 @@ protected RubyModule(Ruby runtime, RubyClass metaClass, boolean objectSpace) {
 
         id = runtime.allocModuleId();
 
-        runtime.addModule(this);
+        // track module instances separately, since they don't descend from BasicObject
+        if (metaClass == runtime.getModule()) {
+            runtime.addModule(this);
+        }
+
         // if (parent == null) parent = runtime.getObject();
         setFlag(NEEDSIMPL_F, !isClass());
         updateGeneration(runtime);

core/src/main/java/org/jruby/RubyObjectSpace.java

@@ -32,9 +32,9 @@
 
 package org.jruby;
 
-import java.util.ArrayList;
 import static org.jruby.RubyEnumerator.enumeratorize;
 
+import java.util.HashSet;
 import java.util.Iterator;
 import java.util.Map;
 import java.util.WeakHashMap;
@@ -175,23 +175,30 @@ public static IRubyObject each_objectInternal(final ThreadContext context, IRuby
         }
         if (rubyClass == classClass(context) || rubyClass == moduleClass(context)) {
 
-            final ArrayList<IRubyObject> modules = new ArrayList<>(96);
-            runtime.eachModule((module) -> {
-                    if (rubyClass.isInstance(module)) {
-                        if (!(module instanceof IncludedModule || module instanceof PrependedModule 
-                                || module == runtime.getJavaSupport().getJavaPackageClass() || module instanceof JavaPackage
-                                || (module instanceof MetaClass && (((MetaClass)module).getAttached() instanceof JavaPackage)))) {
-                            // do nothing for included wrappers or singleton classes
-                            modules.add(module); // store the module to avoid concurrent modification exceptions
-                        }
+            // Use set in case there's any overlaps between eachModule and BasicObject descendants (should not be).
+            final HashSet<IRubyObject> modules = new HashSet<>(96);
+
+            // if Module is requested, iterate over true modules as well as all classes
+            if (rubyClass == moduleClass(context)) {
+                runtime.eachModule((module) -> {
+                    modules.add(module); // store the module to avoid concurrent modification exceptions
+                });
+            }
+
+            // Iterate over all classes
+            runtime.getBasicObject().subclasses(true).forEach((module) -> {
+                if (rubyClass.isInstance(module)) {
+                    if (!(module instanceof IncludedModule || module instanceof PrependedModule
+                            || module == runtime.getJavaSupport().getJavaPackageClass()
+                            || (module instanceof MetaClass && (((MetaClass) module).getAttached() instanceof JavaPackage)))) {
+                        // do nothing for included wrappers or singleton classes
+                        modules.add(module); // store the module to avoid concurrent modification exceptions
                     }
+                }
             });
 
-            final int count = modules.size();
-            for (int i = 0; i<count; i++) {
-                block.yield(context, modules.get(i));
-            }
-            return asFixnum(context, count);
+            modules.forEach((obj) -> block.yield(context, obj));
+            return asFixnum(context, modules.size());
         }
         if (rubyClass.getClass() == MetaClass.class) {
             // each_object(Cls.singleton_class) is basically a walk of Cls and all descendants of Cls.