Fix model unwrapping

9 files changed, 67 insertions, 21 deletions

diff --git a/src/main/java/moe/nea/firmament/mixins/custommodels/ApplyHeadModelInItemRenderer.java b/src/main/java/moe/nea/firmament/mixins/custommodels/ApplyHeadModelInItemRenderer.java
index ffa59c6..dac65fe 100644
--- a/src/main/java/moe/nea/firmament/mixins/custommodels/ApplyHeadModelInItemRenderer.java
+++ b/src/main/java/moe/nea/firmament/mixins/custommodels/ApplyHeadModelInItemRenderer.java
@@ -24,8 +24,8 @@ public class ApplyHeadModelInItemRenderer {
 	                            BakedModel model, boolean useInventoryModel, float z, CallbackInfo ci,
 	                            @Local(argsOnly = true) LocalRef<BakedModel> modelMut
 	) {
-		if (transformationMode == ModelTransformationMode.HEAD
-			&& model instanceof BakedModelExtra extra) {
+		var extra = BakedModelExtra.cast(model);
+		if (transformationMode == ModelTransformationMode.HEAD && extra != null) {
 			var headModel = extra.getHeadModel_firmament();
 			if (headModel != null) {
 				modelMut.set(headModel);
diff --git a/src/main/java/moe/nea/firmament/mixins/custommodels/HeadModelReplacerPatch.java b/src/main/java/moe/nea/firmament/mixins/custommodels/HeadModelReplacerPatch.java
index 7518041..26c331e 100644
--- a/src/main/java/moe/nea/firmament/mixins/custommodels/HeadModelReplacerPatch.java
+++ b/src/main/java/moe/nea/firmament/mixins/custommodels/HeadModelReplacerPatch.java
@@ -35,7 +35,8 @@ public class HeadModelReplacerPatch<S extends LivingEntityRenderState, M extends
 	private Block replaceSkull(BlockItem instance, Operation<Block> original, @Local BakedModel bakedModel) {
 		var oldBlock = original.call(instance);
 		if (oldBlock instanceof AbstractSkullBlock) {
-			if (bakedModel instanceof BakedModelExtra extra && extra.getHeadModel_firmament() != null)
+			var extra = BakedModelExtra.cast(bakedModel);
+			if (extra != null && extra.getHeadModel_firmament() != null)
 				return Blocks.ENCHANTING_TABLE; // Any non skull block. Let's choose the enchanting table because it is very distinct.
 		}
 		return oldBlock;
@@ -48,7 +49,8 @@ public class HeadModelReplacerPatch<S extends LivingEntityRenderState, M extends
 		at = @At(value = "INVOKE", target = "Lnet/minecraft/client/render/entity/feature/ArmorFeatureRenderer;hasModel(Lnet/minecraft/item/ItemStack;Lnet/minecraft/entity/EquipmentSlot;)Z"))
 	private boolean replaceHasModel(ItemStack stack, EquipmentSlot slot, Operation<Boolean> original,
 	                                @Local BakedModel bakedModel) {
-		if (bakedModel instanceof BakedModelExtra extra && extra.getHeadModel_firmament() != null)
+		var extra = BakedModelExtra.cast(bakedModel);
+		if (extra != null && extra.getHeadModel_firmament() != null)
 			return false;
 		return original.call(stack, slot);
 	}
diff --git a/src/main/java/moe/nea/firmament/mixins/custommodels/ItemRendererTintContextPatch.java b/src/main/java/moe/nea/firmament/mixins/custommodels/ItemRendererTintContextPatch.java
index 66feced..8c5411b 100644
--- a/src/main/java/moe/nea/firmament/mixins/custommodels/ItemRendererTintContextPatch.java
+++ b/src/main/java/moe/nea/firmament/mixins/custommodels/ItemRendererTintContextPatch.java
@@ -18,7 +18,8 @@ public class ItemRendererTintContextPatch {
 	@Inject(method = "renderItem(Lnet/minecraft/item/ItemStack;Lnet/minecraft/item/ModelTransformationMode;ZLnet/minecraft/client/util/math/MatrixStack;Lnet/minecraft/client/render/VertexConsumerProvider;IILnet/minecraft/client/render/model/BakedModel;ZF)V",
 		at = @At(value = "HEAD"), allow = 1)
 	private void onStartRendering(ItemStack stack, ModelTransformationMode transformationMode, boolean leftHanded, MatrixStack matrices, VertexConsumerProvider vertexConsumers, int light, int overlay, BakedModel model, boolean useInventoryModel, float z, CallbackInfo ci) {
-		if (model instanceof BakedModelExtra extra) {
+		var extra = BakedModelExtra.cast(model);
+		if (extra != null) {
 			TintOverrides.Companion.enter(extra.getTintOverrides_firmament());
 		}
 	}
@@ -26,7 +27,8 @@ public class ItemRendererTintContextPatch {
 	@Inject(method = "renderItem(Lnet/minecraft/item/ItemStack;Lnet/minecraft/item/ModelTransformationMode;ZLnet/minecraft/client/util/math/MatrixStack;Lnet/minecraft/client/render/VertexConsumerProvider;IILnet/minecraft/client/render/model/BakedModel;ZF)V",
 		at = @At("TAIL"), allow = 1)
 	private void onEndRendering(ItemStack stack, ModelTransformationMode transformationMode, boolean leftHanded, MatrixStack matrices, VertexConsumerProvider vertexConsumers, int light, int overlay, BakedModel model, boolean useInventoryModel, float z, CallbackInfo ci) {
-		if (model instanceof BakedModelExtra extra) {
+		var extra = BakedModelExtra.cast(model);
+		if (extra != null) {
 			TintOverrides.Companion.exit(extra.getTintOverrides_firmament());
 		}
 	}
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 2de2e5e..a5bb34f 100644
--- a/src/main/java/moe/nea/firmament/mixins/custommodels/JsonUnbakedModelDataHolder.java
+++ b/src/main/java/moe/nea/firmament/mixins/custommodels/JsonUnbakedModelDataHolder.java
@@ -84,7 +84,8 @@ public abstract class JsonUnbakedModelDataHolder implements JsonUnbakedModelFirm
 		if (!this.toString().contains("minecraft") && this.toString().contains("crimson")) {
 			System.out.println("Found non minecraft model " + this);
 		}
-		if (bakedModel instanceof BakedModelExtra extra) {
+		var extra = BakedModelExtra.cast(bakedModel);
+		if (extra != null) {
 			var headModel = getHeadModel_firmament();
 			if (headModel != null) {
 				extra.setHeadModel_firmament(baker.bake(headModel, ModelRotation.X0_Y0));
diff --git a/src/main/kotlin/events/CustomItemModelEvent.kt b/src/main/kotlin/events/CustomItemModelEvent.kt
index d5c08f4..4328d77 100644
--- a/src/main/kotlin/events/CustomItemModelEvent.kt
+++ b/src/main/kotlin/events/CustomItemModelEvent.kt
@@ -17,7 +17,7 @@ data class CustomItemModelEvent(
 		val cache =
 			WeakCache.memoize<ItemStack, ItemModels, Optional<BakedModel>>("CustomItemModels") { stack, models ->
 				val modelId = getModelIdentifier(stack) ?: return@memoize Optional.empty()
-				ErrorUtil.softCheck("Model Id needs to have an inventory variant") { modelId.variant() == "inventory" }
+				ErrorUtil.softCheck("Model Id needs to have an inventory variant", modelId.variant() == "inventory")
 				val bakedModel = models.getModel(modelId.id)
 				if (bakedModel == null || bakedModel === models.missingModelSupplier.get()) return@memoize Optional.empty()
 				Optional.of(bakedModel)
diff --git a/src/main/kotlin/features/texturepack/BakedModelExtra.kt b/src/main/kotlin/features/texturepack/BakedModelExtra.kt
index 32f419a..6305748 100644
--- a/src/main/kotlin/features/texturepack/BakedModelExtra.kt
+++ b/src/main/kotlin/features/texturepack/BakedModelExtra.kt
@@ -1,11 +1,30 @@
-
 package moe.nea.firmament.features.texturepack
 
+import net.fabricmc.fabric.api.renderer.v1.model.WrapperBakedModel as WrapperBakedModelFabric
 import net.minecraft.client.render.model.BakedModel
+import net.minecraft.client.render.model.WrapperBakedModel
+import moe.nea.firmament.util.ErrorUtil
 
 interface BakedModelExtra {
+	companion object {
+		@JvmStatic
+		fun cast(originalModel: BakedModel): BakedModelExtra? {
+			var p = originalModel
+			for (i in 0..256) {
+				p = when (p) {
+					is BakedModelExtra -> return p
+					is WrapperBakedModel -> p.wrapped
+					is WrapperBakedModelFabric -> WrapperBakedModelFabric.unwrap(p)
+					else -> break
+				}
+			}
+			ErrorUtil.softError("Could not find a baked model for $originalModel")
+			return null
+		}
+	}
+
 	var tintOverrides_firmament: TintOverrides?
 
 	fun getHeadModel_firmament(): BakedModel?
-    fun setHeadModel_firmament(headModel: BakedModel?)
+	fun setHeadModel_firmament(headModel: BakedModel?)
 }
diff --git a/src/main/kotlin/features/texturepack/TintOverrides.kt b/src/main/kotlin/features/texturepack/TintOverrides.kt
index 8006db8..85fcae4 100644
--- a/src/main/kotlin/features/texturepack/TintOverrides.kt
+++ b/src/main/kotlin/features/texturepack/TintOverrides.kt
@@ -3,7 +3,6 @@ package moe.nea.firmament.features.texturepack
 import com.google.gson.JsonObject
 import com.google.gson.JsonPrimitive
 import moe.nea.firmament.util.ErrorUtil
-import moe.nea.firmament.util.assertNotNullOr
 
 data class TintOverrides(
 	val layerMap: Map<Int, TintOverride> = mapOf()
@@ -14,21 +13,22 @@ data class TintOverrides(
 		val EMPTY = TintOverrides()
 		private val threadLocal = object : ThreadLocal<TintOverrides>() {}
 		fun enter(overrides: TintOverrides?) {
-			ErrorUtil.softCheck("Double entered tintOverrides") {
-				threadLocal.get() == null
-			}
+			ErrorUtil.softCheck("Double entered tintOverrides",
+			                    threadLocal.get() == null)
 			threadLocal.set(overrides ?: EMPTY)
 		}
 
 		fun exit(overrides: TintOverrides?) {
-			ErrorUtil.softCheck("Exited with non matching enter tintOverrides") {
-				threadLocal.get() == (overrides ?: EMPTY)
-			}
+			ErrorUtil.softCheck("Exited with non matching enter tintOverrides",
+			                    threadLocal.get() == (overrides ?: EMPTY))
 			threadLocal.remove()
 		}
 
-		fun getCurrentOverrides() =
-			assertNotNullOr(threadLocal.get(), "Got current tintOverrides without entering") { EMPTY }
+		fun getCurrentOverrides(): TintOverrides {
+			return ErrorUtil.notNullOr(threadLocal.get(), "Got current tintOverrides without entering") {
+				EMPTY
+			}
+		}
 
 		fun parse(jsonObject: JsonObject): TintOverrides {
 			val map = mutableMapOf<Int, TintOverride>()
diff --git a/src/main/kotlin/util/ErrorUtil.kt b/src/main/kotlin/util/ErrorUtil.kt
index afecf25..b06093b 100644
--- a/src/main/kotlin/util/ErrorUtil.kt
+++ b/src/main/kotlin/util/ErrorUtil.kt
@@ -1,25 +1,46 @@
+@file:OptIn(ExperimentalContracts::class)
+
 package moe.nea.firmament.util
 
+import kotlin.contracts.ExperimentalContracts
+import kotlin.contracts.InvocationKind
+import kotlin.contracts.contract
 import moe.nea.firmament.Firmament
 
+@Suppress("NOTHING_TO_INLINE") // Suppressed since i want the logger to not pick up the ErrorUtil stack-frame
 object ErrorUtil {
 	var aggressiveErrors = run {
 		Thread.currentThread().stackTrace.any { it.className.startsWith("org.junit.") } || Firmament.DEBUG
+			|| ErrorUtil::class.java.desiredAssertionStatus()
+	}
+
+	inline fun softCheck(message: String, check: Boolean) {
+		if (!check) softError(message)
 	}
 
-	inline fun softCheck(message: String, func: () -> Boolean) {
+	inline fun lazyCheck(message: String, func: () -> Boolean) {
+		contract {
+			callsInPlace(func, InvocationKind.AT_MOST_ONCE)
+		}
 		if (!aggressiveErrors) return
 		if (func()) return
 		error(message)
 	}
 
-	@Suppress("NOTHING_TO_INLINE") // Suppressed since i want the logger to not pick up the ErrorUtil stack-frame
+	inline fun softError(message: String, exception: Throwable) {
+		if (aggressiveErrors) throw IllegalStateException(message, exception)
+		else Firmament.logger.error(message, exception)
+	}
+
 	inline fun softError(message: String) {
 		if (aggressiveErrors) error(message)
 		else Firmament.logger.error(message)
 	}
 
 	inline fun <T : Any> notNullOr(nullable: T?, message: String, orElse: () -> T): T {
+		contract {
+			callsInPlace(orElse, InvocationKind.AT_MOST_ONCE)
+		}
 		if (nullable == null) {
 			softError(message)
 			return orElse()
diff --git a/src/main/resources/firmament.accesswidener b/src/main/resources/firmament.accesswidener
index 5d64acc..f1a897b 100644
--- a/src/main/resources/firmament.accesswidener
+++ b/src/main/resources/firmament.accesswidener
@@ -28,3 +28,4 @@ accessible method net/minecraft/registry/entry/RegistryEntry$Reference setRegist
 accessible method net/minecraft/entity/LivingEntity getHitbox ()Lnet/minecraft/util/math/Box;
 accessible method net/minecraft/registry/entry/RegistryEntryList$Named <init> (Lnet/minecraft/registry/entry/RegistryEntryOwner;Lnet/minecraft/registry/tag/TagKey;)V
 accessible method net/minecraft/registry/entry/RegistryEntry$Reference setValue (Ljava/lang/Object;)V
+accessible field net/minecraft/client/render/model/WrapperBakedModel wrapped Lnet/minecraft/client/render/model/BakedModel;