Fixed KeyMap#removeIf fails if there are custom keys registered (#1849)

BG-Software-LLC · Aug 26, 2023 · ec7f59a · ec7f59a
1 parent b53dc6b
commit ec7f59a
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 4 deletions.
diff --git a/src/main/java/com/bgsoftware/superiorskyblock/core/key/collections/EntityTypeKeyMap.java b/src/main/java/com/bgsoftware/superiorskyblock/core/key/collections/EntityTypeKeyMap.java
@@ -172,7 +172,24 @@ public Key getKey(Key original, @Nullable Key def) {
 
     @Override
     public boolean removeIf(Predicate<Key> predicate) {
-        return size() != 0 && entrySet().removeIf(entry -> predicate.test(entry.getKey()));
+        if (isEmpty())
+            return false;
+
+        boolean removed = false;
+
+        {
+            Map<EntityTypeKey, V> innerMap = this.innerMap.getIfPresent().orElse(null);
+            if (innerMap != null)
+                removed |= innerMap.entrySet().removeIf(entry -> predicate.test(entry.getKey()));
+        }
+
+        {
+            Map<CustomKey, V> customInnerMap = this.customInnerMap.getIfPresent().orElse(null);
+            if (customInnerMap != null)
+                removed |= customInnerMap.entrySet().removeIf(entry -> predicate.test(entry.getKey()));
+        }
+
+        return removed;
     }
 
     @Override

diff --git a/src/main/java/com/bgsoftware/superiorskyblock/core/key/collections/MaterialKeyMap.java b/src/main/java/com/bgsoftware/superiorskyblock/core/key/collections/MaterialKeyMap.java
@@ -13,11 +13,9 @@
 import javax.annotation.Nullable;
 import java.util.AbstractMap;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
-import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
 
@@ -193,7 +191,24 @@ public Key getKey(Key original, @Nullable Key def) {
 
     @Override
     public boolean removeIf(Predicate<Key> predicate) {
-        return size() != 0 && entrySet().removeIf(entry -> predicate.test(entry.getKey()));
+        if (isEmpty())
+            return false;
+
+        boolean removed = false;
+
+        {
+            Map<MaterialKey, V> innerMap = this.innerMap.getIfPresent().orElse(null);
+            if (innerMap != null)
+                removed |= innerMap.entrySet().removeIf(entry -> predicate.test(entry.getKey()));
+        }
+
+        {
+            Map<CustomKey, V> customInnerMap = this.customInnerMap.getIfPresent().orElse(null);
+            if (customInnerMap != null)
+                removed |= customInnerMap.entrySet().removeIf(entry -> predicate.test(entry.getKey()));
+        }
+
+        return removed;
     }
 
     @Override