CIT Enchantment applying on armor and elytra

#15

4 files changed, 74 insertions, 6 deletions

diff --git a/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citenchantment/ArmorFeatureRendererMixin.java b/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citenchantment/ArmorFeatureRendererMixin.java
new file mode 100644
index 0000000..b4d81c7
--- /dev/null
+++ b/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citenchantment/ArmorFeatureRendererMixin.java
@@ -0,0 +1,33 @@
+package shcm.shsupercm.fabric.citresewn.mixin.citenchantment;
+
+import net.minecraft.client.render.VertexConsumerProvider;
+import net.minecraft.client.render.entity.feature.ArmorFeatureRenderer;
+import net.minecraft.client.render.entity.model.BipedEntityModel;
+import net.minecraft.client.util.math.MatrixStack;
+import net.minecraft.entity.EquipmentSlot;
+import net.minecraft.entity.LivingEntity;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
+import shcm.shsupercm.fabric.citresewn.CITResewn;
+import shcm.shsupercm.fabric.citresewn.config.CITResewnConfig;
+import shcm.shsupercm.fabric.citresewn.pack.cits.CITEnchantment;
+
+@Mixin(ArmorFeatureRenderer.class)
+public class ArmorFeatureRendererMixin<T extends LivingEntity, M extends BipedEntityModel<T>, A extends BipedEntityModel<T>> {
+    @Inject(method = "renderArmor", at = @At("HEAD"))
+    private void setAppliedContextAndStartApplyingArmor(MatrixStack matrices, VertexConsumerProvider vertexConsumers, T livingEntity, EquipmentSlot armorSlot, int light, A model, CallbackInfo ci) {
+        if (CITResewnConfig.INSTANCE().enabled && CITResewn.INSTANCE.activeCITs != null) {
+            CITResewn.INSTANCE.activeCITs.setEnchantmentAppliedContextCached(livingEntity.getEquippedStack(armorSlot), livingEntity.world, livingEntity);
+            CITEnchantment.shouldApply = true;
+        }
+    }
+
+    @Inject(method = "renderArmor", at = @At("RETURN"))
+    private void stopApplyingArmor(MatrixStack matrices, VertexConsumerProvider vertexConsumers, T livingEntity, EquipmentSlot armorSlot, int light, A model, CallbackInfo ci) {
+        CITEnchantment.shouldApply = false;
+        if (CITResewn.INSTANCE.activeCITs != null)
+            CITResewn.INSTANCE.activeCITs.setEnchantmentAppliedContextCached(null, null, null);
+    }
+}
diff --git a/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citenchantment/ElytraFeatureRendererMixin.java b/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citenchantment/ElytraFeatureRendererMixin.java
new file mode 100644
index 0000000..03fb4c0
--- /dev/null
+++ b/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citenchantment/ElytraFeatureRendererMixin.java
@@ -0,0 +1,33 @@
+package shcm.shsupercm.fabric.citresewn.mixin.citenchantment;
+
+import net.minecraft.client.render.VertexConsumerProvider;
+import net.minecraft.client.render.entity.feature.ElytraFeatureRenderer;
+import net.minecraft.client.util.math.MatrixStack;
+import net.minecraft.entity.EquipmentSlot;
+import net.minecraft.entity.LivingEntity;
+import net.minecraft.item.ItemStack;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
+import shcm.shsupercm.fabric.citresewn.CITResewn;
+import shcm.shsupercm.fabric.citresewn.config.CITResewnConfig;
+import shcm.shsupercm.fabric.citresewn.pack.cits.CITEnchantment;
+
+@Mixin(ElytraFeatureRenderer.class)
+public class ElytraFeatureRendererMixin {
+    @Inject(method = "render(Lnet/minecraft/client/util/math/MatrixStack;Lnet/minecraft/client/render/VertexConsumerProvider;ILnet/minecraft/entity/LivingEntity;FFFFFF)V", at = @At("HEAD"))
+    private void setAppliedContextAndStartApplyingElytra(MatrixStack matrixStack, VertexConsumerProvider vertexConsumerProvider, int i, LivingEntity livingEntity, float f, float g, float h, float j, float k, float l, CallbackInfo ci) {
+        if (CITResewnConfig.INSTANCE().enabled && CITResewn.INSTANCE.activeCITs != null) {
+            CITResewn.INSTANCE.activeCITs.setEnchantmentAppliedContextCached(livingEntity.getEquippedStack(EquipmentSlot.CHEST), livingEntity.world, livingEntity);
+            CITEnchantment.shouldApply = true;
+        }
+    }
+
+    @Inject(method = "render(Lnet/minecraft/client/util/math/MatrixStack;Lnet/minecraft/client/render/VertexConsumerProvider;ILnet/minecraft/entity/LivingEntity;FFFFFF)V", at = @At("RETURN"))
+    private void stopApplyingElytra(MatrixStack matrixStack, VertexConsumerProvider vertexConsumerProvider, int i, LivingEntity livingEntity, float f, float g, float h, float j, float k, float l, CallbackInfo ci) {
+        CITEnchantment.shouldApply = false;
+        if (CITResewn.INSTANCE.activeCITs != null)
+            CITResewn.INSTANCE.activeCITs.setEnchantmentAppliedContextCached(null, null, null);
+    }
+}
diff --git a/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citenchantment/ItemRendererMixin.java b/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citenchantment/ItemRendererMixin.java
index 7af6c22..eb9f714 100644
--- a/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citenchantment/ItemRendererMixin.java
+++ b/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citenchantment/ItemRendererMixin.java
@@ -43,7 +43,7 @@ public class ItemRendererMixin {
     private static void getArmorGlintConsumer(VertexConsumerProvider provider, RenderLayer layer, boolean solid, boolean glint, CallbackInfoReturnable<VertexConsumer> cir) {
         if (!CITEnchantment.shouldApply)
             return;
-        VertexConsumer vertexConsumer = solid ? CITEnchantment.GlintRenderLayer.ARMOR_GLINT.tryApply(cir.getReturnValue(), provider) : CITEnchantment.GlintRenderLayer.ARMOR_ENTITY_GLINT.tryApply(cir.getReturnValue(), provider);
+        VertexConsumer vertexConsumer = solid ? CITEnchantment.GlintRenderLayer.ARMOR_GLINT.tryApply(cir.getReturnValue(), layer, provider) : CITEnchantment.GlintRenderLayer.ARMOR_ENTITY_GLINT.tryApply(cir.getReturnValue(), layer, provider);
         if (vertexConsumer != null)
             cir.setReturnValue(vertexConsumer);
     }
@@ -52,7 +52,7 @@ public class ItemRendererMixin {
     private static void getCompassGlintConsumer(VertexConsumerProvider provider, RenderLayer layer, MatrixStack.Entry entry, CallbackInfoReturnable<VertexConsumer> cir) {
         if (!CITEnchantment.shouldApply)
             return;
-        VertexConsumer vertexConsumer = CITEnchantment.GlintRenderLayer.GLINT.tryApply(null, provider);
+        VertexConsumer vertexConsumer = CITEnchantment.GlintRenderLayer.GLINT.tryApply(null, layer, provider);
         if (vertexConsumer != null)
             cir.setReturnValue(VertexConsumers.union(new OverlayVertexConsumer(vertexConsumer, entry.getModel(), entry.getNormal()), cir.getReturnValue()));
     }
@@ -61,7 +61,7 @@ public class ItemRendererMixin {
     private static void getDirectCompassGlintConsumer(VertexConsumerProvider provider, RenderLayer layer, MatrixStack.Entry entry, CallbackInfoReturnable<VertexConsumer> cir) {
         if (!CITEnchantment.shouldApply)
             return;
-        VertexConsumer vertexConsumer = CITEnchantment.GlintRenderLayer.DIRECT_GLINT.tryApply(null, provider);
+        VertexConsumer vertexConsumer = CITEnchantment.GlintRenderLayer.DIRECT_GLINT.tryApply(null, layer, provider);
         if (vertexConsumer != null)
             cir.setReturnValue(VertexConsumers.union(new OverlayVertexConsumer(vertexConsumer, entry.getModel(), entry.getNormal()), cir.getReturnValue()));
     }
@@ -70,7 +70,7 @@ public class ItemRendererMixin {
     private static void getItemGlintConsumer(VertexConsumerProvider provider, RenderLayer layer, boolean solid, boolean glint, CallbackInfoReturnable<VertexConsumer> cir) {
         if (!CITEnchantment.shouldApply)
             return;
-        VertexConsumer vertexConsumer = MinecraftClient.isFabulousGraphicsOrBetter() && layer == TexturedRenderLayers.getItemEntityTranslucentCull() ? CITEnchantment.GlintRenderLayer.GLINT_TRANSLUCENT.tryApply(cir.getReturnValue(), provider) : (solid ? CITEnchantment.GlintRenderLayer.GLINT.tryApply(cir.getReturnValue(), provider) : CITEnchantment.GlintRenderLayer.ENTITY_GLINT.tryApply(cir.getReturnValue(), provider));
+        VertexConsumer vertexConsumer = MinecraftClient.isFabulousGraphicsOrBetter() && layer == TexturedRenderLayers.getItemEntityTranslucentCull() ? CITEnchantment.GlintRenderLayer.GLINT_TRANSLUCENT.tryApply(cir.getReturnValue(), layer, provider) : (solid ? CITEnchantment.GlintRenderLayer.GLINT.tryApply(cir.getReturnValue(), layer, provider) : CITEnchantment.GlintRenderLayer.ENTITY_GLINT.tryApply(cir.getReturnValue(), layer, provider));
         if (vertexConsumer != null)
             cir.setReturnValue(vertexConsumer);
     }
@@ -79,7 +79,7 @@ public class ItemRendererMixin {
     private static void getDirectItemGlintConsumer(VertexConsumerProvider provider, RenderLayer layer, boolean solid, boolean glint, CallbackInfoReturnable<VertexConsumer> cir) {
         if (!CITEnchantment.shouldApply)
             return;
-        VertexConsumer vertexConsumer = solid ? CITEnchantment.GlintRenderLayer.DIRECT_GLINT.tryApply(cir.getReturnValue(), provider) : CITEnchantment.GlintRenderLayer.DIRECT_ENTITY_GLINT.tryApply(cir.getReturnValue(), provider);
+        VertexConsumer vertexConsumer = solid ? CITEnchantment.GlintRenderLayer.DIRECT_GLINT.tryApply(cir.getReturnValue(), layer, provider) : CITEnchantment.GlintRenderLayer.DIRECT_ENTITY_GLINT.tryApply(cir.getReturnValue(), layer, provider);
         if (vertexConsumer != null)
             cir.setReturnValue(vertexConsumer);
     }
diff --git a/src/main/java/shcm/shsupercm/fabric/citresewn/pack/cits/CITEnchantment.java b/src/main/java/shcm/shsupercm/fabric/citresewn/pack/cits/CITEnchantment.java
index cce2f29..2458646 100644
--- a/src/main/java/shcm/shsupercm/fabric/citresewn/pack/cits/CITEnchantment.java
+++ b/src/main/java/shcm/shsupercm/fabric/citresewn/pack/cits/CITEnchantment.java
@@ -157,7 +157,7 @@ public class CITEnchantment extends CIT {
                     layer.build(false));
         }
 
-        public VertexConsumer tryApply(VertexConsumer base, VertexConsumerProvider provider) {
+        public VertexConsumer tryApply(VertexConsumer base, RenderLayer baseLayer, VertexConsumerProvider provider) {
             if (!shouldApply || appliedContext == null)
                 return null;
 
@@ -165,6 +165,8 @@ public class CITEnchantment extends CIT {
                     .map(cit -> provider.getBuffer(cit.renderLayers.get(GlintRenderLayer.this)))
                     .toArray(VertexConsumer[]::new));
 
+            provider.getBuffer(baseLayer); // refresh base layer for armor consumer
+
             return base == null ? applied : VertexConsumers.union(applied, base);
         }
     }