From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: Jake Potrebic <jake.m.potrebic@gmail.com>
Date: Wed, 22 May 2024 10:01:19 -0700
Subject: [PATCH] Fix equipment slot and group API

Adds the following:
- Add test for EquipmentSlotGroup
- Expose LivingEntity#canUseSlot

Co-authored-by: SoSeDiK <mrsosedik@gmail.com>

diff --git a/src/main/java/org/bukkit/craftbukkit/entity/CraftLivingEntity.java b/src/main/java/org/bukkit/craftbukkit/entity/CraftLivingEntity.java
index a62d17b72c675120b447e625cb3dc437681bdf20..f16067b674118a47735ad22797988d50b4415040 100644
--- a/src/main/java/org/bukkit/craftbukkit/entity/CraftLivingEntity.java
+++ b/src/main/java/org/bukkit/craftbukkit/entity/CraftLivingEntity.java
@@ -1230,4 +1230,11 @@ public class CraftLivingEntity extends CraftEntity implements LivingEntity {
         this.getHandle().setYBodyRot(bodyYaw);
     }
     // Paper end - body yaw API
+
+    // Paper start - Expose canUseSlot
+    @Override
+    public boolean canUseEquipmentSlot(org.bukkit.inventory.EquipmentSlot slot) {
+        return this.getHandle().canUseSlot(org.bukkit.craftbukkit.CraftEquipmentSlot.getNMS(slot));
+    }
+    // Paper end - Expose canUseSlot
 }
diff --git a/src/main/java/org/bukkit/craftbukkit/inventory/CraftInventoryPlayer.java b/src/main/java/org/bukkit/craftbukkit/inventory/CraftInventoryPlayer.java
index e62baea16df017f1e394e3c706157e158066eb93..656c9a6d8cd42891141ee29ec91ab5d166051ed6 100644
--- a/src/main/java/org/bukkit/craftbukkit/inventory/CraftInventoryPlayer.java
+++ b/src/main/java/org/bukkit/craftbukkit/inventory/CraftInventoryPlayer.java
@@ -135,6 +135,10 @@ public class CraftInventoryPlayer extends CraftInventory implements org.bukkit.i
             case HEAD:
                 this.setHelmet(item);
                 break;
+                // Paper start
+            case BODY:
+                throw new IllegalArgumentException("BODY is not valid for players!");
+                // Paper end
             default:
                 throw new IllegalArgumentException("Could not set slot " + slot + " - not a valid slot for PlayerInventory");
         }
@@ -162,6 +166,10 @@ public class CraftInventoryPlayer extends CraftInventory implements org.bukkit.i
                 return java.util.Objects.requireNonNullElseGet(this.getChestplate(), () -> new ItemStack(org.bukkit.Material.AIR)); // Paper - make nonnull
             case HEAD:
                 return java.util.Objects.requireNonNullElseGet(this.getHelmet(), () -> new ItemStack(org.bukkit.Material.AIR)); // Paper - make nonnull
+                // Paper start
+            case BODY:
+                throw new IllegalArgumentException("BODY is not valid for players!");
+                // Paper end
             default:
                 throw new IllegalArgumentException("Could not get slot " + slot + " - not a valid slot for PlayerInventory");
         }
