Improve retrieval speed for frequently accessed models

src/main/java/org/embeddedt/vintagefix/dynamicresources/model/DynamicModelCache.java

@@ -0,0 +1,79 @@
+package org.embeddedt.vintagefix.dynamicresources.model;
+
+import it.unimi.dsi.fastutil.objects.Reference2ReferenceLinkedOpenHashMap;
+import net.minecraft.client.renderer.block.model.IBakedModel;
+
+import java.util.concurrent.locks.StampedLock;
+import java.util.function.Function;
+
+/**
+ * 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, IBakedModel> cache = new Reference2ReferenceLinkedOpenHashMap<>();
+    private final StampedLock lock = new StampedLock();
+    private final Function<K, IBakedModel> modelRetriever;
+    private final boolean allowNulls;
+
+    public DynamicModelCache(Function<K, IBakedModel> 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 IBakedModel getModelFromCache(K state) {
+        long stamp = lock.readLock();
+        try {
+            return cache.get(state);
+        } finally {
+            lock.unlock(stamp);
+        }
+    }
+
+    private IBakedModel cacheModel(K state) {
+        IBakedModel 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 IBakedModel get(K key) {
+        IBakedModel model = getModelFromCache(key);
+
+        if(model == null && (!allowNulls || needToPopulate(key))) {
+            model = cacheModel(key);
+        }
+
+        return model;
+    }
+}

src/main/java/org/embeddedt/vintagefix/mixin/dynamic_resources/MixinBlockModelShapes.java

@@ -8,6 +8,7 @@
 import net.minecraft.client.renderer.block.statemap.BlockStateMapper;
 import org.embeddedt.vintagefix.annotation.ClientOnlyMixin;
 import org.embeddedt.vintagefix.dynamicresources.IBlockModelShapes;
+import org.embeddedt.vintagefix.dynamicresources.model.DynamicModelCache;
 import org.spongepowered.asm.mixin.Final;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.Overwrite;
@@ -20,9 +21,9 @@
 public abstract class MixinBlockModelShapes implements IBlockModelShapes {
     @Shadow @Final private ModelManager modelManager;
     @Shadow @Final private BlockStateMapper blockStateMapper;
-    @Shadow public abstract ModelManager getModelManager();
 
     private static Map<IBlockState, ModelResourceLocation> modelLocations;
+    private final DynamicModelCache<IBlockState> vintage$modelCache = new DynamicModelCache<>(this::getModelForStateSlow, false);
 
     /**
      * @reason Don't get all models during init (with dynamic loading, that would
@@ -32,6 +33,15 @@ public abstract class MixinBlockModelShapes implements IBlockModelShapes {
     public void reloadModels() {
         if(modelLocations == null)
             modelLocations = blockStateMapper.putAllStateModelLocations();
+        this.vintage$modelCache.clear();
+    }
+
+    private IBakedModel getModelForStateSlow(IBlockState state) {
+        IBakedModel model = modelManager.getModel(getLocationForState(state));
+        if (model == null) {
+            model = modelManager.getMissingModel();
+        }
+        return model;
     }
 
     /**
@@ -40,11 +50,7 @@ public void reloadModels() {
      **/
     @Overwrite
     public IBakedModel getModelForState(IBlockState state) {
-        IBakedModel model = modelManager.getModel(getLocationForState(state));
-        if (model == null) {
-            model = modelManager.getMissingModel();
-        }
-        return model;
+        return this.vintage$modelCache.get(state);
     }
 
     @Override

src/main/java/org/embeddedt/vintagefix/mixin/dynamic_resources/MixinItemModelMesherForge.java

@@ -2,6 +2,7 @@
 
 import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
 import it.unimi.dsi.fastutil.ints.Int2ObjectOpenHashMap;
+import it.unimi.dsi.fastutil.objects.Reference2ReferenceOpenHashMap;
 import net.minecraft.client.renderer.ItemModelMesher;
 import net.minecraft.client.renderer.block.model.IBakedModel;
 import net.minecraft.client.renderer.block.model.ModelManager;
@@ -10,21 +11,30 @@
 import net.minecraftforge.client.ItemModelMesherForge;
 import net.minecraftforge.registries.IRegistryDelegate;
 import org.embeddedt.vintagefix.annotation.ClientOnlyMixin;
-import org.spongepowered.asm.mixin.Mixin;
-import org.spongepowered.asm.mixin.Overwrite;
-import org.spongepowered.asm.mixin.Shadow;
+import org.embeddedt.vintagefix.dynamicresources.model.DynamicModelCache;
+import org.spongepowered.asm.mixin.*;
 
 import java.util.Map;
 
 @Mixin(ItemModelMesherForge.class)
 @ClientOnlyMixin
 public class MixinItemModelMesherForge extends ItemModelMesher {
-    @Shadow Map<IRegistryDelegate<Item>, Int2ObjectMap<ModelResourceLocation>> locations;
+    @Shadow @Final @Mutable
+    Map<IRegistryDelegate<Item>, Int2ObjectMap<ModelResourceLocation>> locations = new Reference2ReferenceOpenHashMap<>();
+
+    // This is a pretty clever trick to speed up the model lookups - we know that our location objects per-item are unique,
+    // so we can just do reference lookup on them
+    private final DynamicModelCache<ModelResourceLocation> vintage$itemModelCache = new DynamicModelCache<>(this::getItemModelByLocationSlow, true);
 
     public MixinItemModelMesherForge(ModelManager modelManager) {
         super(modelManager);
     }
 
+
+    private IBakedModel getItemModelByLocationSlow(ModelResourceLocation mrl) {
+        return getModelManager().getModel(mrl);
+    }
+
     /**
      * @reason Get the stored location for that item and meta, and get the model
      * from that location from the model manager.
@@ -38,7 +48,7 @@ protected IBakedModel getItemModel(Item item, int meta) {
         ModelResourceLocation location = map.get(meta);
         if(location == null)
             return null;
-        return getModelManager().getModel(location);
+        return this.vintage$itemModelCache.get(location);
     }
 
     /**
@@ -63,5 +73,7 @@ public void register(Item item, int meta, ModelResourceLocation location) {
      **/
     @Overwrite
     @Override
-    public void rebuildCache() {}
+    public void rebuildCache() {
+        this.vintage$itemModelCache.clear();
+    }
 }