Implement caching on the blockstate and item to model maps

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

@@ -6,6 +6,7 @@
 import net.minecraft.client.resources.model.ModelResourceLocation;
 import net.minecraft.world.level.block.state.BlockState;
 import org.embeddedt.modernfix.annotation.ClientOnlyMixin;
+import org.embeddedt.modernfix.dynamicresources.DynamicModelCache;
 import org.embeddedt.modernfix.dynamicresources.ModelLocationCache;
 import org.embeddedt.modernfix.util.DynamicOverridableMap;
 import org.spongepowered.asm.mixin.*;
@@ -24,6 +25,8 @@ public class BlockModelShaperMixin {
     @Shadow @Final @Mutable
     private Map<BlockState, BakedModel> modelByStateCache;
 
+    private final DynamicModelCache<BlockState> mfix$modelCache = new DynamicModelCache<>(k -> this.cacheBlockModel((BlockState)k), false);
+
     @Inject(method = "<init>", at = @At("RETURN"))
     private void replaceModelMap(CallbackInfo ci) {
         // replace the backing map for mods which will access it
@@ -32,23 +35,30 @@ private void replaceModelMap(CallbackInfo ci) {
 
     /**
      * @author embeddedt
-     * @reason no need to rebuild model cache, and location cache is done elsewhere
+     * @reason no need to rebuild vanilla model cache
      */
     @Overwrite
     public void rebuildCache() {
+        this.mfix$modelCache.clear();
     }
 
-    /**
-     * @author embeddedt
-     * @reason get the model from the dynamic model provider
-     */
-    @Overwrite
-    public BakedModel getBlockModel(BlockState state) {
+    private BakedModel cacheBlockModel(BlockState state) {
+        // Do all model system accesses in the unlocked path
         ModelResourceLocation mrl = ModelLocationCache.get(state);
         BakedModel model = mrl == null ? null : modelManager.getModel(mrl);
         if (model == null) {
             model = modelManager.getMissingModel();
         }
+
         return model;
     }
+
+    /**
+     * @author embeddedt
+     * @reason get the model from the dynamic model provider
+     */
+    @Overwrite
+    public BakedModel getBlockModel(BlockState state) {
+        return this.mfix$modelCache.get(state);
+    }
 }

common/src/main/java/org/embeddedt/modernfix/dynamicresources/DynamicModelCache.java

@@ -0,0 +1,79 @@
+package org.embeddedt.modernfix.dynamicresources;
+
+import it.unimi.dsi.fastutil.Function;
+import it.unimi.dsi.fastutil.objects.Reference2ReferenceLinkedOpenHashMap;
+import net.minecraft.client.resources.model.BakedModel;
+
+import java.util.concurrent.locks.StampedLock;
+
+/**
+ * The Mojang Triple-based baked cache system is too slow to be hitting on every model retrieval, so
+ * we need a fast, concurrency-safe wrapper on top.
+ */
+public class DynamicModelCache<K> {
+    private final Reference2ReferenceLinkedOpenHashMap<K, BakedModel> cache = new Reference2ReferenceLinkedOpenHashMap<>();
+    private final StampedLock lock = new StampedLock();
+    private final Function<K, BakedModel> modelRetriever;
+    private final boolean allowNulls;
+
+    public DynamicModelCache(Function<K, BakedModel> modelRetriever, boolean allowNulls) {
+        this.modelRetriever = modelRetriever;
+        this.allowNulls = allowNulls;
+    }
+
+    public void clear() {
+        long stamp = lock.writeLock();
+        try {
+            cache.clear();
+        } finally {
+            lock.unlock(stamp);
+        }
+    }
+
+    private boolean needToPopulate(K state) {
+        long stamp = lock.readLock();
+        try {
+            return !cache.containsKey(state);
+        } finally {
+            lock.unlock(stamp);
+        }
+    }
+
+    private BakedModel getModelFromCache(K state) {
+        long stamp = lock.readLock();
+        try {
+            return cache.get(state);
+        } finally {
+            lock.unlock(stamp);
+        }
+    }
+
+    private BakedModel cacheModel(K state) {
+        BakedModel model = modelRetriever.apply(state);
+
+        // Lock and modify our local, faster cache
+        long stamp = lock.writeLock();
+
+        try {
+            cache.putAndMoveToFirst(state, model);
+            // TODO: choose less arbitrary number
+            if(cache.size() >= 1000) {
+                cache.removeLast();
+            }
+        } finally {
+            lock.unlock(stamp);
+        }
+
+        return model;
+    }
+
+    public BakedModel get(K key) {
+        BakedModel model = getModelFromCache(key);
+
+        if(model == null && (!allowNulls || needToPopulate(key))) {
+            model = cacheModel(key);
+        }
+
+        return model;
+    }
+}

fabric/src/test/java/org/embeddedt/modernfix/dynamicresources/DynamicModelCacheTest.java

@@ -0,0 +1,26 @@
+package org.embeddedt.modernfix.dynamicresources;
+
+import net.minecraft.client.resources.model.BakedModel;
+import net.minecraft.client.resources.model.BuiltInModel;
+import net.minecraft.world.item.Item;
+import net.minecraft.world.item.Items;
+import org.embeddedt.modernfix.testing.util.BootstrapMinecraft;
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+@BootstrapMinecraft
+public class DynamicModelCacheTest {
+    @Test
+    public void testCacheReturnsNullForNullGetter() {
+        DynamicModelCache<Item> cache = new DynamicModelCache(k -> null, true);
+        assertNull(cache.get(Items.STONE));
+    }
+
+    @Test
+    public void testCacheFunctions() {
+        BakedModel model = new BuiltInModel(null, null, null, false);
+        DynamicModelCache<Item> cache = new DynamicModelCache(k -> model, true);
+        assertEquals(model, cache.get(Items.STONE));
+    }
+}

forge/src/main/java/org/embeddedt/modernfix/forge/mixin/perf/dynamic_resources/ItemModelMesherForgeMixin.java

@@ -7,6 +7,7 @@
 import net.minecraft.world.item.Item;
 import net.minecraftforge.registries.IRegistryDelegate;
 import org.embeddedt.modernfix.annotation.ClientOnlyMixin;
+import org.embeddedt.modernfix.dynamicresources.DynamicModelCache;
 import org.embeddedt.modernfix.dynamicresources.ModelLocationCache;
 import org.embeddedt.modernfix.util.ItemMesherMap;
 import org.spongepowered.asm.mixin.*;
@@ -24,6 +25,8 @@ public abstract class ItemModelMesherForgeMixin extends ItemModelShaper {
 
     private Map<IRegistryDelegate<Item>, ModelResourceLocation> overrideLocations;
 
+    private final DynamicModelCache<IRegistryDelegate<Item>> mfix$modelCache = new DynamicModelCache<>(k -> this.mfix$getModelSlow((IRegistryDelegate<Item>)k), true);
+
     public ItemModelMesherForgeMixin(ModelManager arg) {
         super(arg);
     }
@@ -47,6 +50,11 @@ private void replaceLocationMap(CallbackInfo ci) {
         return map;
     }
 
+    private BakedModel mfix$getModelSlow(IRegistryDelegate<Item> key) {
+        ModelResourceLocation map = mfix$getLocationForge(key);
+        return map == null ? null : getModelManager().getModel(map);
+    }
+
     /**
      * @author embeddedt
      * @reason Get the stored location for that item and meta, and get the model
@@ -55,8 +63,7 @@ private void replaceLocationMap(CallbackInfo ci) {
     @Overwrite
     @Override
     public BakedModel getItemModel(Item item) {
-        ModelResourceLocation map = mfix$getLocationForge(item.delegate);
-        return map == null ? null : getModelManager().getModel(map);
+        return this.mfix$modelCache.get(item.delegate);
     }
 
     /**
@@ -77,5 +84,7 @@ public void register(Item item, ModelResourceLocation location) {
      **/
     @Overwrite
     @Override
-    public void rebuildCache() {}
+    public void rebuildCache() {
+        this.mfix$modelCache.clear();
+    }
 }