diff --git a/src/main/java/org/bukkit/craftbukkit/inventory/CraftMetaItem.java b/src/main/java/org/bukkit/craftbukkit/inventory/CraftMetaItem.java
index a6e2281bfac94f1e19836d9c8415d8270387b16d..a4d6c95e55722c9a0e381d7b84916787240cea8e 100644
--- a/src/main/java/org/bukkit/craftbukkit/inventory/CraftMetaItem.java
+++ b/src/main/java/org/bukkit/craftbukkit/inventory/CraftMetaItem.java
@@ -1703,7 +1703,7 @@ class CraftMetaItem implements ItemMeta, Damageable, Repairable, BlockDataMeta {
         if (this.attributeModifiers == null) return LinkedHashMultimap.create(); // Paper - don't change the components
         SetMultimap<Attribute, AttributeModifier> result = LinkedHashMultimap.create();
         for (Map.Entry<Attribute, AttributeModifier> entry : this.attributeModifiers.entries()) {
-            if (entry.getValue().getSlot() == null || entry.getValue().getSlot() == slot) {
+            if (entry.getValue().getSlotGroup().test(slot)) { // Paper - correctly test slot against group
                 result.put(entry.getKey(), entry.getValue());
             }
         }
@@ -1777,9 +1777,7 @@ class CraftMetaItem implements ItemMeta, Damageable, Repairable, BlockDataMeta {
 
         while (iter.hasNext()) {
             Map.Entry<Attribute, AttributeModifier> entry = iter.next();
-            // Explicitly match against null because (as of MC 1.13) AttributeModifiers without a -
-            // set slot are active in any slot.
-            if (entry.getValue().getSlot() == null || entry.getValue().getSlot() == slot) {
+            if (entry.getValue().getSlotGroup().test(slot)) { // Paper - correctly test slot against group
                 iter.remove();
                 ++removed;
             }
diff --git a/src/test/java/io/papermc/paper/inventory/item/EquipmentSlotGroupTest.java b/src/test/java/io/papermc/paper/inventory/item/EquipmentSlotGroupTest.java
new file mode 100644
index 0000000000000000000000000000000000000000..a26dc8f4787276468417196440cb8f73670a56c7
--- /dev/null
+++ b/src/test/java/io/papermc/paper/inventory/item/EquipmentSlotGroupTest.java
@@ -0,0 +1,53 @@
+package io.papermc.paper.inventory.item;
+
+import java.lang.reflect.Field;
+import java.lang.reflect.Modifier;
+import java.util.ArrayList;
+import java.util.List;
+import net.minecraft.world.entity.EquipmentSlot;
+import org.bukkit.craftbukkit.CraftEquipmentSlot;
+import org.bukkit.inventory.EquipmentSlotGroup;
+import org.bukkit.support.environment.Normal;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+import org.junit.jupiter.params.provider.MethodSource;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+
+@Normal
+class EquipmentSlotGroupTest {
+
+    static List<EquipmentSlotGroup> apiValues() throws ReflectiveOperationException {
+        final List<EquipmentSlotGroup> apiValues = new ArrayList<>();
+        for (final Field field : EquipmentSlotGroup.class.getDeclaredFields()) {
+            if (!Modifier.isStatic(field.getModifiers()) || !Modifier.isFinal(field.getModifiers()) || !field.getType().equals(EquipmentSlotGroup.class)) {
+                continue;
+            }
+            apiValues.add((EquipmentSlotGroup) field.get(null));
+        }
+        if (apiValues.isEmpty()) {
+            throw new RuntimeException("Didn't find any api " + EquipmentSlotGroup.class.getSimpleName());
+        }
+        return apiValues;
+    }
+
+    @ParameterizedTest
+    @MethodSource("apiValues")
+    void testBukkitToNms(final EquipmentSlotGroup slotGroup) {
+        final net.minecraft.world.entity.EquipmentSlotGroup nmsGroup = CraftEquipmentSlot.getNMSGroup(slotGroup);
+        assertNotNull(nmsGroup, "No nms slot group found for " + slotGroup);
+        assertEquals(nmsGroup.getSerializedName(), slotGroup.toString(), "slot group name mismatch");
+        for (final EquipmentSlot slot : EquipmentSlot.values()) {
+            assertEquals(nmsGroup.test(slot), slotGroup.test(CraftEquipmentSlot.getSlot(slot)));
+        }
+    }
+
+    @ParameterizedTest
+    @EnumSource(net.minecraft.world.entity.EquipmentSlotGroup.class)
+    void testNmsToBukkit(final net.minecraft.world.entity.EquipmentSlotGroup slotGroup) {
+        final EquipmentSlotGroup apiGroup = CraftEquipmentSlot.getSlot(slotGroup);
+        assertNotNull(apiGroup, "No api slot group found for " + slotGroup);
+        assertEquals(apiGroup.toString(), slotGroup.getSerializedName(), "slot group name mismatch");
+    }
+}