Store delegates on registry objects to avoid map lookup

forge/src/main/java/org/embeddedt/modernfix/forge/mixin/perf/forge_registry_alloc/DelegateHolderMixin.java

@@ -0,0 +1,30 @@
+package org.embeddedt.modernfix.forge.mixin.perf.forge_registry_alloc;
+
+import net.minecraft.core.Holder;
+import net.minecraft.core.Registry;
+import net.minecraft.resources.ResourceKey;
+import net.minecraft.world.item.Item;
+import net.minecraft.world.level.block.Block;
+import org.embeddedt.modernfix.forge.registry.DelegateHolder;
+import org.spongepowered.asm.mixin.Mixin;
+
+@Mixin({Block.class, Item.class})
+public class DelegateHolderMixin<T> implements DelegateHolder<T> {
+    private Holder.Reference<T> mfix$delegate;
+    private ResourceKey<Registry<T>> mfix$key;
+
+    @Override
+    public Holder.Reference<T> mfix$getDelegate(ResourceKey<Registry<T>> registryKey) {
+        if(mfix$key == registryKey) {
+            return mfix$delegate;
+        } else {
+            return null;
+        }
+    }
+
+    @Override
+    public void mfix$setDelegate(ResourceKey<Registry<T>> registryKey, Holder.Reference<T> holder) {
+        this.mfix$delegate = holder;
+        this.mfix$key = registryKey;
+    }
+}

forge/src/main/java/org/embeddedt/modernfix/forge/mixin/perf/forge_registry_alloc/ForgeRegistryMixin.java

@@ -3,13 +3,19 @@ package org.embeddedt.modernfix.forge.mixin.perf.forge_registry_alloc;
 import it.unimi.dsi.fastutil.Hash;
 import it.unimi.dsi.fastutil.objects.Object2ObjectOpenHashMap;
 import net.minecraft.core.Holder;
+import net.minecraft.core.Registry;
 import net.minecraft.resources.ResourceKey;
 import net.minecraft.resources.ResourceLocation;
 import net.minecraftforge.registries.ForgeRegistry;
+import net.minecraftforge.registries.RegistryManager;
+import org.embeddedt.modernfix.forge.registry.DelegateHolder;
 import org.spongepowered.asm.mixin.Final;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.Overwrite;
 import org.spongepowered.asm.mixin.Shadow;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
 
 import java.util.Locale;
 import java.util.Map;
@@ -22,6 +28,10 @@ public abstract class ForgeRegistryMixin<V> {
 
     @Shadow @Final private Map<V, Holder.Reference<V>> delegatesByValue = new Object2ObjectOpenHashMap<>(Hash.DEFAULT_INITIAL_SIZE, 0.5F);
 
+    @Shadow public abstract ResourceKey<Registry<V>> getRegistryKey();
+
+    @Shadow @Final private RegistryManager stage;
+
     /**
      * @author embeddedt
      * @reason stop allocating so many unneeded objects. stop.
@@ -54,14 +64,34 @@ public abstract class ForgeRegistryMixin<V> {
 
     /**
      * @author embeddedt
-     * @reason see above
+     * @reason store delegates that are accessed extremely regularly on the registry entry itself, rather than
+     * going through a map lookup
+     */
+    @Inject(method = "bindDelegate", at = @At("RETURN"))
+    private void attachDelegate(ResourceKey<V> rkey, V value, CallbackInfoReturnable<Holder.Reference<V>> cir) {
+        // Only attach delegates for the ACTIVE registry. The Forge registry system is weird and seems to keep multiple
+        // copies of itself at once.
+        if(this.stage == RegistryManager.ACTIVE && value instanceof DelegateHolder<?>) {
+            ((DelegateHolder<V>)value).mfix$setDelegate(this.getRegistryKey(), cir.getReturnValue());
+        }
+    }
+
+    /**
+     * @author embeddedt
+     * @reason skip map lookup for hot delegates, avoid allocations otherwise
      */
     @Overwrite
     public Holder.Reference<V> getDelegateOrThrow(V value) {
-        Holder.Reference<V> holder = delegatesByValue.get(value);
+        Holder.Reference<V> holder = null;
+        if (this.stage == RegistryManager.ACTIVE && value instanceof DelegateHolder<?>) {
+            holder = ((DelegateHolder<V>)value).mfix$getDelegate(this.getRegistryKey());
+        }
 
         if (holder == null) {
-            throw new IllegalArgumentException(String.format(Locale.ENGLISH, "No delegate exists for value %s", value));
+            holder = delegatesByValue.get(value);
+            if (holder == null) {
+                throw new IllegalArgumentException(String.format(Locale.ENGLISH, "No delegate exists for value %s", value));
+            }
         }
 
         return holder;

forge/src/main/java/org/embeddedt/modernfix/forge/registry/DelegateHolder.java

@@ -0,0 +1,10 @@
+package org.embeddedt.modernfix.forge.registry;
+
+import net.minecraft.core.Holder;
+import net.minecraft.core.Registry;
+import net.minecraft.resources.ResourceKey;
+
+public interface DelegateHolder<T> {
+    Holder.Reference<T> mfix$getDelegate(ResourceKey<Registry<T>> registryKey);
+    void mfix$setDelegate(ResourceKey<Registry<T>> registryKey, Holder.Reference<T> holder);
+}