Deduplicate capability provider lists

neoforge/src/main/java/org/embeddedt/modernfix/neoforge/caps/CapProviderGetter.java

@@ -0,0 +1,91 @@
+package org.embeddedt.modernfix.neoforge.caps;
+
+import it.unimi.dsi.fastutil.objects.ObjectOpenHashSet;
+import net.neoforged.neoforge.capabilities.BaseCapability;
+import net.neoforged.neoforge.capabilities.BlockCapability;
+import net.neoforged.neoforge.capabilities.EntityCapability;
+import net.neoforged.neoforge.capabilities.ICapabilityProvider;
+import net.neoforged.neoforge.capabilities.ItemCapability;
+
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
+import java.lang.reflect.Field;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+public class CapProviderGetter {
+    private static final MethodHandle BLOCK_CAP_PROVIDERS, ITEM_CAP_PROVIDERS, ENTITY_CAP_PROVIDERS;
+    private static final boolean FOUND_PROVIDERS;
+
+    static {
+        MethodHandle bProvider, iProvider, eProvider;
+        boolean found;
+        try {
+            bProvider = obtainForClass(BlockCapability.class);
+            iProvider = obtainForClass(ItemCapability.class);
+            eProvider = obtainForClass(EntityCapability.class);
+            found = true;
+        } catch(ReflectiveOperationException e) {
+            e.printStackTrace();
+            bProvider = null;
+            iProvider = null;
+            eProvider = null;
+            found = false;
+        }
+        BLOCK_CAP_PROVIDERS = bProvider;
+        ITEM_CAP_PROVIDERS = iProvider;
+        ENTITY_CAP_PROVIDERS = eProvider;
+        FOUND_PROVIDERS = found;
+    }
+
+    private static MethodHandle obtainForClass(Class<? extends BaseCapability> clz) throws ReflectiveOperationException{
+        Field field = clz.getDeclaredField("providers");
+        field.setAccessible(true);
+        return MethodHandles.publicLookup().unreflectGetter(field);
+    }
+
+    private static final Map<?, List<ICapabilityProvider<?, ?, ?>>> DUMMY_MAP = new HashMap<>();
+
+    public static <T extends BaseCapability> Map<?, List<ICapabilityProvider<?, ?, ?>>> getProviderMap(T cap) {
+        if(!FOUND_PROVIDERS) {
+            return DUMMY_MAP;
+        }
+
+        try {
+            if(cap instanceof BlockCapability<?,?> blockCap) {
+                return (Map<?, List<ICapabilityProvider<?, ?, ?>>>)BLOCK_CAP_PROVIDERS.invokeExact(blockCap);
+            } else if(cap instanceof ItemCapability<?,?> itemCap) {
+                return (Map<?, List<ICapabilityProvider<?, ?, ?>>>)ITEM_CAP_PROVIDERS.invokeExact(itemCap);
+            } else if(cap instanceof EntityCapability<?,?> entityCap) {
+                return (Map<?, List<ICapabilityProvider<?, ?, ?>>>)ENTITY_CAP_PROVIDERS.invokeExact(entityCap);
+            } else {
+                return DUMMY_MAP;
+            }
+        } catch(Throwable e) {
+            throw new RuntimeException("Couldn't get map", e);
+        }
+    }
+
+    public static void deduplicateCap(BaseCapability<?, ?> cap) {
+        var map = getProviderMap(cap);
+        if(map.isEmpty()) {
+            return;
+        }
+        ObjectOpenHashSet<List<ICapabilityProvider<?, ?, ?>>> capLists = new ObjectOpenHashSet<>();
+        int hits = 0;
+        for(Map.Entry<?, List<ICapabilityProvider<?, ?, ?>>> entry : map.entrySet()) {
+            var v = entry.getValue();
+            var canonicalList = capLists.get(v);
+            if(canonicalList == null) {
+                canonicalList = List.copyOf(v);
+                // This works because List.equals/hashCode is well-defined across any List implementation
+                capLists.add(canonicalList);
+            } else {
+                hits++;
+            }
+            entry.setValue(canonicalList);
+        }
+        //ModernFix.LOGGER.info("Deduplicated {}/{} lists", hits, map.size());
+    }
+}

neoforge/src/main/java/org/embeddedt/modernfix/neoforge/caps/ITrackingCapEvent.java

@@ -0,0 +1,9 @@
+package org.embeddedt.modernfix.neoforge.caps;
+
+import net.neoforged.neoforge.capabilities.BaseCapability;
+
+import java.util.Set;
+
+public interface ITrackingCapEvent {
+    Set<BaseCapability<?, ?>> mfix$getTrackedCaps();
+}

