Also handle aqua affinity

Author: kennytv

bukkit/src/main/java/com/viaversion/viaversion/bukkit/listeners/v1_20_5to1_21/PlayerChangeItemListener.java

@@ -41,6 +41,7 @@
 public final class PlayerChangeItemListener extends ViaBukkitListener {
 
     private final Enchantment efficiency = Enchantment.getByKey(NamespacedKey.minecraft("efficiency"));
+    private final Enchantment aquaAffinity = Enchantment.getByKey(NamespacedKey.minecraft("aqua_affinity"));
     private final Enchantment depthStrider = Enchantment.getByKey(NamespacedKey.minecraft("depth_strider"));
     private final Enchantment soulSpeed = Enchantment.getByKey(NamespacedKey.minecraft("soul_speed"));
     private final Enchantment swiftSneak = Enchantment.getByKey(NamespacedKey.minecraft("swift_sneak"));
@@ -61,6 +62,8 @@ public void onPlayerInventorySlotChangedEvent(final PlayerInventorySlotChangeEve
             sendAttributeUpdate(player, item, Slot.BOOTS);
         } else if (slot == 37) {
             sendAttributeUpdate(player, item, Slot.LEGGINGS);
+        } else if (slot == 39) {
+            sendAttributeUpdate(player, item, Slot.HELMET);
         }
     }
 
@@ -84,11 +87,13 @@ private void sendAttributeUpdate(final Player player, @Nullable final ItemStack
 
         final EfficiencyAttributeStorage.ActiveEnchants activeEnchants = storage.activeEnchants();
         int efficiencyLevel = activeEnchants.efficiency().level();
+        int aquaAffinityLevel = activeEnchants.aquaAffinity().level();
         int soulSpeedLevel = activeEnchants.soulSpeed().level();
         int swiftSneakLevel = activeEnchants.swiftSneak().level();
         int depthStriderLevel = activeEnchants.depthStrider().level();
         switch (slot) {
             case HAND -> efficiencyLevel = item != null ? item.getEnchantmentLevel(efficiency) : 0;
+            case HELMET -> aquaAffinityLevel = item != null ? item.getEnchantmentLevel(aquaAffinity) : 0;
             case LEGGINGS -> swiftSneakLevel = item != null && swiftSneak != null ? item.getEnchantmentLevel(swiftSneak) : 0;
             case BOOTS -> {
                 depthStriderLevel = item != null && depthStrider != null ? item.getEnchantmentLevel(depthStrider) : 0;
@@ -97,10 +102,10 @@ private void sendAttributeUpdate(final Player player, @Nullable final ItemStack
                 //soulSpeedLevel = item != null && soulSpeed != null ? item.getEnchantmentLevel(soulSpeed) : 0;
             }
         }
-        storage.setEnchants(player.getEntityId(), connection, efficiencyLevel, soulSpeedLevel, swiftSneakLevel, depthStriderLevel);
+        storage.setEnchants(player.getEntityId(), connection, efficiencyLevel, soulSpeedLevel, swiftSneakLevel, aquaAffinityLevel, depthStriderLevel);
     }
 
     private enum Slot {
-        HAND, BOOTS, LEGGINGS
+        HAND, BOOTS, LEGGINGS, HELMET
     }
 }

common/src/main/java/com/viaversion/viaversion/protocols/v1_20_5to1_21/storage/EfficiencyAttributeStorage.java

