From 0915afe43b179f630a78d4166bbff61132a16bb2 Mon Sep 17 00:00:00 2001
From: Nassim Jahnke <nassim@njahnke.dev>
Date: Fri, 19 Jul 2024 11:26:28 +0200
Subject: [PATCH] Handle 1.21 attribute id changes Fixes #4005

---
 .../viaversion/api/data/MappingDataBase.java        |  4 ++--
 .../rewriter/BlockItemPacketRewriter1_21.java       | 13 ++++++++++---
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/api/src/main/java/com/viaversion/viaversion/api/data/MappingDataBase.java b/api/src/main/java/com/viaversion/viaversion/api/data/MappingDataBase.java
index cd3cf026c..e0fed837d 100644
--- a/api/src/main/java/com/viaversion/viaversion/api/data/MappingDataBase.java
+++ b/api/src/main/java/com/viaversion/viaversion/api/data/MappingDataBase.java
@@ -46,6 +46,7 @@ public class MappingDataBase implements MappingData {
     protected ParticleMappings particleMappings;
     protected BiMappings itemMappings;
     protected BiMappings blockMappings;
+    protected BiMappings attributeMappings;
     protected Mappings blockStateMappings;
     protected Mappings blockEntityMappings;
     protected Mappings soundMappings;
@@ -53,7 +54,6 @@ public class MappingDataBase implements MappingData {
     protected Mappings enchantmentMappings;
     protected Mappings paintingMappings;
     protected Mappings menuMappings;
-    protected Mappings attributeMappings;
     protected Map<RegistryType, List<TagData>> tags;
 
     public MappingDataBase(final String unmappedVersion, final String mappedVersion) {
@@ -76,7 +76,7 @@ public class MappingDataBase implements MappingData {
         menuMappings = loadMappings(data, "menus");
         enchantmentMappings = loadMappings(data, "enchantments");
         paintingMappings = loadMappings(data, "paintings");
-        attributeMappings = loadMappings(data, "attributes");
+        attributeMappings = loadBiMappings(data, "attributes");
 
 
         final CompoundTag unmappedIdentifierData = readUnmappedIdentifiersFile("identifiers-" + unmappedVersion + ".nbt");
diff --git a/common/src/main/java/com/viaversion/viaversion/protocols/v1_20_5to1_21/rewriter/BlockItemPacketRewriter1_21.java b/common/src/main/java/com/viaversion/viaversion/protocols/v1_20_5to1_21/rewriter/BlockItemPacketRewriter1_21.java
index eb96cc491..a552a6b09 100644
--- a/common/src/main/java/com/viaversion/viaversion/protocols/v1_20_5to1_21/rewriter/BlockItemPacketRewriter1_21.java
+++ b/common/src/main/java/com/viaversion/viaversion/protocols/v1_20_5to1_21/rewriter/BlockItemPacketRewriter1_21.java
@@ -41,6 +41,7 @@ import com.viaversion.viaversion.rewriter.BlockRewriter;
 import com.viaversion.viaversion.rewriter.StructuredItemRewriter;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Objects;
 
 public final class BlockItemPacketRewriter1_21 extends StructuredItemRewriter<ClientboundPacket1_20_5, ServerboundPacket1_20_5, Protocol1_20_5To1_21> {
 
@@ -139,9 +140,10 @@ public final class BlockItemPacketRewriter1_21 extends StructuredItemRewriter<Cl
         dataContainer.replaceKey(StructuredDataKey.BUNDLE_CONTENTS1_20_5, StructuredDataKey.BUNDLE_CONTENTS1_21);
         dataContainer.replace(StructuredDataKey.ATTRIBUTE_MODIFIERS1_20_5, StructuredDataKey.ATTRIBUTE_MODIFIERS1_21, attributeModifiers -> {
             final AttributeModifiers1_21.AttributeModifier[] modifiers = Arrays.stream(attributeModifiers.modifiers()).map(modifier -> {
+                final int mappedAttributeId = Protocol1_20_5To1_21.MAPPINGS.getNewAttributeId(modifier.attribute());
                 final AttributeModifiers1_20_5.ModifierData modData = modifier.modifier();
                 final AttributeModifiers1_21.ModifierData updatedModData = new AttributeModifiers1_21.ModifierData(Protocol1_20_5To1_21.mapAttributeUUID(modData.uuid(), modData.name()), modData.amount(), modData.operation());
-                return new AttributeModifiers1_21.AttributeModifier(modifier.attribute(), updatedModData, modifier.slotType());
+                return new AttributeModifiers1_21.AttributeModifier(mappedAttributeId, updatedModData, modifier.slotType());
             }).toArray(AttributeModifiers1_21.AttributeModifier[]::new);
             return new AttributeModifiers1_21(modifiers, attributeModifiers.showInTooltip());
         });
@@ -168,6 +170,11 @@ public final class BlockItemPacketRewriter1_21 extends StructuredItemRewriter<Cl
         dataContainer.remove(StructuredDataKey.JUKEBOX_PLAYABLE);
         dataContainer.replace(StructuredDataKey.ATTRIBUTE_MODIFIERS1_21, StructuredDataKey.ATTRIBUTE_MODIFIERS1_20_5, attributeModifiers -> {
             final AttributeModifiers1_20_5.AttributeModifier[] modifiers = Arrays.stream(attributeModifiers.modifiers()).map(modifier -> {
+                final int mappedAttributeId = Protocol1_20_5To1_21.MAPPINGS.getAttributeMappings().inverse().getNewId(modifier.attribute());
+                if (mappedAttributeId == -1) {
+                    return null;
+                }
+
                 final AttributeModifiers1_21.ModifierData modData = modifier.modifier();
                 final String name = AttributeModifierMappings1_21.idToName(modData.id());
                 final AttributeModifiers1_20_5.ModifierData updatedModData = new AttributeModifiers1_20_5.ModifierData(
@@ -176,8 +183,8 @@ public final class BlockItemPacketRewriter1_21 extends StructuredItemRewriter<Cl
                     modData.amount(),
                     modData.operation()
                 );
-                return new AttributeModifiers1_20_5.AttributeModifier(modifier.attribute(), updatedModData, modifier.slotType());
-            }).toArray(AttributeModifiers1_20_5.AttributeModifier[]::new);
+                return new AttributeModifiers1_20_5.AttributeModifier(mappedAttributeId, updatedModData, modifier.slotType());
+            }).filter(Objects::nonNull).toArray(AttributeModifiers1_20_5.AttributeModifier[]::new);
             return new AttributeModifiers1_20_5(modifiers, attributeModifiers.showInTooltip());
         });
     }