From 2c3c62d246c641ce483f00fefbbd14060663139a Mon Sep 17 00:00:00 2001
From: EnZaXD <florian.michael07@gmail.com>
Date: Fri, 1 Nov 2024 18:58:44 +0100
Subject: [PATCH] Update VV API usage (#906)

Co-authored-by: Nassim Jahnke <nassim@njahnke.dev>
---
 .../BackwardsStructuredItemRewriter.java      |  1 +
 .../rewriter/BlockItemPacketRewriter1_21.java | 22 ++++++++++++-------
 2 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/common/src/main/java/com/viaversion/viabackwards/api/rewriters/BackwardsStructuredItemRewriter.java b/common/src/main/java/com/viaversion/viabackwards/api/rewriters/BackwardsStructuredItemRewriter.java
index 9869c7b9..e146c111 100644
--- a/common/src/main/java/com/viaversion/viabackwards/api/rewriters/BackwardsStructuredItemRewriter.java
+++ b/common/src/main/java/com/viaversion/viabackwards/api/rewriters/BackwardsStructuredItemRewriter.java
@@ -114,6 +114,7 @@ public class BackwardsStructuredItemRewriter<C extends ClientboundPacketType, S
         if (customData != null) {
             if (customData.remove(nbtTagName("id")) instanceof final IntTag originalTag) {
                 item.setIdentifier(originalTag.asInt());
+                removeCustomTag(dataContainer, customData);
             }
         }
 
diff --git a/common/src/main/java/com/viaversion/viabackwards/protocol/v1_21to1_20_5/rewriter/BlockItemPacketRewriter1_21.java b/common/src/main/java/com/viaversion/viabackwards/protocol/v1_21to1_20_5/rewriter/BlockItemPacketRewriter1_21.java
index 4be17c25..96bc8554 100644
--- a/common/src/main/java/com/viaversion/viabackwards/protocol/v1_21to1_20_5/rewriter/BlockItemPacketRewriter1_21.java
+++ b/common/src/main/java/com/viaversion/viabackwards/protocol/v1_21to1_20_5/rewriter/BlockItemPacketRewriter1_21.java
@@ -18,6 +18,7 @@
 package com.viaversion.viabackwards.protocol.v1_21to1_20_5.rewriter;
 
 import com.viaversion.nbt.tag.ByteTag;
+import com.viaversion.nbt.tag.CompoundTag;
 import com.viaversion.nbt.tag.StringTag;
 import com.viaversion.nbt.tag.Tag;
 import com.viaversion.viabackwards.api.rewriters.BackwardsStructuredItemRewriter;
@@ -53,7 +54,6 @@ import java.util.ArrayList;
 import java.util.List;
 
 import static com.viaversion.viaversion.protocols.v1_20_5to1_21.rewriter.BlockItemPacketRewriter1_21.downgradeItemData;
-import static com.viaversion.viaversion.protocols.v1_20_5to1_21.rewriter.BlockItemPacketRewriter1_21.resetRarityValues;
 import static com.viaversion.viaversion.protocols.v1_20_5to1_21.rewriter.BlockItemPacketRewriter1_21.updateItemData;
 
 public final class BlockItemPacketRewriter1_21 extends BackwardsStructuredItemRewriter<ClientboundPacket1_21, ServerboundPacket1_20_5, Protocol1_21To1_20_5> {
@@ -190,15 +190,13 @@ public final class BlockItemPacketRewriter1_21 extends BackwardsStructuredItemRe
         super.handleItemToClient(connection, item);
         downgradeItemData(item);
 
-        final StructuredDataContainer dataContainer = item.dataContainer();
-        if (dataContainer.has(StructuredDataKey.RARITY)) {
+        if (data.has(StructuredDataKey.RARITY)) {
             return item;
         }
-
         // Change rarity of trident and piglin banner pattern
         final boolean trident = identifier == 1188;
         if (trident || identifier == 1200) {
-            dataContainer.set(StructuredDataKey.RARITY, trident ? 3 : 1); // Epic or Uncommon
+            data.set(StructuredDataKey.RARITY, trident ? 3 : 1); // Epic or Uncommon
             saveTag(createCustomTag(item), new ByteTag(true), "rarity");
         }
         return item;
@@ -210,8 +208,8 @@ public final class BlockItemPacketRewriter1_21 extends BackwardsStructuredItemRe
             return item;
         }
 
-        final StructuredDataContainer dataContainer = item.dataContainer();
-        dataContainer.setIdLookup(protocol, false);
+        final StructuredDataContainer data = item.dataContainer();
+        data.setIdLookup(protocol, false);
 
         // Rewrite enchantments
         final EnchantmentsPaintingsStorage storage = connection.get(EnchantmentsPaintingsStorage.class);
@@ -224,7 +222,15 @@ public final class BlockItemPacketRewriter1_21 extends BackwardsStructuredItemRe
         // Order is important
         super.handleItemToServer(connection, item);
         updateItemData(item);
-        resetRarityValues(item, nbtTagName("rarity"));
+
+        final CompoundTag customData = data.get(StructuredDataKey.CUSTOM_DATA);
+        if (customData == null) {
+            return item;
+        }
+        if (customData.remove(nbtTagName("rarity")) != null) {
+            data.remove(StructuredDataKey.RARITY);
+            removeCustomTag(data, customData);
+        }
         return item;
     }