From 6374d90d4a03a56a0b6077272df151184d375c9c Mon Sep 17 00:00:00 2001
From: embeddedt <42941056+embeddedt@users.noreply.github.com>
Date: Mon, 10 Apr 2023 15:58:12 -0400
Subject: [PATCH] Properly hold onto unbaked models as they are loaded

---
 .../dynamic_resources/ModelBakeryMixin.java   | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/src/main/java/org/embeddedt/modernfix/mixin/perf/dynamic_resources/ModelBakeryMixin.java b/src/main/java/org/embeddedt/modernfix/mixin/perf/dynamic_resources/ModelBakeryMixin.java
index a4784a9b..e81e8733 100644
--- a/src/main/java/org/embeddedt/modernfix/mixin/perf/dynamic_resources/ModelBakeryMixin.java
+++ b/src/main/java/org/embeddedt/modernfix/mixin/perf/dynamic_resources/ModelBakeryMixin.java
@@ -92,6 +92,8 @@ public abstract class ModelBakeryMixin {
     private Cache<Triple<ResourceLocation, Transformation, Boolean>, BakedModel> loadedBakedModels;
     private Cache<ResourceLocation, UnbakedModel> loadedModels;
 
+    private HashMap<ResourceLocation, UnbakedModel> smallLoadingCache = new HashMap<>();
+
 
     @Inject(method = "<init>(Lnet/minecraft/server/packs/resources/ResourceManager;Lnet/minecraft/client/color/block/BlockColors;Z)V", at = @At("RETURN"))
     private void replaceTopLevelBakedModels(ResourceManager manager, BlockColors colors, boolean vanillaBakery, CallbackInfo ci) {
@@ -230,6 +232,11 @@ public abstract class ModelBakeryMixin {
         return unbakedCache.get(rl);
     }
 
+    @Inject(method = "cacheAndQueueDependencies", at = @At("RETURN"))
+    private void addToSmallLoadingCache(ResourceLocation location, UnbakedModel model, CallbackInfo ci) {
+        smallLoadingCache.put(location, model);
+    }
+
 
     /**
      * @author embeddedt
@@ -258,21 +265,27 @@ public abstract class ModelBakeryMixin {
                                 if(debugDynamicModelLoading)
                                     LOGGER.info("Loading {}", resourcelocation);
                                 this.loadModel(resourcelocation);
-                                // TODO: in theory the cache can get evicted right here and we lose the model
-                                // very unlikely to occur though
                             }
                         } catch (ModelBakery.BlockStateDefinitionException var9) {
                             LOGGER.warn(var9.getMessage());
                             this.unbakedCache.put(resourcelocation, iunbakedmodel);
+                            smallLoadingCache.put(resourcelocation, iunbakedmodel);
                         } catch (Exception var10) {
                             LOGGER.warn("Unable to load model: '{}' referenced from: {}: {}", resourcelocation, modelLocation, var10);
                             this.unbakedCache.put(resourcelocation, iunbakedmodel);
+                            smallLoadingCache.put(resourcelocation, iunbakedmodel);
                         } finally {
                             this.loadingStack.remove(resourcelocation);
                         }
                     }
 
-                    return this.unbakedCache.getOrDefault(modelLocation, iunbakedmodel);
+                    // We have to get the result from the temporary cache used for a model load
+                    // As in pathological cases (e.g. Pedestals on 1.19) unbakedCache can lose
+                    // the model immediately
+                    UnbakedModel result = smallLoadingCache.getOrDefault(modelLocation, iunbakedmodel);
+                    // We are done with loading, so clear this cache to allow GC of any unneeded models
+                    smallLoadingCache.clear();
+                    return result;
                 }
             }
         }