Properly remove custom_data component if we created it in 1.20.5+ protocols (ViaVersion#4229)

Author: FlorianMichael

common/src/main/java/com/viaversion/viaversion/protocols/v1_20_5to1_21/rewriter/BlockItemPacketRewriter1_21.java

@@ -136,14 +136,13 @@ public Item handleItemToClient(final UserConnection connection, final Item item)
         super.handleItemToClient(connection, item);
         updateItemData(item);
 
-        final StructuredDataContainer dataContainer = item.dataContainer();
-        if (dataContainer.has(StructuredDataKey.RARITY)) {
+        final StructuredDataContainer data = item.dataContainer();
+        if (data.has(StructuredDataKey.RARITY)) {
             return item;
         }
-
         // Change rarity of trident and piglin banner pattern
         if (item.identifier() == 1188 || item.identifier() == 1200) {
-            dataContainer.set(StructuredDataKey.RARITY, 0); // Common
+            data.set(StructuredDataKey.RARITY, 0); // Common
             saveTag(createCustomTag(item), new ByteTag(true), "rarity");
         }
         return item;
@@ -174,7 +173,16 @@ public Item handleItemToServer(final UserConnection connection, final Item item)
 
         super.handleItemToServer(connection, item);
         downgradeItemData(item);
-        resetRarityValues(item, nbtTagName("rarity"));
+
+        final StructuredDataContainer data = item.dataContainer();
+        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;
     }
 
@@ -206,21 +214,6 @@ public static void downgradeItemData(final Item item) {
         });
     }
 
-    public static void resetRarityValues(final Item item, final String tagName) {
-        final StructuredDataContainer dataContainer = item.dataContainer();
-        final CompoundTag customData = dataContainer.get(StructuredDataKey.CUSTOM_DATA);
-        if (customData == null) {
-            return;
-        }
-
-        if (customData.remove(tagName) != null) {
-            dataContainer.remove(StructuredDataKey.RARITY);
-            if (customData.isEmpty()) {
-                dataContainer.remove(StructuredDataKey.CUSTOM_DATA);
-            }
-        }
-    }
-
     private int itemToJubeboxSong(final int id) {
         String identifier = Protocol1_20_5To1_21.MAPPINGS.getFullItemMappings().identifier(id);
         if (!identifier.contains("music_disc_")) {

common/src/main/java/com/viaversion/viaversion/protocols/v1_21to1_21_2/rewriter/BlockItemPacketRewriter1_21_2.java

@@ -484,9 +484,7 @@ public Item handleItemToServer(final UserConnection connection, final Item item)
 
         if (customData.remove(nbtTagName("remove_custom_name")) != null) {
             dataContainer.remove(StructuredDataKey.CUSTOM_NAME);
-            if (customData.isEmpty()) {
-                dataContainer.remove(StructuredDataKey.CUSTOM_DATA);
-            }
+            removeCustomTag(dataContainer, customData);
         }
 
         final IntArrayTag emptyEnchantments = customData.getIntArrayTag(nbtTagName("0_enchants"));
@@ -504,9 +502,7 @@ public Item handleItemToServer(final UserConnection connection, final Item item)
             if (customData.remove(nbtTagName("remove_glint")) != null) {
                 dataContainer.remove(StructuredDataKey.ENCHANTMENT_GLINT_OVERRIDE);
             }
-            if (customData.isEmpty()) {
-                dataContainer.remove(StructuredDataKey.CUSTOM_DATA);
-            }
+            removeCustomTag(dataContainer, customData);
         }
         return item;
     }

common/src/main/java/com/viaversion/viaversion/rewriter/StructuredItemRewriter.java

@@ -41,6 +41,8 @@
 public class StructuredItemRewriter<C extends ClientboundPacketType, S extends ServerboundPacketType,
     T extends Protocol<C, ?, ?, S>> extends ItemRewriter<C, S, T> {
 
+    public static final String MARKER_KEY = "VV|custom_data";
+
     public StructuredItemRewriter(
         T protocol,
         Type<Item> itemType, Type<Item[]> itemArrayType, Type<Item> mappedItemType, Type<Item[]> mappedItemArrayType,
@@ -205,15 +207,18 @@ protected void restoreTextComponents(final Item item) {
         // Remove custom name
         if (customData.remove(nbtTagName("added_custom_name")) != null) {
             data.remove(StructuredDataKey.CUSTOM_NAME);
+            removeCustomTag(data, customData);
         } else {
             final Tag customName = removeBackupTag(customData, "custom_name");
             if (customName != null) {
                 data.set(StructuredDataKey.CUSTOM_NAME, customName);
+                removeCustomTag(data, customData);
             }
 
             final Tag itemName = removeBackupTag(customData, "item_name");
             if (itemName != null) {
                 data.set(StructuredDataKey.ITEM_NAME, itemName);
+                removeCustomTag(data, customData);
             }
         }
     }
@@ -223,6 +228,7 @@ protected CompoundTag createCustomTag(final Item item) {
         CompoundTag customData = data.get(StructuredDataKey.CUSTOM_DATA);
         if (customData == null) {
             customData = new CompoundTag();
+            customData.putBoolean(MARKER_KEY, true);
             data.set(StructuredDataKey.CUSTOM_DATA, customData);
         }
         return customData;
@@ -239,6 +245,13 @@ protected void saveTag(final CompoundTag customData, final Tag tag, final String
         return customData.remove(nbtTagName(tagName));
     }
 
+    protected void removeCustomTag(final StructuredDataContainer data, final CompoundTag customData) {
+        // Only remove if we initially added it and only the marker is left
+        if (customData.contains(MARKER_KEY) && customData.size() == 1) {
+            data.remove(StructuredDataKey.CUSTOM_DATA);
+        }
+    }
+
     @FunctionalInterface
     private interface ItemHandler {