Fixed bug with Cosmetic Armor support

Closes #23

3 files changed, 10 insertions, 6 deletions

diff --git a/src/main/java/shcm/shsupercm/fabric/citresewn/OptionalCompat.java b/src/main/java/shcm/shsupercm/fabric/citresewn/OptionalCompat.java
index 7558d03..ea77f5e 100644
--- a/src/main/java/shcm/shsupercm/fabric/citresewn/OptionalCompat.java
+++ b/src/main/java/shcm/shsupercm/fabric/citresewn/OptionalCompat.java
@@ -4,6 +4,7 @@ import io.github.apace100.cosmetic_armor.CosmeticArmor;
 import net.minecraft.entity.EquipmentSlot;
 import net.minecraft.entity.LivingEntity;
 import net.minecraft.item.ItemStack;
+import net.minecraft.item.Items;
 
 import java.util.function.Predicate;
 
@@ -18,10 +19,10 @@ public final class OptionalCompat {
         compatCosmeticArmor = isLoaded.test("cosmetic-armor") ? CompatCosmeticArmor.impl() : null;
     }
 
-    public static ItemStack getCosmeticArmor(ItemStack original, LivingEntity entity, EquipmentSlot slot) {
+    public static ItemStack getCosmeticArmor(ItemStack original, LivingEntity entity, EquipmentSlot slot, boolean elytra) {
         if (INSTANCE().compatCosmeticArmor != null) {
             ItemStack stackInCosmeticSlot = INSTANCE().compatCosmeticArmor.getStackInCosmeticSlot(entity, slot);
-            if (!stackInCosmeticSlot.isEmpty())
+            if (!stackInCosmeticSlot.isEmpty() && (!elytra || stackInCosmeticSlot.isOf(Items.ELYTRA)))
                 return stackInCosmeticSlot;
         }
 
diff --git a/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citarmor/ArmorFeatureRendererMixin.java b/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citarmor/ArmorFeatureRendererMixin.java
index d8b78a4..f4f4fca 100644
--- a/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citarmor/ArmorFeatureRendererMixin.java
+++ b/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citarmor/ArmorFeatureRendererMixin.java
@@ -33,7 +33,7 @@ public class ArmorFeatureRendererMixin<T extends LivingEntity, M extends BipedEn
         ItemStack itemStack = entity.getEquippedStack(armorSlot);
 
         //compat Cosmetic Armor
-        itemStack = OptionalCompat.getCosmeticArmor(itemStack, entity, armorSlot);
+        itemStack = OptionalCompat.getCosmeticArmor(itemStack, entity, armorSlot, false);
 
         Map<String, Identifier> armorTextures = CITResewn.INSTANCE.activeCITs.getArmorTextures(itemStack, entity.world, entity);
         if (armorTextures != null) {
diff --git a/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citelytra/ElytraFeatureRendererMixin.java b/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citelytra/ElytraFeatureRendererMixin.java
index b5ef126..2214598 100644
--- a/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citelytra/ElytraFeatureRendererMixin.java
+++ b/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citelytra/ElytraFeatureRendererMixin.java
@@ -7,6 +7,7 @@ import net.minecraft.client.util.math.MatrixStack;
 import net.minecraft.entity.EquipmentSlot;
 import net.minecraft.entity.LivingEntity;
 import net.minecraft.item.ItemStack;
+import net.minecraft.item.Items;
 import net.minecraft.util.Identifier;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.injection.At;
@@ -32,7 +33,7 @@ public class ElytraFeatureRendererMixin {
         ItemStack itemStack = livingEntity.getEquippedStack(EquipmentSlot.CHEST);
 
         //compat Cosmetic Armor
-        itemStack = OptionalCompat.getCosmeticArmor(itemStack, livingEntity, EquipmentSlot.CHEST);
+        itemStack = OptionalCompat.getCosmeticArmor(itemStack, livingEntity, EquipmentSlot.CHEST, true);
 
         this.elytraItemCached = new WeakReference<>(itemStack);
         this.livingEntityCached = new WeakReference<>(livingEntity);
@@ -44,9 +45,11 @@ public class ElytraFeatureRendererMixin {
             return RenderLayer.getArmorCutoutNoCull(originalIdentifier);
 
         ItemStack itemStack = this.elytraItemCached.get();
-        LivingEntity livingEntity = livingEntityCached.get();
-        if (itemStack != null && livingEntity != null) {
+        LivingEntity livingEntity = this.livingEntityCached.get();
+        if (itemStack != null && itemStack.isOf(Items.ELYTRA) && livingEntity != null) {
             Identifier elytraTexture = CITResewn.INSTANCE.activeCITs.getElytraTexture(itemStack, livingEntity.world, livingEntity);
+            this.elytraItemCached = new WeakReference<>(null);
+            this.livingEntityCached = new WeakReference<>(null);
             if (elytraTexture != null)
                 return RenderLayer.getArmorCutoutNoCull(elytraTexture);
         }