@@ -30,23 +30,27 @@ public final class EfficiencyAttributeStorage implements StorableObject {
     private static final EnchantAttributeModifier EFFICIENCY = new EnchantAttributeModifier("minecraft:enchantment.efficiency/mainhand", 19, 0, level -> (level * level) + 1);
     private static final EnchantAttributeModifier SOUL_SPEED = new EnchantAttributeModifier("minecraft:enchantment.soul_speed", 21, 0.1, level -> 0.04D + ((level - 1) * 0.01D));
     private static final EnchantAttributeModifier SWIFT_SNEAK = new EnchantAttributeModifier("minecraft:enchantment.swift_sneak", 25, 0.3, level -> level * 0.15D);
+    private static final EnchantAttributeModifier AQUA_AFFINITY = new EnchantAttributeModifier("minecraft:enchantment.aqua_affinity", 28, 0.2, level -> level * 4, (byte) 2);
     private static final EnchantAttributeModifier DEPTH_STRIDER = new EnchantAttributeModifier("minecraft:enchantment.depth_strider", 30, 0, level -> level / 3D);
     private static final ActiveEnchants DEFAULT = new ActiveEnchants(-1,
         new ActiveEnchant(EFFICIENCY, 0),
         new ActiveEnchant(SOUL_SPEED, 0),
         new ActiveEnchant(SWIFT_SNEAK, 0),
+        new ActiveEnchant(AQUA_AFFINITY, 0),
         new ActiveEnchant(DEPTH_STRIDER, 0)
     );
     private final Object lock = new Object();
-    private volatile ActiveEnchants activeEnchants = DEFAULT;
     private volatile boolean attributesSent = true;
     private volatile boolean loginSent;
+    private ActiveEnchants activeEnchants = DEFAULT;
 
-    public void setEnchants(final int entityId, final UserConnection connection, final int efficiency, final int soulSpeed, final int swiftSneak, final int depthStrider) {
+    public void setEnchants(final int entityId, final UserConnection connection, final int efficiency, final int soulSpeed,
+                            final int swiftSneak, final int aquaAffinity, final int depthStrider) {
         // Always called from the main thread
         if (efficiency == activeEnchants.efficiency.level
             && soulSpeed == activeEnchants.soulSpeed.level
             && swiftSneak == activeEnchants.swiftSneak.level
+            && aquaAffinity == activeEnchants.aquaAffinity.level
             && depthStrider == activeEnchants.depthStrider.level) {
             return;
         }
@@ -56,6 +60,7 @@ public void setEnchants(final int entityId, final UserConnection connection, fin
                 new ActiveEnchant(EFFICIENCY, efficiency),
                 new ActiveEnchant(SOUL_SPEED, soulSpeed),
                 new ActiveEnchant(SWIFT_SNEAK, swiftSneak),
+                new ActiveEnchant(AQUA_AFFINITY, aquaAffinity),
                 new ActiveEnchant(DEPTH_STRIDER, depthStrider)
             );
             this.attributesSent = false;
@@ -99,7 +104,7 @@ private void sendAttributesPacket(final UserConnection connection) {
                 attributesPacket.write(Types.VAR_INT, 1); // Modifiers
                 attributesPacket.write(Types.STRING, modifier.key);
                 attributesPacket.write(Types.DOUBLE, enchant.modifier.modifierFunction.get(enchant.level));
-                attributesPacket.write(Types.BYTE, (byte) 0); // 'Add' operation
+                attributesPacket.write(Types.BYTE, modifier.operation);
             } else {
                 attributesPacket.write(Types.VAR_INT, 0); // Modifiers
             }
@@ -108,14 +113,18 @@ private void sendAttributesPacket(final UserConnection connection) {
         attributesPacket.scheduleSend(Protocol1_20_5To1_21.class);
     }
 
-    public record ActiveEnchants(int entityId, ActiveEnchant efficiency, ActiveEnchant soulSpeed,
-                                 ActiveEnchant swiftSneak, ActiveEnchant depthStrider) {
+    public record ActiveEnchants(int entityId, ActiveEnchant efficiency, ActiveEnchant aquaAffinity,
+                                 ActiveEnchant soulSpeed, ActiveEnchant swiftSneak, ActiveEnchant depthStrider) {
     }
 
     public record ActiveEnchant(EnchantAttributeModifier modifier, int level) {
     }
 
-    public record EnchantAttributeModifier(String key, int attributeId, double baseValue, LevelToModifier modifierFunction) {
+    public record EnchantAttributeModifier(String key, int attributeId, double baseValue, LevelToModifier modifierFunction, byte operation) {
+
+        private EnchantAttributeModifier(String key, int attributeId, double baseValue, LevelToModifier modifierFunction) {
+            this(key, attributeId, baseValue, modifierFunction, (byte) 0);
+        }
     }
 
     @FunctionalInterface