Fix tint overrides

3 files changed, 63 insertions, 5 deletions

diff --git a/src/main/java/moe/nea/firmament/mixins/custommodels/JsonUnbakedModelDataHolder.java b/src/main/java/moe/nea/firmament/mixins/custommodels/JsonUnbakedModelDataHolder.java
index 12eea5b..2de2e5e 100644
--- a/src/main/java/moe/nea/firmament/mixins/custommodels/JsonUnbakedModelDataHolder.java
+++ b/src/main/java/moe/nea/firmament/mixins/custommodels/JsonUnbakedModelDataHolder.java
@@ -5,6 +5,7 @@ import com.llamalad7.mixinextras.sugar.Local;
 import moe.nea.firmament.features.texturepack.BakedModelExtra;
 import moe.nea.firmament.features.texturepack.JsonUnbakedModelFirmExtra;
 import moe.nea.firmament.features.texturepack.TintOverrides;
+import moe.nea.firmament.util.ErrorUtil;
 import net.minecraft.client.render.model.BakedModel;
 import net.minecraft.client.render.model.Baker;
 import net.minecraft.client.render.model.ModelRotation;
@@ -79,14 +80,11 @@ public abstract class JsonUnbakedModelDataHolder implements JsonUnbakedModelFirm
 		}
 	}
 
-	@ModifyReturnValue(
-		method = "bake(Lnet/minecraft/client/render/model/Baker;Ljava/util/function/Function;Lnet/minecraft/client/render/model/ModelBakeSettings;)Lnet/minecraft/client/render/model/BakedModel;",
-		at = @At(value = "RETURN"))
-	private BakedModel bakeExtraInfo(BakedModel original, @Local(argsOnly = true) Baker baker) {
+	private void addExtraBakeInfo(BakedModel bakedModel, Baker baker) {
 		if (!this.toString().contains("minecraft") && this.toString().contains("crimson")) {
 			System.out.println("Found non minecraft model " + this);
 		}
-		if (original instanceof BakedModelExtra extra) {
+		if (bakedModel instanceof BakedModelExtra extra) {
 			var headModel = getHeadModel_firmament();
 			if (headModel != null) {
 				extra.setHeadModel_firmament(baker.bake(headModel, ModelRotation.X0_Y0));
@@ -95,6 +93,37 @@ public abstract class JsonUnbakedModelDataHolder implements JsonUnbakedModelFirm
 				extra.setTintOverrides_firmament(getTintOverrides_firmament());
 			}
 		}
+	}
+
+	/**
+	 * @see ProvideBakerToJsonUnbakedModelPatch
+	 */
+	@Override
+	public void storeExtraBaker_firmament(@NotNull Baker baker) {
+		this.storedBaker = baker;
+	}
+
+	@Unique
+	private Baker storedBaker;
+
+	@ModifyReturnValue(
+		method = "bake(Ljava/util/function/Function;Lnet/minecraft/client/render/model/ModelBakeSettings;Z)Lnet/minecraft/client/render/model/BakedModel;",
+		at = @At("RETURN"))
+	private BakedModel bakeExtraInfoWithoutBaker(BakedModel original) {
+		if (storedBaker != null) {
+			addExtraBakeInfo(original, storedBaker);
+			storedBaker = null;
+		}
+		return original;
+	}
+
+	@ModifyReturnValue(
+		method = {
+			"bake(Lnet/minecraft/client/render/model/Baker;Ljava/util/function/Function;Lnet/minecraft/client/render/model/ModelBakeSettings;)Lnet/minecraft/client/render/model/BakedModel;"
+		},
+		at = @At(value = "RETURN"))
+	private BakedModel bakeExtraInfo(BakedModel original, @Local(argsOnly = true) Baker baker) {
+		addExtraBakeInfo(original, baker);
 		return original;
 	}
 }
diff --git a/src/main/java/moe/nea/firmament/mixins/custommodels/ProvideBakerToJsonUnbakedModelPatch.java b/src/main/java/moe/nea/firmament/mixins/custommodels/ProvideBakerToJsonUnbakedModelPatch.java
new file mode 100644
index 0000000..c1ac119
--- /dev/null
+++ b/src/main/java/moe/nea/firmament/mixins/custommodels/ProvideBakerToJsonUnbakedModelPatch.java
@@ -0,0 +1,27 @@
+package moe.nea.firmament.mixins.custommodels;
+
+import com.llamalad7.mixinextras.injector.wrapoperation.Operation;
+import com.llamalad7.mixinextras.injector.wrapoperation.WrapOperation;
+import moe.nea.firmament.features.texturepack.JsonUnbakedModelFirmExtra;
+import net.minecraft.client.render.model.BakedModel;
+import net.minecraft.client.render.model.Baker;
+import net.minecraft.client.render.model.ModelBakeSettings;
+import net.minecraft.client.render.model.json.JsonUnbakedModel;
+import net.minecraft.client.texture.Sprite;
+import net.minecraft.client.util.SpriteIdentifier;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.injection.At;
+
+import java.util.function.Function;
+
+/**
+ * @see JsonUnbakedModelDataHolder#storeExtraBaker_firmament
+ */
+@Mixin(targets = "net.minecraft.client.render.model.ModelBaker$BakerImpl")
+public abstract class ProvideBakerToJsonUnbakedModelPatch implements Baker {
+	@WrapOperation(method = "bake(Lnet/minecraft/client/render/model/UnbakedModel;Lnet/minecraft/client/render/model/ModelBakeSettings;)Lnet/minecraft/client/render/model/BakedModel;", at = @At(value = "INVOKE", target = "Lnet/minecraft/client/render/model/json/JsonUnbakedModel;bake(Ljava/util/function/Function;Lnet/minecraft/client/render/model/ModelBakeSettings;Z)Lnet/minecraft/client/render/model/BakedModel;"))
+	private BakedModel provideExtraBakerToModel(JsonUnbakedModel instance, Function<SpriteIdentifier, Sprite> function, ModelBakeSettings modelBakeSettings, boolean bl, Operation<BakedModel> original) {
+		((JsonUnbakedModelFirmExtra) instance).storeExtraBaker_firmament(this);
+		return original.call(instance, function, modelBakeSettings, bl);
+	}
+}
diff --git a/src/main/kotlin/features/texturepack/JsonUnbakedModelFirmExtra.kt b/src/main/kotlin/features/texturepack/JsonUnbakedModelFirmExtra.kt
index 0d0f8f2..9f641b8 100644
--- a/src/main/kotlin/features/texturepack/JsonUnbakedModelFirmExtra.kt
+++ b/src/main/kotlin/features/texturepack/JsonUnbakedModelFirmExtra.kt
@@ -1,9 +1,11 @@
 
 package moe.nea.firmament.features.texturepack
 
+import net.minecraft.client.render.model.Baker
 import net.minecraft.util.Identifier
 
 interface JsonUnbakedModelFirmExtra {
+	fun storeExtraBaker_firmament(baker: Baker)
 
     fun setHeadModel_firmament(identifier: Identifier?)
     fun getHeadModel_firmament(): Identifier?