From 539a798db74c8c9a3dfc041750098a1623aef5e3 Mon Sep 17 00:00:00 2001
From: BlayTheNinth <1933180+BlayTheNinth@users.noreply.github.com>
Date: Mon, 19 Feb 2024 20:34:57 +0100
Subject: [PATCH] fix: Use ModelLoadingPlugin to fix additional model loading
 in Fabric

---
 .../balm/fabric/client/FabricBalmClient.java  |  5 +++
 .../client/rendering/FabricBalmModels.java    | 37 ++++++++-----------
 2 files changed, 21 insertions(+), 21 deletions(-)

diff --git a/fabric/src/main/java/net/blay09/mods/balm/fabric/client/FabricBalmClient.java b/fabric/src/main/java/net/blay09/mods/balm/fabric/client/FabricBalmClient.java
index f9d28fa2..866181fb 100644
--- a/fabric/src/main/java/net/blay09/mods/balm/fabric/client/FabricBalmClient.java
+++ b/fabric/src/main/java/net/blay09/mods/balm/fabric/client/FabricBalmClient.java
@@ -1,10 +1,13 @@
 package net.blay09.mods.balm.fabric.client;
 
 import net.blay09.mods.balm.api.Balm;
+import net.blay09.mods.balm.api.client.BalmClient;
 import net.blay09.mods.balm.api.event.client.DisconnectedFromServerEvent;
+import net.blay09.mods.balm.fabric.client.rendering.FabricBalmModels;
 import net.blay09.mods.balm.fabric.network.FabricBalmNetworking;
 import net.fabricmc.api.ClientModInitializer;
 import net.fabricmc.fabric.api.client.event.lifecycle.v1.ClientLifecycleEvents;
+import net.fabricmc.fabric.api.client.model.loading.v1.ModelLoadingPlugin;
 
 public class FabricBalmClient implements ClientModInitializer {
     @Override
@@ -12,5 +15,7 @@ public void onInitializeClient() {
         ClientLifecycleEvents.CLIENT_STARTED.register(client -> FabricBalmNetworking.initializeClientHandlers());
 
         Balm.getEvents().onEvent(DisconnectedFromServerEvent.class, event -> Balm.getConfig().resetToBackingConfigs());
+
+        ModelLoadingPlugin.register((FabricBalmModels) BalmClient.getModels());
     }
 }
diff --git a/fabric/src/main/java/net/blay09/mods/balm/fabric/client/rendering/FabricBalmModels.java b/fabric/src/main/java/net/blay09/mods/balm/fabric/client/rendering/FabricBalmModels.java
index 0ac7c177..cb0dec68 100644
--- a/fabric/src/main/java/net/blay09/mods/balm/fabric/client/rendering/FabricBalmModels.java
+++ b/fabric/src/main/java/net/blay09/mods/balm/fabric/client/rendering/FabricBalmModels.java
@@ -1,10 +1,10 @@
 package net.blay09.mods.balm.fabric.client.rendering;
 
 import com.mojang.datafixers.util.Pair;
-import com.mojang.logging.LogUtils;
 import com.mojang.math.Transformation;
 import net.blay09.mods.balm.api.DeferredObject;
 import net.blay09.mods.balm.api.client.rendering.BalmModels;
+import net.fabricmc.fabric.api.client.model.loading.v1.ModelLoadingPlugin;
 import net.fabricmc.loader.api.FabricLoader;
 import net.fabricmc.loader.api.MappingResolver;
 import net.minecraft.client.renderer.RenderType;
@@ -16,7 +16,6 @@
 import net.minecraft.world.level.block.state.BlockState;
 import org.jetbrains.annotations.Nullable;
 import org.joml.Matrix4f;
-import org.slf4j.Logger;
 
 import java.lang.reflect.Constructor;
 import java.lang.reflect.InvocationTargetException;
@@ -26,38 +25,37 @@
 import java.util.function.Function;
 import java.util.function.Supplier;
 
-public class FabricBalmModels implements BalmModels {
-
-    private static final Logger LOGGER = LogUtils.getLogger();
+public class FabricBalmModels implements BalmModels, ModelLoadingPlugin {
 
     private static abstract class DeferredModel extends DeferredObject<BakedModel> {
         public DeferredModel(ResourceLocation identifier) {
             super(identifier);
         }
 
-        public void resolveAndSet(ModelBakery modelBakery, BiFunction<ResourceLocation, Material, TextureAtlasSprite> spriteBiFunction) {
-            try {
-                set(resolve(modelBakery, spriteBiFunction));
-            } catch (Exception exception) {
-                LOGGER.warn("Unable to bake model: '{}':", getIdentifier(), exception);
-                set(modelBakery.getBakedTopLevelModels().get(ModelBakery.MISSING_MODEL_LOCATION));
-            }
-        }
-
         public abstract BakedModel resolve(ModelBakery modelBakery, BiFunction<ResourceLocation, Material, TextureAtlasSprite> spriteBiFunction);
-    }
 
+        @Override
+        public void set(BakedModel object) {
+            super.set(object);
+        }
+    }
 
+    private final List<ResourceLocation> additionalModels = Collections.synchronizedList(new ArrayList<>());
     private final List<DeferredModel> modelsToBake = Collections.synchronizedList(new ArrayList<>());
     public final List<Pair<Supplier<Block>, Supplier<BakedModel>>> overrides = Collections.synchronizedList(new ArrayList<>());
     private ModelBakery modelBakery;
 
+    @Override
+    public void onInitializeModelLoader(Context context) {
+        context.addModels(additionalModels);
+    }
+
     public void onBakeModels(ModelBakery modelBakery, BiFunction<ResourceLocation, Material, TextureAtlasSprite> spriteBiFunction) {
         this.modelBakery = modelBakery;
 
         synchronized (modelsToBake) {
             for (DeferredModel model : modelsToBake) {
-                model.resolveAndSet(modelBakery, spriteBiFunction);
+                model.set(model.resolve(modelBakery, spriteBiFunction));
             }
         }
 
@@ -78,13 +76,10 @@ public DeferredObject<BakedModel> loadModel(final ResourceLocation identifier) {
         DeferredModel deferredModel = new DeferredModel(identifier) {
             @Override
             public BakedModel resolve(ModelBakery bakery, BiFunction<ResourceLocation, Material, TextureAtlasSprite> spriteBiFunction) {
-                UnbakedModel model = bakery.getModel(identifier);
-                ModelBaker baker = createBaker(identifier, spriteBiFunction);
-                Function<Material, TextureAtlasSprite> modelTextureGetter = createModelTextureGetter(identifier, spriteBiFunction);
-                return model.bake(baker, modelTextureGetter, getModelState(Transformation.identity()), identifier);
+                return modelBakery.getBakedTopLevelModels().get(identifier);
             }
         };
-        modelsToBake.add(deferredModel);
+        additionalModels.add(identifier);
         return deferredModel;
     }