Reimplement optimized lookup map for block state models

Author: embeddedt

src/main/java/org/embeddedt/modernfix/common/mixin/perf/dynamic_resources/MixinBlockState.java

@@ -0,0 +1,26 @@
+package org.embeddedt.modernfix.common.mixin.perf.dynamic_resources;
+
+import net.minecraft.client.renderer.block.model.BlockStateModel;
+import net.minecraft.world.level.block.state.BlockBehaviour;
+import org.embeddedt.modernfix.annotation.ClientOnlyMixin;
+import org.embeddedt.modernfix.duck.IModelHoldingBlockState;
+import org.spongepowered.asm.mixin.Mixin;
+
+import java.lang.ref.SoftReference;
+
+@Mixin(BlockBehaviour.BlockStateBase.class)
+@ClientOnlyMixin
+public class MixinBlockState implements IModelHoldingBlockState {
+    private volatile SoftReference<BlockStateModel> mfix$model;
+
+    @Override
+    public BlockStateModel mfix$getModel() {
+        var ref = mfix$model;
+        return ref != null ? ref.get() : null;
+    }
+
+    @Override
+    public void mfix$setModel(BlockStateModel model) {
+        mfix$model = model != null ? new SoftReference<>(model) : null;
+    }
+}

src/main/java/org/embeddedt/modernfix/common/mixin/perf/dynamic_resources/MixinModelManager.java

@@ -12,6 +12,7 @@
 import net.minecraft.server.packs.resources.Resource;
 import net.minecraft.world.level.block.state.BlockState;
 import org.embeddedt.modernfix.annotation.ClientOnlyMixin;
+import org.embeddedt.modernfix.dynresources.BlockStateModelMap;
 import org.embeddedt.modernfix.dynresources.DynamicModelSystem;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.Overwrite;
@@ -62,8 +63,7 @@ private static Object2IntMap<BlockState> buildModelGroups(BlockColors blockColor
      */
     @Overwrite
     private static Map<BlockState, BlockStateModel> createBlockStateToModelDispatch(Map<BlockState, BlockStateModel> blockStateModels, BlockStateModel missingModel) {
-        return Maps.asMap(DynamicModelSystem.getAllBlockStates(), state -> {
-            return blockStateModels.getOrDefault(state, missingModel);
-        });
+        BlockStateModelMap.resetCache();
+        return new BlockStateModelMap(blockStateModels, missingModel);
     }
 }

src/main/java/org/embeddedt/modernfix/dynresources/BlockStateModelMap.java

@@ -0,0 +1,105 @@
+package org.embeddedt.modernfix.dynresources;
+
+import net.minecraft.client.renderer.block.model.BlockStateModel;
+import net.minecraft.world.level.block.Block;
+import net.minecraft.world.level.block.state.BlockState;
+import org.embeddedt.modernfix.duck.IModelHoldingBlockState;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.util.Collection;
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * Optimized blockstate->model dispatch map that stores the models directly on the block state objects using a helper
+ * field. This relies on the fact that Minecraft should only have one model map in flight at a time.
+ */
+public record BlockStateModelMap(Map<BlockState, BlockStateModel> modelMap,
+                                 BlockStateModel fallbackModel) implements Map<BlockState, BlockStateModel> {
+
+    @Override
+    public int size() {
+        return Block.BLOCK_STATE_REGISTRY.size();
+    }
+
+    @Override
+    public boolean isEmpty() {
+        return false;
+    }
+
+    @Override
+    public boolean containsKey(Object o) {
+        return o instanceof BlockState;
+    }
+
+    @Override
+    public boolean containsValue(Object o) {
+        return modelMap.containsValue(o);
+    }
+
+    @Override
+    public BlockStateModel get(Object o) {
+        if (o instanceof IModelHoldingBlockState modelHolder) {
+            BlockStateModel model = modelHolder.mfix$getModel();
+
+            if(model != null) {
+                return model;
+            }
+
+            model = modelMap.getOrDefault(o, fallbackModel);
+            modelHolder.mfix$setModel(model);
+            return model;
+        } else {
+            return modelMap.getOrDefault(o, fallbackModel);
+        }
+    }
+
+    @Override
+    public @Nullable BlockStateModel put(BlockState blockState, BlockStateModel blockStateModel) {
+        var oldModel = modelMap.put(blockState, blockStateModel);
+        ((IModelHoldingBlockState)blockState).mfix$setModel(null);
+        return oldModel;
+    }
+
+    @Override
+    public BlockStateModel remove(Object o) {
+        var old = modelMap.remove(o);
+        if (o instanceof IModelHoldingBlockState holder) {
+            holder.mfix$setModel(null);
+        }
+        return old;
+    }
+
+    @Override
+    public void putAll(@NotNull Map<? extends BlockState, ? extends BlockStateModel> map) {
+        map.forEach(this::put);
+    }
+
+    @Override
+    public void clear() {
+        modelMap.clear();
+        resetCache();
+    }
+
+    @Override
+    public @NotNull Set<BlockState> keySet() {
+        return modelMap.keySet();
+    }
+
+    @Override
+    public @NotNull Collection<BlockStateModel> values() {
+        return modelMap.values();
+    }
+
+    @Override
+    public @NotNull Set<Entry<BlockState, BlockStateModel>> entrySet() {
+        return modelMap.entrySet();
+    }
+
+    public static void resetCache() {
+        for (var state : Block.BLOCK_STATE_REGISTRY) {
+            ((IModelHoldingBlockState) state).mfix$setModel(null);
+        }
+    }
+}

src/main/java/org/embeddedt/modernfix/dynresources/DynamicModelSystem.java

@@ -11,6 +11,7 @@
 import it.unimi.dsi.fastutil.objects.ObjectSets;
 import it.unimi.dsi.fastutil.objects.ReferenceSets;
 import net.minecraft.client.color.block.BlockColors;
+import net.minecraft.client.renderer.block.model.BlockStateModel;
 import net.minecraft.client.renderer.block.model.ItemModelGenerator;
 import net.minecraft.client.resources.model.BlockStateModelLoader;
 import net.minecraft.client.resources.model.ClientItemInfoLoader;
@@ -35,6 +36,7 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.function.BiFunction;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 
 public class DynamicModelSystem {
@@ -172,13 +174,13 @@ public Object load(K key) throws Exception {
                 }
             }
         });
-        return new DynamicRegistryMap<>(input.keySet(),k -> {
+        return new DynamicRegistryMap<>(input.keySet(), k -> {
             if (k != null) {
                 Object value = bakedCache.getUnchecked(k);
                 if (value == NULL_BAKED) {
                     value = null;
                 }
-                return (V)value;
+                return (V) value;
             } else {
                 return null;
             }