From 0a2601257c1a91f684f7e326c4530e73d6871e31 Mon Sep 17 00:00:00 2001
From: embeddedt <42941056+embeddedt@users.noreply.github.com>
Date: Wed, 5 Jul 2023 12:04:10 -0400
Subject: [PATCH] Fix dynamic resources losing nested model loads sometimes

---
 .../mixin/perf/dynamic_resources/ModelBakeryMixin.java   | 9 +++++++--
 .../mixin/perf/dynamic_resources/ModelBakeryMixin.java   | 8 ++++++--
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/fabric/src/main/java/org/embeddedt/modernfix/fabric/mixin/perf/dynamic_resources/ModelBakeryMixin.java b/fabric/src/main/java/org/embeddedt/modernfix/fabric/mixin/perf/dynamic_resources/ModelBakeryMixin.java
index 58474c27..b7a9d8dc 100644
--- a/fabric/src/main/java/org/embeddedt/modernfix/fabric/mixin/perf/dynamic_resources/ModelBakeryMixin.java
+++ b/fabric/src/main/java/org/embeddedt/modernfix/fabric/mixin/perf/dynamic_resources/ModelBakeryMixin.java
@@ -348,6 +348,8 @@ public abstract class ModelBakeryMixin implements IExtendedModelBakery {
         return model;
     }
 
+    private int mfix$nestedLoads = 0;
+
     /**
      * @author embeddedt
      * @reason synchronize
@@ -372,6 +374,7 @@ public abstract class ModelBakeryMixin implements IExtendedModelBakery {
                     while(!this.loadingStack.isEmpty()) {
                         ResourceLocation resourcelocation = this.loadingStack.iterator().next();
 
+                        mfix$nestedLoads++;
                         try {
                             existing = this.unbakedCache.get(resourcelocation);
                             if (existing == null) {
@@ -389,6 +392,7 @@ public abstract class ModelBakeryMixin implements IExtendedModelBakery {
                             this.unbakedCache.put(resourcelocation, iunbakedmodel);
                             smallLoadingCache.put(resourcelocation, iunbakedmodel);
                         } finally {
+                            mfix$nestedLoads--;
                             this.loadingStack.remove(resourcelocation);
                         }
                     }
@@ -398,7 +402,8 @@ public abstract class ModelBakeryMixin implements IExtendedModelBakery {
                     // 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();
+                    if(mfix$nestedLoads == 0)
+                        smallLoadingCache.clear();
                     cir.setReturnValue(result);
                 }
             }
@@ -464,7 +469,7 @@ public abstract class ModelBakeryMixin implements IExtendedModelBakery {
                 if (iunbakedmodel instanceof BlockModel) {
                     BlockModel blockmodel = (BlockModel)iunbakedmodel;
                     if (blockmodel.getRootModel() == GENERATION_MARKER) {
-                        ibakedmodel = ITEM_MODEL_GENERATOR.generateBlockModel(textureGetter, blockmodel).bake((ModelBakery)(Object)this, blockmodel, this.atlasSet::getSprite, arg2, arg, false);
+                        ibakedmodel = ITEM_MODEL_GENERATOR.generateBlockModel(textureGetter, blockmodel).bake((ModelBakery)(Object)this, blockmodel, textureGetter, arg2, arg, false);
                     }
                 }
                 if(ibakedmodel == null) {
diff --git a/forge/src/main/java/org/embeddedt/modernfix/forge/mixin/perf/dynamic_resources/ModelBakeryMixin.java b/forge/src/main/java/org/embeddedt/modernfix/forge/mixin/perf/dynamic_resources/ModelBakeryMixin.java
index 3fe1d7cf..6fb6fa29 100644
--- a/forge/src/main/java/org/embeddedt/modernfix/forge/mixin/perf/dynamic_resources/ModelBakeryMixin.java
+++ b/forge/src/main/java/org/embeddedt/modernfix/forge/mixin/perf/dynamic_resources/ModelBakeryMixin.java
@@ -233,6 +233,7 @@ public abstract class ModelBakeryMixin implements IExtendedModelBakery {
         smallLoadingCache.put(location, model);
     }
 
+    private int mfix$nestedLoads = 0;
 
     /**
      * @author embeddedt
@@ -258,6 +259,7 @@ public abstract class ModelBakeryMixin implements IExtendedModelBakery {
                     while(!this.loadingStack.isEmpty()) {
                         ResourceLocation resourcelocation = this.loadingStack.iterator().next();
 
+                        mfix$nestedLoads++;
                         try {
                             existing = this.unbakedCache.get(resourcelocation);
                             if (existing == null) {
@@ -275,6 +277,7 @@ public abstract class ModelBakeryMixin implements IExtendedModelBakery {
                             this.unbakedCache.put(resourcelocation, iunbakedmodel);
                             smallLoadingCache.put(resourcelocation, iunbakedmodel);
                         } finally {
+                            mfix$nestedLoads--;
                             this.loadingStack.remove(resourcelocation);
                         }
                     }
@@ -284,7 +287,8 @@ public abstract class ModelBakeryMixin implements IExtendedModelBakery {
                     // 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();
+                    if(mfix$nestedLoads == 0)
+                        smallLoadingCache.clear();
                     cir.setReturnValue(result);
                 }
             }
@@ -337,7 +341,7 @@ public abstract class ModelBakeryMixin implements IExtendedModelBakery {
                 if (iunbakedmodel instanceof BlockModel) {
                     BlockModel blockmodel = (BlockModel)iunbakedmodel;
                     if (blockmodel.getRootModel() == GENERATION_MARKER) {
-                        ibakedmodel = ITEM_MODEL_GENERATOR.generateBlockModel(textureGetter, blockmodel).bake((ModelBakery)(Object)this, blockmodel, this.atlasSet::getSprite, arg2, arg, false);
+                        ibakedmodel = ITEM_MODEL_GENERATOR.generateBlockModel(textureGetter, blockmodel).bake((ModelBakery)(Object)this, blockmodel, textureGetter, arg2, arg, false);
                     }
                 }
                 if(ibakedmodel == null) {