From 5e3222c55dff0cc66c309030fa2f7fa95e18a59e Mon Sep 17 00:00:00 2001
From: Alexander Brandes <mc.cache@web.de>
Date: Sat, 20 May 2023 21:00:35 +0200
Subject: [PATCH] Ensure blockmap uniqueness across TE values

---
 .../util/collection/Int2BaseBlockMap.java     |  2 +-
 .../util/collection/BlockMapTest.java         | 20 +++++++++++++++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/worldedit-core/src/main/java/com/sk89q/worldedit/util/collection/Int2BaseBlockMap.java b/worldedit-core/src/main/java/com/sk89q/worldedit/util/collection/Int2BaseBlockMap.java
index bf8e3bed5..fca089a68 100644
--- a/worldedit-core/src/main/java/com/sk89q/worldedit/util/collection/Int2BaseBlockMap.java
+++ b/worldedit-core/src/main/java/com/sk89q/worldedit/util/collection/Int2BaseBlockMap.java
@@ -170,7 +170,7 @@ class Int2BaseBlockMap extends AbstractInt2ObjectMap<BaseBlock> {
             return old;
         }
         int oldId = commonMap.put(key, internalId);
-        return assumeAsBlock(oldId);
+        return BlockStateIdAccess.isValidInternalId(oldId) ? assumeAsBlock(oldId) : uncommonMap.remove(key);
     }
 
     @Override
diff --git a/worldedit-core/src/test/java/com/sk89q/worldedit/util/collection/BlockMapTest.java b/worldedit-core/src/test/java/com/sk89q/worldedit/util/collection/BlockMapTest.java
index 595993919..9558ac0bb 100644
--- a/worldedit-core/src/test/java/com/sk89q/worldedit/util/collection/BlockMapTest.java
+++ b/worldedit-core/src/test/java/com/sk89q/worldedit/util/collection/BlockMapTest.java
@@ -21,6 +21,8 @@
 package com.sk89q.worldedit.util.collection;
 
 import com.google.common.collect.ImmutableMap;
+import com.sk89q.jnbt.CompoundTag;
+import com.sk89q.jnbt.StringTag;
 import com.sk89q.worldedit.WorldEdit;
 import com.sk89q.worldedit.event.platform.PlatformsRegisteredEvent;
 import com.sk89q.worldedit.extension.platform.Capability;
@@ -100,6 +102,7 @@ class BlockMapTest {
 
         registerBlock("minecraft:air");
         registerBlock("minecraft:oak_wood");
+        registerBlock("minecraft:chest");
     }
 
     @AfterAll
@@ -123,6 +126,7 @@ class BlockMapTest {
 
     private final BaseBlock air = checkNotNull(BlockTypes.AIR).getDefaultState().toBaseBlock();
     private final BaseBlock oakWood = checkNotNull(BlockTypes.OAK_WOOD).getDefaultState().toBaseBlock();
+    private final BaseBlock chestWithNbt = checkNotNull(BlockTypes.CHEST).getDefaultState().toBaseBlock(new CompoundTag(ImmutableMap.of("dummy", new StringTag("value"))));
 
     private AutoCloseable mocks;
 
@@ -748,6 +752,22 @@ class BlockMapTest {
             });
         }
 
+        @SuppressWarnings("OverwrittenKey")
+        @Test
+        @DisplayName("put with valid and invalid keys doesn't duplicate")
+        void putWithInvalidAndValid() {
+            generator.makeVectorsStream().forEach(vec -> {
+                BlockMap<BaseBlock> map = BlockMap.createForBaseBlock();
+                // This tests https://github.com/EngineHub/WorldEdit/issues/2250
+                // Due to two internal maps, a bug existed where both could have the same value
+                map.put(vec, chestWithNbt);
+                map.put(vec, air);
+
+                assertEquals(1, map.size());
+                assertEquals(air, map.get(vec));
+            });
+        }
+
     }
 
     @Test