neoforge/src/main/java/org/embeddedt/modernfix/neoforge/mixin/perf/capability_list_compaction/CapabilityHooksMixin.java

@@ -0,0 +1,27 @@
+package org.embeddedt.modernfix.neoforge.mixin.perf.capability_list_compaction;
+
+import com.llamalad7.mixinextras.injector.wrapoperation.Operation;
+import com.llamalad7.mixinextras.injector.wrapoperation.WrapOperation;
+import net.neoforged.bus.api.Event;
+import net.neoforged.neoforge.capabilities.BaseCapability;
+import net.neoforged.neoforge.capabilities.CapabilityHooks;
+import org.embeddedt.modernfix.neoforge.caps.CapProviderGetter;
+import org.embeddedt.modernfix.neoforge.caps.ITrackingCapEvent;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.injection.At;
+
+@Mixin(value = CapabilityHooks.class, remap = false)
+public class CapabilityHooksMixin {
+    @WrapOperation(method = "init", at = @At(value = "INVOKE", target = "Lnet/neoforged/fml/ModLoader;postEventWrapContainerInModOrder(Lnet/neoforged/bus/api/Event;)V"))
+    private static void deduplicateCaps(Event event, Operation<Void> original) {
+        original.call(event);
+        if(event instanceof ITrackingCapEvent) {
+            //var stopwatch = Stopwatch.createStarted();
+            for(BaseCapability<?, ?> cap : ((ITrackingCapEvent)event).mfix$getTrackedCaps()) {
+                CapProviderGetter.deduplicateCap(cap);
+            }
+            //stopwatch.stop();
+            //ModernFix.LOGGER.info("Deduplicated capability lists in {}", stopwatch);
+        }
+    }
+}

neoforge/src/main/java/org/embeddedt/modernfix/neoforge/mixin/perf/capability_list_compaction/RegisterCapabilitiesEventMixin.java

@@ -0,0 +1,51 @@
+package org.embeddedt.modernfix.neoforge.mixin.perf.capability_list_compaction;
+
+import net.minecraft.world.entity.EntityType;
+import net.minecraft.world.level.ItemLike;
+import net.minecraft.world.level.block.Block;
+import net.minecraft.world.level.block.entity.BlockEntityType;
+import net.neoforged.neoforge.capabilities.BaseCapability;
+import net.neoforged.neoforge.capabilities.BlockCapability;
+import net.neoforged.neoforge.capabilities.EntityCapability;
+import net.neoforged.neoforge.capabilities.IBlockCapabilityProvider;
+import net.neoforged.neoforge.capabilities.ICapabilityProvider;
+import net.neoforged.neoforge.capabilities.ItemCapability;
+import net.neoforged.neoforge.capabilities.RegisterCapabilitiesEvent;
+import org.embeddedt.modernfix.neoforge.caps.ITrackingCapEvent;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
+
+import java.util.HashSet;
+import java.util.Set;
+
+@Mixin(value = RegisterCapabilitiesEvent.class, remap = false)
+public class RegisterCapabilitiesEventMixin implements ITrackingCapEvent {
+    private final Set<BaseCapability<?, ?>> mfix$trackedCapabilities = new HashSet<>();
+
+    @Inject(method = "registerBlock", at = @At("HEAD"))
+    private void trackBlockCap(BlockCapability<?, ?> capability, IBlockCapabilityProvider<?, ?> provider, Block[] blocks, CallbackInfo ci) {
+        mfix$trackedCapabilities.add(capability);
+    }
+
+    @Inject(method = "registerBlockEntity", at = @At("HEAD"))
+    private void trackBlockEntityCap(BlockCapability<?, ?> capability, BlockEntityType<?> type, ICapabilityProvider<?, ?, ?> provider, CallbackInfo ci) {
+        mfix$trackedCapabilities.add(capability);
+    }
+
+    @Inject(method = "registerItem", at = @At("HEAD"))
+    private void trackItemCap(ItemCapability<?, ?> capability, ICapabilityProvider<?, ?, ?> provider, ItemLike[] items, CallbackInfo ci) {
+        mfix$trackedCapabilities.add(capability);
+    }
+
+    @Inject(method = "registerEntity", at = @At("HEAD"))
+    private void trackEntityCap(EntityCapability<?, ?> capability, EntityType<?> type, ICapabilityProvider<?, ?, ?> provider, CallbackInfo ci) {
+        mfix$trackedCapabilities.add(capability);
+    }
+
+    @Override
+    public Set<BaseCapability<?, ?>> mfix$getTrackedCaps() {
+        return mfix$trackedCapabilities;
+    }
+}