From ef528a4f3a2c3de474e9e2ea8f9c7461514764be Mon Sep 17 00:00:00 2001
From: embeddedt <42941056+embeddedt@users.noreply.github.com>
Date: Sun, 9 Jun 2024 20:40:21 -0400
Subject: [PATCH] Avoid thread-local map by storing model reference on
 BlockState object directly

---
 .../BlockModelShaperMixin.java                | 35 +++++++++++--------
 .../BlockStateBaseMixin.java                  | 26 ++++++++++++++
 .../duck/IModelHoldingBlockState.java         |  8 +++++
 3 files changed, 55 insertions(+), 14 deletions(-)
 create mode 100644 common/src/main/java/org/embeddedt/modernfix/common/mixin/perf/dynamic_resources/BlockStateBaseMixin.java
 create mode 100644 common/src/main/java/org/embeddedt/modernfix/duck/IModelHoldingBlockState.java

diff --git a/common/src/main/java/org/embeddedt/modernfix/common/mixin/perf/dynamic_resources/BlockModelShaperMixin.java b/common/src/main/java/org/embeddedt/modernfix/common/mixin/perf/dynamic_resources/BlockModelShaperMixin.java
index 9ba89eeb..d146bfb2 100644
--- a/common/src/main/java/org/embeddedt/modernfix/common/mixin/perf/dynamic_resources/BlockModelShaperMixin.java
+++ b/common/src/main/java/org/embeddedt/modernfix/common/mixin/perf/dynamic_resources/BlockModelShaperMixin.java
@@ -1,12 +1,14 @@
 package org.embeddedt.modernfix.common.mixin.perf.dynamic_resources;
 
-import it.unimi.dsi.fastutil.objects.Reference2ReferenceLinkedOpenHashMap;
 import net.minecraft.client.renderer.block.BlockModelShaper;
 import net.minecraft.client.resources.model.BakedModel;
 import net.minecraft.client.resources.model.ModelManager;
 import net.minecraft.client.resources.model.ModelResourceLocation;
+import net.minecraft.core.registries.BuiltInRegistries;
+import net.minecraft.world.level.block.Block;
 import net.minecraft.world.level.block.state.BlockState;
 import org.embeddedt.modernfix.annotation.ClientOnlyMixin;
+import org.embeddedt.modernfix.duck.IModelHoldingBlockState;
 import org.embeddedt.modernfix.dynamicresources.ModelLocationCache;
 import org.embeddedt.modernfix.util.DynamicOverridableMap;
 import org.spongepowered.asm.mixin.*;
@@ -24,13 +26,18 @@ public class BlockModelShaperMixin {
     @Shadow
     private Map<BlockState, BakedModel> modelByStateCache;
 
-    private ThreadLocal<Reference2ReferenceLinkedOpenHashMap<BlockState, BakedModel>> mfix$modelCache = ThreadLocal.withInitial(Reference2ReferenceLinkedOpenHashMap::new);
-
     @Inject(method = { "<init>", "replaceCache" }, at = @At("RETURN"))
     private void replaceModelMap(CallbackInfo ci) {
         // replace the backing map for mods which will access it
         this.modelByStateCache = new DynamicOverridableMap<>(state -> modelManager.getModel(ModelLocationCache.get(state)));
-        this.mfix$modelCache = ThreadLocal.withInitial(Reference2ReferenceLinkedOpenHashMap::new);
+        // Clear the cached models on blockstate objects
+        for(Block block : BuiltInRegistries.BLOCK) {
+            for(BlockState state : block.getStateDefinition().getPossibleStates()) {
+                if(state instanceof IModelHoldingBlockState modelHolder) {
+                    modelHolder.mfix$setModel(null);
+                }
+            }
+        }
     }
 
     private BakedModel cacheBlockModel(BlockState state) {
@@ -50,18 +57,18 @@ public class BlockModelShaperMixin {
      */
     @Overwrite
     public BakedModel getBlockModel(BlockState state) {
-        Reference2ReferenceLinkedOpenHashMap<BlockState, BakedModel> map = this.mfix$modelCache.get();
-        BakedModel model = map.get(state);
+        if(state instanceof IModelHoldingBlockState modelHolder) {
+            BakedModel model = modelHolder.mfix$getModel();
 
-        if(model != null) {
+            if(model != null) {
+                return model;
+            }
+
+            model = this.cacheBlockModel(state);
+            modelHolder.mfix$setModel(model);
             return model;
+        } else {
+            return this.cacheBlockModel(state);
         }
-
-        model = this.cacheBlockModel(state);
-        map.putAndMoveToFirst(state, model);
-        if(map.size() > 500) {
-            map.removeLast();
-        }
-        return model;
     }
 }
diff --git a/common/src/main/java/org/embeddedt/modernfix/common/mixin/perf/dynamic_resources/BlockStateBaseMixin.java b/common/src/main/java/org/embeddedt/modernfix/common/mixin/perf/dynamic_resources/BlockStateBaseMixin.java
new file mode 100644
index 00000000..0221e374
--- /dev/null
+++ b/common/src/main/java/org/embeddedt/modernfix/common/mixin/perf/dynamic_resources/BlockStateBaseMixin.java
@@ -0,0 +1,26 @@
+package org.embeddedt.modernfix.common.mixin.perf.dynamic_resources;
+
+import net.minecraft.client.resources.model.BakedModel;
+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 BlockStateBaseMixin implements IModelHoldingBlockState {
+    private volatile SoftReference<BakedModel> mfix$model;
+
+    @Override
+    public BakedModel mfix$getModel() {
+        var ref = mfix$model;
+        return ref != null ? ref.get() : null;
+    }
+
+    @Override
+    public void mfix$setModel(BakedModel model) {
+        mfix$model = model != null ? new SoftReference<>(model) : null;
+    }
+}
diff --git a/common/src/main/java/org/embeddedt/modernfix/duck/IModelHoldingBlockState.java b/common/src/main/java/org/embeddedt/modernfix/duck/IModelHoldingBlockState.java
new file mode 100644
index 00000000..0471e64c
--- /dev/null
+++ b/common/src/main/java/org/embeddedt/modernfix/duck/IModelHoldingBlockState.java
@@ -0,0 +1,8 @@
+package org.embeddedt.modernfix.duck;
+
+import net.minecraft.client.resources.model.BakedModel;
+
+public interface IModelHoldingBlockState {
+    BakedModel mfix$getModel();
+    void mfix$setModel(BakedModel model);
+}