From 38af6f883b22e4a5e22923d98faedd6cf4a3245f Mon Sep 17 00:00:00 2001
From: Nassim Jahnke <nassim@njahnke.dev>
Date: Thu, 20 Jun 2024 13:13:50 +0200
Subject: [PATCH] Fix 1.21 enchantment remapping Fixes #795

---
 .../StructuredEnchantmentRewriter.java        | 19 ++++++++++++-------
 .../rewriter/BlockItemPacketRewriter1_21.java | 19 ++++++++++++++++---
 2 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/common/src/main/java/com/viaversion/viabackwards/api/rewriters/StructuredEnchantmentRewriter.java b/common/src/main/java/com/viaversion/viabackwards/api/rewriters/StructuredEnchantmentRewriter.java
index 9015c09f..bab84734 100644
--- a/common/src/main/java/com/viaversion/viabackwards/api/rewriters/StructuredEnchantmentRewriter.java
+++ b/common/src/main/java/com/viaversion/viabackwards/api/rewriters/StructuredEnchantmentRewriter.java
@@ -30,7 +30,6 @@ import com.viaversion.viaversion.api.minecraft.data.StructuredDataKey;
 import com.viaversion.viaversion.api.minecraft.item.Item;
 import com.viaversion.viaversion.api.minecraft.item.data.Enchantments;
 import com.viaversion.viaversion.libs.fastutil.ints.Int2IntMap;
-import com.viaversion.viaversion.libs.fastutil.ints.IntIntPair;
 import com.viaversion.viaversion.libs.fastutil.objects.ObjectIterator;
 import com.viaversion.viaversion.rewriter.IdRewriteFunction;
 import com.viaversion.viaversion.util.ComponentUtil;
@@ -86,20 +85,20 @@ public class StructuredEnchantmentRewriter {
         boolean changed = false;
 
         final ObjectIterator<Int2IntMap.Entry> iterator = enchantments.enchantments().int2IntEntrySet().iterator();
-        final List<IntIntPair> updatedIds = new ArrayList<>();
+        final List<PendingIdChange> updatedIds = new ArrayList<>();
         while (iterator.hasNext()) {
             final Int2IntMap.Entry entry = iterator.next();
             final int id = entry.getIntKey();
             final int mappedId = rewriteFunction.rewrite(id);
+            final int level = entry.getIntValue();
             if (mappedId != -1) {
                 if (rewriteIds) {
                     // Update the map after to iteration to preserve the current ids before possibly saving the original, avoid CME
-                    updatedIds.add(IntIntPair.of(id, mappedId));
+                    updatedIds.add(new PendingIdChange(id, mappedId, level));
                 }
                 continue;
             }
 
-            final int level = entry.getIntValue();
             final Tag description = descriptionSupplier.get(id, level);
             if (description != null) {
                 if (!changed) {
@@ -113,9 +112,12 @@ public class StructuredEnchantmentRewriter {
             }
         }
 
-        for (final IntIntPair pair : updatedIds) {
-            final int level = enchantments.enchantments().remove(pair.firstInt());
-            enchantments.add(pair.secondInt(), level);
+        // Remove all first, then add the new ones
+        for (final PendingIdChange change : updatedIds) {
+            enchantments.remove(change.id());
+        }
+        for (final PendingIdChange change : updatedIds) {
+            enchantments.add(change.mappedId(), change.level());
         }
 
         if (loreToAdd.isEmpty()) {
@@ -198,4 +200,7 @@ public class StructuredEnchantmentRewriter {
 
         Tag get(int id, int level);
     }
+
+    private record PendingIdChange(int id, int mappedId, int level) {
+    }
 }
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 bc5aecfb..1a0d0ca3 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
@@ -48,6 +48,7 @@ import com.viaversion.viaversion.rewriter.BlockRewriter;
 import com.viaversion.viaversion.rewriter.IdRewriteFunction;
 import java.util.ArrayList;
 import java.util.Arrays;
+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.updateItemData;
@@ -186,7 +187,8 @@ public final class BlockItemPacketRewriter1_21 extends BackwardsStructuredItemRe
         }
 
         final Enchantments enchantments = enchantmentsData.value();
-        for (final Int2IntMap.Entry entry : new ArrayList<>(enchantments.enchantments().int2IntEntrySet())) {
+        final List<PendingIdChange> updatedIds = new ArrayList<>();
+        for (final Int2IntMap.Entry entry : enchantments.enchantments().int2IntEntrySet()) {
             final int id = entry.getIntKey();
             final String enchantmentKey = Enchantments1_20_5.idToKey(id);
             if (enchantmentKey == null) {
@@ -195,9 +197,20 @@ public final class BlockItemPacketRewriter1_21 extends BackwardsStructuredItemRe
 
             final int mappedId = storage.enchantments().keyToId(enchantmentKey);
             if (id != mappedId) {
-                enchantments.enchantments().remove(id);
-                enchantments.enchantments().put(mappedId, entry.getIntValue());
+                final int level = entry.getIntValue();
+                updatedIds.add(new PendingIdChange(id, mappedId, level));
             }
         }
+
+        // Remove first, then add updated entries
+        for (final PendingIdChange change : updatedIds) {
+            enchantments.remove(change.id);
+        }
+        for (final PendingIdChange change : updatedIds) {
+            enchantments.add(change.mappedId, change.level);
+        }
+    }
+
+    private record PendingIdChange(int id, int mappedId, int level) {
     }
 }
\ No newline at end of file