Modify how ModLoadingStage handles dispatch to Registry Events. Stops the objectholder spam and generally improves performance on larger packs.

Signed-off-by: cpw <cpw+github@weeksfamily.ca>

src/main/java/net/minecraftforge/fml/ModContainer.java

@@ -19,6 +19,7 @@
 
 package net.minecraftforge.fml;
 
+import cpw.mods.modlauncher.api.LamdbaExceptionUtils;
 import net.minecraftforge.eventbus.api.Event;
 import net.minecraftforge.fml.config.ModConfig;
 import net.minecraftforge.fml.event.lifecycle.IModBusEvent;
@@ -102,14 +103,13 @@ public abstract class ModContainer
     public static <T extends Event & IModBusEvent> CompletableFuture<Void> buildTransitionHandler(
             final ModContainer target,
             final ModLoadingStage.EventGenerator<T> eventGenerator,
-            final ModLoadingStage.EventDispatcher<T> eventDispatcher,
             final BiFunction<ModLoadingStage, Throwable, ModLoadingStage> stateChangeHandler,
             final Executor executor) {
         return CompletableFuture
                 .runAsync(() -> {
                     ModLoadingContext.get().setActiveContainer(target, target.contextExtension.get());
                     target.activityMap.getOrDefault(target.modLoadingStage, ()->{}).run();
-                    eventDispatcher.apply(target::acceptEvent).accept(eventGenerator.apply(target));
+                    target.acceptEvent(eventGenerator.apply(target));
                 }, executor)
                 .whenComplete((mc, exception) -> {
                     target.modLoadingStage = stateChangeHandler.apply(target.modLoadingStage, exception);

src/main/java/net/minecraftforge/fml/ModList.java

@@ -120,11 +120,10 @@ public class ModList
 
     <T extends Event & IModBusEvent> Function<Executor, CompletableFuture<List<Throwable>>> futureVisitor(
             final ModLoadingStage.EventGenerator<T> eventGenerator,
-            final ModLoadingStage.EventDispatcher<T> eventManager,
             final BiFunction<ModLoadingStage, Throwable, ModLoadingStage> stateChange) {
         return executor -> gather(
                 this.mods.stream()
-                .map(m -> ModContainer.buildTransitionHandler(m, eventGenerator, eventManager, stateChange, executor))
+                .map(mod -> ModContainer.buildTransitionHandler(mod, eventGenerator, stateChange, executor))
                 .collect(Collectors.toList()))
             .thenComposeAsync(ModList::completableFutureFromExceptionList, executor);
     }

src/main/java/net/minecraftforge/fml/ModLoadingStage.java

@@ -34,6 +34,7 @@ import net.minecraftforge.fml.event.lifecycle.ParallelDispatchEvent;
 import net.minecraftforge.registries.GameData;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
@@ -50,8 +51,8 @@ public enum ModLoadingStage
     ERROR(),
     VALIDATE(),
     CONSTRUCT(FMLConstructModEvent.class),
-    CREATE_REGISTRIES(()->Stream.of(EventGenerator.fromFunction(RegistryEvent.NewRegistry::new)), EventDispatcher.identity()),
+    CREATE_REGISTRIES(()->Stream.of(EventGenerator.fromFunction(RegistryEvent.NewRegistry::new))),
-    LOAD_REGISTRIES(GameData::generateRegistryEvents, GameData.buildRegistryEventDispatch()),
+    LOAD_REGISTRIES(GameData::generateRegistryEvents, GameData::preRegistryEventDispatch, GameData::postRegistryEventDispatch),
     COMMON_SETUP(FMLCommonSetupEvent.class),
     SIDED_SETUP(DistExecutor.unsafeRunForDist(()->()->FMLClientSetupEvent.class, ()->()->FMLDedicatedServerSetupEvent.class)),
     ENQUEUE_IMC(InterModEnqueueEvent.class),
@@ -60,30 +61,37 @@ public enum ModLoadingStage
     DONE();
 
     private final Supplier<Stream<EventGenerator<?>>> eventFunctionStream;
-    private final EventDispatcher<?> eventManager;
     private final Optional<Class<? extends ParallelDispatchEvent>> parallelEventClass;
     private final ThreadSelector threadSelector;
     private final BiFunction<Executor, CompletableFuture<List<Throwable>>, CompletableFuture<List<Throwable>>> finalActivityGenerator;
+    private final BiFunction<Executor, ? extends EventGenerator<?>, CompletableFuture<List<Throwable>>> preDispatchHook;
+    private final BiFunction<Executor, ? extends EventGenerator<?>, CompletableFuture<List<Throwable>>> postDispatchHook;
     private DeferredWorkQueue deferredWorkQueue;
 
     ModLoadingStage(Class<? extends ParallelDispatchEvent> parallelClass) {
         final EventGenerator<?> event = EventGenerator.fromFunction(LamdbaExceptionUtils.rethrowFunction((ModContainer mc) -> parallelClass.getConstructor(ModContainer.class).newInstance(mc)));
         this.eventFunctionStream = ()->Stream.of(event);
         this.threadSelector = ThreadSelector.PARALLEL;
-        this.eventManager = EventDispatcher.identity();
         this.parallelEventClass = Optional.of(parallelClass);
         deferredWorkQueue = new DeferredWorkQueue(this, parallelClass);
         this.finalActivityGenerator = (e, prev) -> prev.thenApplyAsync((List<Throwable> t) -> {
             deferredWorkQueue.runTasks();
             return t;
         }, e);
+        this.preDispatchHook = (t,f)->CompletableFuture.completedFuture(Collections.emptyList());
+        this.postDispatchHook = (t,f)->CompletableFuture.completedFuture(Collections.emptyList());
     }
 
-    <T extends Event & IModBusEvent> ModLoadingStage(Supplier<Stream<EventGenerator<?>>> eventStream, EventDispatcher<?> eventManager) {
+    ModLoadingStage(Supplier<Stream<EventGenerator<?>>> eventStream) {
+        this(eventStream, (t,f)->CompletableFuture.completedFuture(Collections.emptyList()), (t,f)->CompletableFuture.completedFuture(Collections.emptyList()));
+    }
+    @SuppressWarnings("unchecked")
+    <T extends Event & IModBusEvent> ModLoadingStage(Supplier<Stream<EventGenerator<?>>> eventStream, final BiFunction<Executor, ? extends EventGenerator<T>, CompletableFuture<List<Throwable>>> preDispatchHook, final BiFunction<Executor, ? extends EventGenerator<T>, CompletableFuture<List<Throwable>>> postDispatchHook) {
         this.eventFunctionStream = eventStream;
         this.parallelEventClass = Optional.empty();
-        this.eventManager = eventManager;
         this.threadSelector = ThreadSelector.SYNC;
+        this.preDispatchHook = preDispatchHook;
+        this.postDispatchHook = postDispatchHook;
         this.finalActivityGenerator = (e, prev) ->prev.thenApplyAsync(Function.identity(), e);
     }
 
@@ -97,17 +105,10 @@ public enum ModLoadingStage
     @SuppressWarnings("unchecked")
     public <T extends Event & IModBusEvent> CompletableFuture<List<Throwable>> buildTransition(final Executor syncExecutor, final Executor parallelExecutor, Function<Executor, CompletableFuture<Void>> preSyncTask, Function<Executor, CompletableFuture<Void>> postSyncTask) {
         List<CompletableFuture<List<Throwable>>> cfs = new ArrayList<>();
-        EventDispatcher<T> em = (EventDispatcher<T>) this.eventManager;
+        this.eventFunctionStream.get()
-        eventFunctionStream.get()
                 .map(f->(EventGenerator<T>)f)
-                .reduce((head, tail)->{
+                .reduce((head, tail)-> addCompletableFutureTaskForModDispatch(syncExecutor, parallelExecutor, cfs, head, ModLoadingStage::currentState, tail))
-                    cfs.add(ModList.get()
+                .ifPresent(last-> addCompletableFutureTaskForModDispatch(syncExecutor, parallelExecutor, cfs, last, ModLoadingStage::nextState, null));
-                            .futureVisitor(head, em, ModLoadingStage::currentState)
-                            .apply(threadSelector.apply(syncExecutor, parallelExecutor))
-                    );
-                    return tail;
-                })
-                .ifPresent(last->cfs.add(ModList.get().futureVisitor(last, em, ModLoadingStage::nextState).apply(threadSelector.apply(syncExecutor, parallelExecutor))));
         final CompletableFuture<Void> preSyncTaskCF = preSyncTask.apply(syncExecutor);
         final CompletableFuture<List<Throwable>> eventDispatchCF = ModList.gather(cfs).thenCompose(ModList::completableFutureFromExceptionList);
         final CompletableFuture<List<Throwable>> postEventDispatchCF = preSyncTaskCF.thenComposeAsync(n -> eventDispatchCF, parallelExecutor).thenApply(r -> {
@@ -117,6 +118,14 @@ public enum ModLoadingStage
         return this.finalActivityGenerator.apply(syncExecutor, postEventDispatchCF);
     }
 
+    @SuppressWarnings("unchecked")
+    private <T extends Event & IModBusEvent> EventGenerator<T> addCompletableFutureTaskForModDispatch(final Executor syncExecutor, final Executor parallelExecutor, final List<CompletableFuture<List<Throwable>>> completeableFutures, final EventGenerator<T> eventGenerator, BiFunction<ModLoadingStage, Throwable, ModLoadingStage> nextState, final EventGenerator<T> nextGenerator) {
+        completeableFutures.add(((BiFunction<Executor, EventGenerator<T>, CompletableFuture<List<Throwable>>>)preDispatchHook).apply(threadSelector.apply(syncExecutor, parallelExecutor), eventGenerator));
+        completeableFutures.add(ModList.get().futureVisitor(eventGenerator, nextState).apply(threadSelector.apply(syncExecutor, parallelExecutor)));
+        completeableFutures.add(((BiFunction<Executor, EventGenerator<T>, CompletableFuture<List<Throwable>>>)postDispatchHook).apply(threadSelector.apply(syncExecutor, parallelExecutor), eventGenerator));
+        return nextGenerator;
+    }
+
     ModLoadingStage nextState(Throwable exception) {
         return exception != null ? ERROR : values()[this.ordinal()+1];
     }
@@ -153,9 +162,4 @@ public enum ModLoadingStage
             return fn::apply;
         }
     }
-    public interface EventDispatcher<T extends Event & IModBusEvent> extends Function<Consumer<? super T>, Consumer<? super T>> {
-        static <FN extends Event & IModBusEvent> EventDispatcher<FN> identity() {
-            return consumer -> consumer;
-        }
-    }
 }

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

@@ -22,8 +22,8 @@ package net.minecraftforge.registries;
 import com.google.common.collect.*;
 import com.mojang.serialization.Lifecycle;
 
-import java.util.HashSet;
+import java.util.*;
-import java.util.Set;
+
 import net.minecraft.block.AirBlock;
 import net.minecraft.block.Block;
 import net.minecraft.block.BlockState;
@@ -92,12 +92,8 @@ import org.apache.logging.log4j.Logger;
 import javax.annotation.Nullable;
 
 import java.lang.reflect.Field;
-import java.util.Collection;
+import java.util.concurrent.CompletableFuture;
-import java.util.IdentityHashMap;
+import java.util.concurrent.Executor;
-import java.util.LinkedHashMap;
-import java.util.List;
-import java.util.Locale;
-import java.util.Map;
 import java.util.function.BiConsumer;
 import java.util.function.BiFunction;
 import java.util.function.Function;
@@ -346,22 +342,30 @@ public class GameData
         keys.add(0, BLOCKS.func_240901_a_());
         keys.add(1, ITEMS.func_240901_a_());
 
-        final Function<ResourceLocation, ? extends RegistryEvent.Register<?>> modContainerEventGeneratorFunction = rl -> RegistryManager.ACTIVE.getRegistry(rl).getRegisterEvent(rl);
+        final Function<ResourceLocation, ? extends RegistryEvent.Register<?>> registerEventGenerator = rl -> RegistryManager.ACTIVE.getRegistry(rl).getRegisterEvent(rl);
-        return keys.stream().map(rl -> ModLoadingStage.EventGenerator.fromFunction(mc -> modContainerEventGeneratorFunction.apply(rl)));
+        return keys.stream().map(rl -> ModLoadingStage.EventGenerator.fromFunction(mc -> registerEventGenerator.apply(rl)));
     }
 
-    public static ModLoadingStage.EventDispatcher<RegistryEvent.Register<?>> buildRegistryEventDispatch() {
+    public static CompletableFuture<List<Throwable>> preRegistryEventDispatch(final Executor executor, final ModLoadingStage.EventGenerator<? extends RegistryEvent.Register<?>> eventGenerator) {
-        return eventConsumer -> eventToSend -> {
+        return CompletableFuture.runAsync(()-> {
-            final ResourceLocation rl = eventToSend.getName();
+                    final RegistryEvent.Register<?> event = eventGenerator.apply(null);
-            ForgeRegistry<?> fr = (ForgeRegistry<?>) eventToSend.getRegistry();
+                    final ResourceLocation rl = event.getName();
-            StartupMessageManager.modLoaderConsumer().ifPresent(s->s.accept("REGISTERING "+rl));
+                    ForgeRegistry<?> fr = (ForgeRegistry<?>) event.getRegistry();
-            fr.unfreeze();
+                    StartupMessageManager.modLoaderConsumer().ifPresent(s -> s.accept("REGISTERING " + rl));
-            eventConsumer.accept(eventToSend);
+                    fr.unfreeze();
+                }, executor).thenApply(v->Collections.emptyList());
+    }
+    
+    public static CompletableFuture<List<Throwable>> postRegistryEventDispatch(final Executor executor, final ModLoadingStage.EventGenerator<? extends RegistryEvent.Register<?>> eventGenerator) {
+        return CompletableFuture.runAsync(()-> {
+            final RegistryEvent.Register<?> event = eventGenerator.apply(null);
+            final ResourceLocation rl = event.getName();
+            ForgeRegistry<?> fr = (ForgeRegistry<?>) event.getRegistry();
             fr.freeze();
-            LOGGER.debug(REGISTRIES,"Applying holder lookups: {}", rl.toString());
+            LOGGER.debug(REGISTRIES, "Applying holder lookups: {}", rl.toString());
             ObjectHolderRegistry.applyObjectHolders(rl::equals);
-            LOGGER.debug(REGISTRIES,"Holder lookups applied: {}", rl.toString());
+            LOGGER.debug(REGISTRIES, "Holder lookups applied: {}", rl.toString());
-        };
+        }, executor).thenApply(v->Collections.emptyList());
     }
 
     public static void setCustomTagTypesFromRegistries()