From 87d1aad3d16e8029115fc75d9dc0bca93702389a Mon Sep 17 00:00:00 2001
From: embeddedt <42941056+embeddedt@users.noreply.github.com>
Date: Mon, 5 Aug 2024 09:48:10 -0400
Subject: [PATCH] Backport unbaked model caching from 1.21

---
 .../dynamic_resources/ModelManagerMixin.java   |  7 +++++--
 .../embeddedt/modernfix/util/CacheUtil.java    | 18 ++++++++++++++++++
 2 files changed, 23 insertions(+), 2 deletions(-)
 create mode 100644 common/src/main/java/org/embeddedt/modernfix/util/CacheUtil.java

diff --git a/common/src/main/java/org/embeddedt/modernfix/common/mixin/perf/dynamic_resources/ModelManagerMixin.java b/common/src/main/java/org/embeddedt/modernfix/common/mixin/perf/dynamic_resources/ModelManagerMixin.java
index 392ba1c7..2a1ce44b 100644
--- a/common/src/main/java/org/embeddedt/modernfix/common/mixin/perf/dynamic_resources/ModelManagerMixin.java
+++ b/common/src/main/java/org/embeddedt/modernfix/common/mixin/perf/dynamic_resources/ModelManagerMixin.java
@@ -13,6 +13,7 @@ import net.minecraft.world.level.block.state.BlockState;
 import net.minecraft.world.level.block.state.StateDefinition;
 import org.embeddedt.modernfix.ModernFix;
 import org.embeddedt.modernfix.annotation.ClientOnlyMixin;
+import org.embeddedt.modernfix.util.CacheUtil;
 import org.embeddedt.modernfix.util.LambdaMap;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.Shadow;
@@ -45,12 +46,14 @@ public class ModelManagerMixin {
 
     @Redirect(method = "reload", at = @At(value = "INVOKE", target = "Lnet/minecraft/client/resources/model/ModelManager;loadBlockModels(Lnet/minecraft/server/packs/resources/ResourceManager;Ljava/util/concurrent/Executor;)Ljava/util/concurrent/CompletableFuture;"))
     private CompletableFuture<Map<ResourceLocation, BlockModel>> deferBlockModelLoad(ResourceManager manager, Executor executor) {
-        return CompletableFuture.completedFuture(new LambdaMap<>(location -> loadSingleBlockModel(manager, location)));
+        var cache = CacheUtil.<ResourceLocation, BlockModel>simpleCacheForLambda(location -> loadSingleBlockModel(manager, location), 100L);
+        return CompletableFuture.completedFuture(new LambdaMap<>(location -> cache.getUnchecked(location)));
     }
 
     @Redirect(method = "reload", at = @At(value = "INVOKE", target = "Lnet/minecraft/client/resources/model/ModelManager;loadBlockStates(Lnet/minecraft/server/packs/resources/ResourceManager;Ljava/util/concurrent/Executor;)Ljava/util/concurrent/CompletableFuture;"))
     private CompletableFuture<Map<ResourceLocation, List<ModelBakery.LoadedJson>>> deferBlockStateLoad(ResourceManager manager, Executor executor) {
-        return CompletableFuture.completedFuture(new LambdaMap<>(location -> loadSingleBlockState(manager, location)));
+        var cache = CacheUtil.<ResourceLocation, List<ModelBakery.LoadedJson>>simpleCacheForLambda(location -> loadSingleBlockState(manager, location), 100L);
+        return CompletableFuture.completedFuture(new LambdaMap<>(location -> cache.getUnchecked(location)));
     }
 
     @Redirect(method = "loadModels", at = @At(value = "INVOKE", target = "Lnet/minecraft/world/level/block/state/StateDefinition;getPossibleStates()Lcom/google/common/collect/ImmutableList;"))
diff --git a/common/src/main/java/org/embeddedt/modernfix/util/CacheUtil.java b/common/src/main/java/org/embeddedt/modernfix/util/CacheUtil.java
new file mode 100644
index 00000000..bdeebbe4
--- /dev/null
+++ b/common/src/main/java/org/embeddedt/modernfix/util/CacheUtil.java
@@ -0,0 +1,18 @@
+package org.embeddedt.modernfix.util;
+
+import com.google.common.cache.CacheBuilder;
+import com.google.common.cache.CacheLoader;
+import com.google.common.cache.LoadingCache;
+
+import java.util.function.Function;
+
+public class CacheUtil {
+    public static <K, V> LoadingCache<K, V> simpleCacheForLambda(Function<K, V> function, long maxSize) {
+        return CacheBuilder.newBuilder().maximumSize(maxSize).build(new CacheLoader<K, V>() {
+            @Override
+            public V load(K key) throws Exception {
+                return function.apply(key);
+            }
+        });
+    }
+}
\ No newline at end of file