From a3e7435c615922cb00318f70dc1d616274b5e631 Mon Sep 17 00:00:00 2001
From: embeddedt <42941056+embeddedt@users.noreply.github.com>
Date: Tue, 25 Apr 2023 11:10:09 -0400
Subject: [PATCH] Correctly emulate nullishness of baked top level model map

---
 .../DynamicBakedModelProvider.java            | 24 ++++++++++++++++---
 .../dynamic_resources/ModelBakeryMixin.java   | 16 ++++++++++++-
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/src/main/java/org/embeddedt/modernfix/dynamicresources/DynamicBakedModelProvider.java b/src/main/java/org/embeddedt/modernfix/dynamicresources/DynamicBakedModelProvider.java
index 74fca3d5..3b21d552 100644
--- a/src/main/java/org/embeddedt/modernfix/dynamicresources/DynamicBakedModelProvider.java
+++ b/src/main/java/org/embeddedt/modernfix/dynamicresources/DynamicBakedModelProvider.java
@@ -5,6 +5,7 @@ import it.unimi.dsi.fastutil.objects.Object2ObjectOpenHashMap;
 import it.unimi.dsi.fastutil.objects.ObjectOpenHashSet;
 import net.minecraft.client.resources.model.BakedModel;
 import net.minecraft.client.resources.model.BlockModelRotation;
+import net.minecraft.client.resources.model.BuiltInModel;
 import net.minecraft.client.resources.model.ModelBakery;
 import net.minecraft.resources.ResourceLocation;
 import org.apache.commons.lang3.tuple.Triple;
@@ -22,12 +23,19 @@ public class DynamicBakedModelProvider implements Map<ResourceLocation, BakedMod
     private final ModelBakery bakery;
     private final Map<Triple<ResourceLocation, Transformation, Boolean>, BakedModel> bakedCache;
     private final Map<ResourceLocation, BakedModel> permanentOverrides;
+    private BakedModel missingModel;
+    private static final BakedModel SENTINEL = new BuiltInModel(null, null, null, false);
 
     public DynamicBakedModelProvider(ModelBakery bakery, Map<Triple<ResourceLocation, Transformation, Boolean>, BakedModel> cache) {
         this.bakery = bakery;
         this.bakedCache = cache;
         this.permanentOverrides = new Object2ObjectOpenHashMap<>();
     }
+
+    public void setMissingModel(BakedModel model) {
+        this.missingModel = model;
+    }
+
     private static Triple<ResourceLocation, Transformation, Boolean> vanillaKey(Object o) {
         return Triple.of((ResourceLocation)o, BlockModelRotation.X0_Y0.getRotation(), false);
     }
@@ -43,7 +51,7 @@ public class DynamicBakedModelProvider implements Map<ResourceLocation, BakedMod
 
     @Override
     public boolean containsKey(Object o) {
-        return true; //permanentOverrides.containsKey(o) || bakedCache.containsKey(vanillaKey(o));
+        return permanentOverrides.getOrDefault(o, SENTINEL) != null;
     }
 
     @Override
@@ -53,8 +61,18 @@ public class DynamicBakedModelProvider implements Map<ResourceLocation, BakedMod
 
     @Override
     public BakedModel get(Object o) {
-        BakedModel model = permanentOverrides.get(o);
-        return model != null ? model : bakery.bake((ResourceLocation)o, BlockModelRotation.X0_Y0);
+        BakedModel model = permanentOverrides.getOrDefault(o, SENTINEL);
+        if(model != SENTINEL)
+            return model;
+        else {
+            model = bakery.bake((ResourceLocation)o, BlockModelRotation.X0_Y0);
+            if(model == missingModel) {
+                // to correctly emulate the original map, we return null for missing models
+                permanentOverrides.put((ResourceLocation) o, null);
+                return null;
+            } else
+                return model;
+        }
     }
 
     @Nullable
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 240cffce..ea1928e3 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
@@ -103,6 +103,8 @@ public abstract class ModelBakeryMixin implements IExtendedModelBakery {
 
     @Shadow public abstract UnbakedModel getModel(ResourceLocation modelLocation);
 
+    @Shadow @Nullable public abstract BakedModel getBakedModel(ResourceLocation arg, ModelState arg2, Function<Material, TextureAtlasSprite> textureGetter);
+
     private Cache<Triple<ResourceLocation, Transformation, Boolean>, BakedModel> loadedBakedModels;
     private Cache<ResourceLocation, UnbakedModel> loadedModels;
 
@@ -322,6 +324,8 @@ public abstract class ModelBakeryMixin implements IExtendedModelBakery {
     @Inject(method = "uploadTextures", at = @At(value = "FIELD", target = "Lnet/minecraft/client/resources/model/ModelBakery;topLevelModels:Ljava/util/Map;", ordinal = 0), cancellable = true)
     private void skipBake(TextureManager resourceManager, ProfilerFiller profiler, CallbackInfoReturnable<AtlasSet> cir) {
         profiler.pop();
+        // ensure missing model is a permanent override
+        this.bakedTopLevelModels.put(MISSING_MODEL_LOCATION, this.getBakedModel(MISSING_MODEL_LOCATION, BlockModelRotation.X0_Y0, this.atlasSet::getSprite));
         cir.setReturnValue(atlasSet);
     }
 
@@ -446,6 +450,8 @@ public abstract class ModelBakeryMixin implements IExtendedModelBakery {
         return loadOnlyRelevantBlockState(stateDefinition, location);
     }
 
+    private BakedModel bakedMissingModel = null;
+
     @Inject(method = "getBakedModel", at = @At("HEAD"), cancellable = true)
     public void getOrLoadBakedModelDynamic(ResourceLocation arg, ModelState arg2, Function<Material, TextureAtlasSprite> textureGetter, CallbackInfoReturnable<BakedModel> cir) {
         Triple<ResourceLocation, Transformation, Boolean> triple = Triple.of(arg, arg2.getRotation(), arg2.isUvLocked());
@@ -468,7 +474,15 @@ public abstract class ModelBakeryMixin implements IExtendedModelBakery {
                     }
                 }
                 if(ibakedmodel == null) {
-                    ibakedmodel = iunbakedmodel.bake((ModelBakery) (Object) this, textureGetter, arg2, arg);
+                    if(iunbakedmodel == missingModel) {
+                        // use a shared baked missing model
+                        if(bakedMissingModel == null) {
+                            bakedMissingModel = iunbakedmodel.bake((ModelBakery) (Object) this, textureGetter, arg2, arg);
+                            ((DynamicBakedModelProvider)this.bakedTopLevelModels).setMissingModel(bakedMissingModel);
+                        }
+                        ibakedmodel = bakedMissingModel;
+                    } else
+                        ibakedmodel = iunbakedmodel.bake((ModelBakery) (Object) this, textureGetter, arg2, arg);
                 }
                 DynamicModelBakeEvent event = new DynamicModelBakeEvent(arg, iunbakedmodel, ibakedmodel, (ModelLoader)(Object)this);
                 MinecraftForge.EVENT_BUS.post(event);