send changes to equipment correctly on event cancel or modification (#…

…3758)
SpongePowered · Oct 12, 2022 · e73f78a · e73f78a
1 parent 2ea2682
commit e73f78a
Showing 1 changed file with 23 additions and 2 deletions.
diff --git a/...va/org/spongepowered/common/mixin/inventory/event/entity/LivingEntityMixin_Inventory.java b/...va/org/spongepowered/common/mixin/inventory/event/entity/LivingEntityMixin_Inventory.java
@@ -84,7 +84,7 @@ protected LivingEntityMixin_Inventory(final EntityType<?> param0, final Level pa
     @Inject(method = "handleEquipmentChanges",
             at = @At(value = "INVOKE", remap = false, target = "Ljava/util/Map;forEach(Ljava/util/function/BiConsumer;)V"))
     protected void inventory$onHandleEquipmentChanges(final Map<EquipmentSlot, ItemStack> map, final CallbackInfo ci) {
-        map.entrySet().removeIf(entry -> {
+        map.entrySet().forEach(entry -> {
             final Slot slotAdapter = this.impl$getSpongeSlot(entry.getKey());
             ItemStack oldStack = null;
             switch (entry.getKey().getType()) {
@@ -95,7 +95,7 @@ protected LivingEntityMixin_Inventory(final EntityType<?> param0, final Level pa
                     oldStack = this.shadow$getLastArmorItem(entry.getKey());
                     break;
             }
-            return this.impl$throwEquipmentEvent(entry.getKey(), slotAdapter, entry.getValue(), oldStack);
+            entry.setValue(this.impl$callEquipmentEvent(entry.getKey(), slotAdapter, entry.getValue(), oldStack));
         });
     }
 
@@ -104,6 +104,27 @@ protected LivingEntityMixin_Inventory(final EntityType<?> param0, final Level pa
         this.shadow$completeUsingItem();
     }
 
+    private ItemStack impl$callEquipmentEvent(final EquipmentSlot equipmentslottype, final Slot slot, final ItemStack newStack, final ItemStack oldStack) {
+        final ChangeEntityEquipmentEvent event = InventoryEventFactory.callChangeEntityEquipmentEvent((LivingEntity) (Object) this,
+                ItemStackUtil.snapshotOf(oldStack), ItemStackUtil.snapshotOf(newStack), slot);
+        if (event.isCancelled()) {
+            this.shadow$setItemSlot(equipmentslottype, oldStack);
+            return oldStack;
+        }
+        final Transaction<@NonNull ItemStackSnapshot> transaction = event.transaction();
+        if (!transaction.isValid()) {
+            this.shadow$setItemSlot(equipmentslottype, oldStack);
+            return oldStack;
+        }
+        final Optional<ItemStackSnapshot> optional = transaction.custom();
+        if (optional.isPresent()) {
+            final ItemStack custom = ItemStackUtil.fromSnapshotToNative(optional.get());
+            this.shadow$setItemSlot(equipmentslottype, custom);
+            return custom;
+        }
+        return newStack;
+    }
+
     private boolean impl$throwEquipmentEvent(final EquipmentSlot equipmentslottype, final Slot slot, final ItemStack newStack, final ItemStack oldStack) {
         final ChangeEntityEquipmentEvent event = InventoryEventFactory.callChangeEntityEquipmentEvent((LivingEntity) (Object) this,
                 ItemStackUtil.snapshotOf(oldStack), ItemStackUtil.snapshotOf(newStack), slot);