Store object holder handlers per registry name (#10899)

src/main/java/net/minecraftforge/registries/GameData.java

@@ -350,7 +350,7 @@ public class GameData {
                 if (forgeRegistry != null)
                     forgeRegistry.freeze();
                 LOGGER.debug(REGISTRIES, "Applying holder lookups: {}", registryKey.identifier());
-                ObjectHolderRegistry.applyObjectHolders(registryKey.identifier()::equals);
+                ObjectHolderRegistry.applyObjectHolders(registryKey.identifier());
                 LOGGER.debug(REGISTRIES, "Holder lookups applied: {}", registryKey.identifier());
             } catch (Throwable t) {
                 aggregate.addSuppressed(t);

src/main/java/net/minecraftforge/registries/ObjectHolderRegistry.java

@@ -5,10 +5,12 @@
 
 package net.minecraftforge.registries;
 
+import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
 import java.util.Set;
-import java.util.function.Consumer;
-import java.util.function.Predicate;
 
 import net.minecraft.resources.Identifier;
 import org.apache.logging.log4j.LogManager;
@@ -16,32 +18,34 @@ import org.apache.logging.log4j.Logger;
 import org.jetbrains.annotations.ApiStatus;
 
 /**
- * Internal registry for tracking {@link ObjectHolder} references
+ * Internal registry for tracking {@link RegistryObject} references
  */
 @ApiStatus.Internal
 class ObjectHolderRegistry {
     /**
-     * Exposed to allow modders to register their own notification handlers.
-     * This runnable will be called after a registry snapshot has been injected and finalized.
-     * The internal list is backed by a HashSet so it is HIGHLY recommended you implement a proper equals
-     * and hashCode function to de-duplicate callers here.
-     * The default @ObjectHolder implementation uses the hashCode/equals for the field the annotation is on.
+     * Registers a handler that is run after the named registry has been injected and finalized.
+     * Handlers are kept in a HashSet per registry, so the same instance is only ever registered once.
      */
-    static synchronized void addHandler(Consumer<Predicate<Identifier>> ref) {
-        objectHolders.add(ref);
+    static synchronized void addHandler(Identifier registry, Runnable ref) {
+        namedHandlers.computeIfAbsent(registry, id -> new HashSet<>()).add(ref);
     }
 
     /**
-     * Removed the specified handler from the notification list.
-     *
-     * The internal list is backed by a hash set, and so proper hashCode and equals operations are required for success.
-     *
-     * The default @ObjectHolder implementation uses the hashCode/equals for the field the annotation is on.
+     * Removes a handler registered with {@link #addHandler(Identifier, Runnable)}.
      *
      * @return true if handler was matched and removed.
      */
-    static synchronized boolean removeHandler(Consumer<Predicate<Identifier>> ref) {
-        return objectHolders.remove(ref);
+    static synchronized boolean removeHandler(Identifier registry, Runnable ref) {
+        var handlers = namedHandlers.get(registry);
+        return handlers != null && handlers.remove(ref);
+    }
+
+    /**
+     * Registers a check that is run once, before the handlers of the next holder pass.
+     * Used to report holders that point at a registry that does not exist.
+     */
+    static synchronized void addValidator(Runnable ref) {
+        validators.add(ref);
     }
 
     //==============================================================
@@ -49,12 +53,19 @@ class ObjectHolderRegistry {
     //==============================================================
 
     private static final Logger LOGGER = LogManager.getLogger();
-    private static final Set<Consumer<Predicate<Identifier>>> objectHolders = new HashSet<>();
+    // Order is not guaranteed, the old implementation kept every handler in one HashSet.
+    private static final Map<Identifier, Set<Runnable>> namedHandlers = new HashMap<>();
+    private static final List<Runnable> validators = new ArrayList<>();
 
+    /** Runs the handlers of every registry. */
     static void applyObjectHolders() {
         try {
             LOGGER.debug(ForgeRegistry.REGISTRIES, "Applying holder lookups");
-            applyObjectHolders(key -> true);
+            var exceptions = new ArrayList<Exception>();
+            validate(exceptions);
+            for (var handlers : namedHandlers.values())
+                apply(handlers, exceptions);
+            throwIfNeeded(exceptions);
             LOGGER.debug(ForgeRegistry.REGISTRIES, "Holder lookups applied");
         } catch (RuntimeException e) {
             // It is more important that the calling contexts continue without exception to prevent further cascading errors
@@ -62,17 +73,52 @@ class ObjectHolderRegistry {
         }
     }
 
-    static void applyObjectHolders(Predicate<Identifier> filter) {
-        RuntimeException aggregate = new RuntimeException("Failed to apply some object holders, see suppressed exceptions for details");
-        for (Consumer<Predicate<Identifier>> objectHolder : objectHolders) {
-            try {
-                objectHolder.accept(filter);
-            } catch (Exception e) {
-                aggregate.addSuppressed(e);
-            }
+    /** Runs the handlers of one registry. */
+    static void applyObjectHolders(Identifier registry) {
+        var exceptions = new ArrayList<Exception>();
+        validate(exceptions);
+        var handlers = namedHandlers.get(registry);
+        if (handlers != null)
+            apply(handlers, exceptions);
+        throwIfNeeded(exceptions);
+    }
+
+    // Each validator runs once, at the first holder pass after it was registered.
+    private static void validate(List<Exception> exceptions) {
+        List<Runnable> pending;
+        synchronized (ObjectHolderRegistry.class) {
+            if (validators.isEmpty())
+                return;
+            pending = new ArrayList<>(validators);
+            validators.clear();
         }
 
-        if (aggregate.getSuppressed().length > 0)
-            throw aggregate;
+        for (var validator : pending) {
+            try {
+                validator.run();
+            } catch (Exception e) {
+                exceptions.add(e);
+            }
+        }
+    }
+
+    private static void apply(Set<Runnable> handlers, List<Exception> exceptions) {
+        for (var handler : handlers) {
+            try {
+                handler.run();
+            } catch (Exception e) {
+                exceptions.add(e);
+            }
+        }
+    }
+
+    private static void throwIfNeeded(List<Exception> exceptions) {
+        if (exceptions.isEmpty())
+            return;
+
+        var aggregate = new RuntimeException("Failed to apply some object holders, see suppressed exceptions for details");
+        for (var e : exceptions)
+            aggregate.addSuppressed(e);
+        throw aggregate;
     }
 }

src/main/java/net/minecraftforge/registries/RegistryObject.java

@@ -27,7 +27,6 @@ public final class RegistryObject<T> implements Supplier<T> {
     private final Identifier name;
     @Nullable
     private final ResourceKey<T> key;
-    private final boolean optionalRegistry;
     @Nullable
     private T value;
     @Nullable
@@ -129,7 +128,6 @@ public final class RegistryObject<T> implements Supplier<T> {
     private RegistryObject() {
         this.name = null;
         this.key = null;
-        this.optionalRegistry = false;
     }
 
     @SuppressWarnings("unchecked")
@@ -138,40 +136,21 @@ public final class RegistryObject<T> implements Supplier<T> {
             throw new IllegalArgumentException("Invalid registry argument, must not be null");
         this.name = name;
         this.key = (ResourceKey<T>) ResourceKey.create(registry.getRegistryKey(), name);
-        this.optionalRegistry = false;
-        ObjectHolderRegistry.addHandler(pred -> {
-            if (pred.test(registry.getRegistryName()))
-                this.updateReference((IForgeRegistry<? extends T>) registry);
-        });
+        ObjectHolderRegistry.addHandler(registry.getRegistryName(), () -> this.updateReference((IForgeRegistry<? extends T>) registry));
         this.updateReference((IForgeRegistry<? extends T>) registry);
     }
 
     private RegistryObject(final Identifier name, final Identifier registryName, final String modid, boolean optionalRegistry) {
         this.name = name;
         this.key = ResourceKey.create(ResourceKey.createRegistryKey(registryName), name);
-        this.optionalRegistry = optionalRegistry;
-        final Throwable callerStack = new Throwable("Calling Site from mod: " + modid);
-        ObjectHolderRegistry.addHandler(new Consumer<>() {
-            private boolean registryExists = false;
-            private boolean invalidRegistry = false;
-
-            @Override
-            public void accept(Predicate<Identifier> pred) {
-                if (invalidRegistry)
-                    return;
-
-                if (!RegistryObject.this.optionalRegistry && !registryExists) {
-                    if (!registryExists(registryName)) {
-                        invalidRegistry = true;
-                        throw new IllegalStateException("Unable to find registry with key " + registryName + " for mod \"" + modid + "\". Check the 'caused by' to see further stack.", callerStack);
-                    }
-                    registryExists = true;
-                }
-
-                if (pred.test(registryName))
-                    RegistryObject.this.updateReference(registryName);
-            }
-        });
+        if (!optionalRegistry) {
+            final Throwable callerStack = new Throwable("Calling Site from mod: " + modid);
+            ObjectHolderRegistry.addValidator(() -> {
+                if (!registryExists(registryName))
+                    throw new IllegalStateException("Unable to find registry with key " + registryName + " for mod \"" + modid + "\". Check the 'caused by' to see further stack.", callerStack);
+            });
+        }
+        ObjectHolderRegistry.addHandler(registryName, () -> this.updateReference(registryName));
         this.updateReference(registryName);
     }
 
@@ -181,7 +160,6 @@ public final class RegistryObject<T> implements Supplier<T> {
         @SuppressWarnings("unchecked")
         var typed = (ResourceKey<Registry<T>>)registry;
         this.key = ResourceKey.create(typed, name);
-        this.optionalRegistry = true;
     }
 
     /**