PaperMC/patches/server/0900-Improve-Registry.patch

From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: Jake Potrebic <jake.m.potrebic@gmail.com>
Date: Wed, 20 Dec 2023 02:03:05 -0800
Subject: [PATCH] Improve Registry


diff --git a/src/main/java/org/bukkit/craftbukkit/CraftRegistry.java b/src/main/java/org/bukkit/craftbukkit/CraftRegistry.java
index 249f0dcad04a35244da6dab837a461bb42aad00a..273844c9071b8d5cf6009c6c94a6c47a9d0cc700 100644
--- a/src/main/java/org/bukkit/craftbukkit/CraftRegistry.java
+++ b/src/main/java/org/bukkit/craftbukkit/CraftRegistry.java
@@ -152,6 +152,7 @@ public class CraftRegistry<B extends Keyed, M> implements Registry<B> {
 
     private final Class<?> bukkitClass; // Paper - relax preload class
     private final Map<NamespacedKey, B> cache = new HashMap<>();
+    private final Map<B, NamespacedKey> byValue = new java.util.IdentityHashMap<>(); // Paper - improve Registry
     private final net.minecraft.core.Registry<M> minecraftRegistry;
     private final BiFunction<NamespacedKey, M, B> minecraftToBukkit;
     private final BiFunction<NamespacedKey, ApiVersion, NamespacedKey> serializationUpdater; // Paper - rename to make it *clear* what it is *only* for
@@ -200,6 +201,7 @@ public class CraftRegistry<B extends Keyed, M> implements Registry<B> {
         }
 
         this.cache.put(namespacedKey, bukkit);
+        this.byValue.put(bukkit, namespacedKey); // Paper - improve Registry
 
         return bukkit;
     }
@@ -232,4 +234,11 @@ public class CraftRegistry<B extends Keyed, M> implements Registry<B> {
 
         return this.minecraftToBukkit.apply(namespacedKey, minecraft);
     }
+
+    // Paper start - improve Registry
+    @Override
+    public NamespacedKey getKey(final B value) {
+        return this.byValue.get(value);
+    }
+    // Paper end - improve Registry
 }
diff --git a/src/main/java/org/bukkit/craftbukkit/inventory/trim/CraftTrimMaterial.java b/src/main/java/org/bukkit/craftbukkit/inventory/trim/CraftTrimMaterial.java
index caf7e4312e95e90dd0822355c8832006e69a2700..38578ef887227ecc8e8bba281eae17a849b4fbe6 100644
--- a/src/main/java/org/bukkit/craftbukkit/inventory/trim/CraftTrimMaterial.java
+++ b/src/main/java/org/bukkit/craftbukkit/inventory/trim/CraftTrimMaterial.java
@@ -54,6 +54,7 @@ public class CraftTrimMaterial implements TrimMaterial, Handleable<net.minecraft
     @Override
     @NotNull
     public NamespacedKey getKey() {
+        if (true) return java.util.Objects.requireNonNull(org.bukkit.Registry.TRIM_MATERIAL.getKey(this), () -> this + " doesn't have a key"); // Paper
         return this.key;
     }
 
diff --git a/src/main/java/org/bukkit/craftbukkit/inventory/trim/CraftTrimPattern.java b/src/main/java/org/bukkit/craftbukkit/inventory/trim/CraftTrimPattern.java
index f91577c9239c8d5ed4b72b23dde9c053b4beae0b..36df80a8be8a2485823f699e45c99674dbe71507 100644
--- a/src/main/java/org/bukkit/craftbukkit/inventory/trim/CraftTrimPattern.java
+++ b/src/main/java/org/bukkit/craftbukkit/inventory/trim/CraftTrimPattern.java
@@ -54,6 +54,7 @@ public class CraftTrimPattern implements TrimPattern, Handleable<net.minecraft.w
     @Override
     @NotNull
     public NamespacedKey getKey() {
+        if (true) return java.util.Objects.requireNonNull(org.bukkit.Registry.TRIM_PATTERN.getKey(this), () -> this + " doesn't have a key"); // Paper
         return this.key;
     }
 
diff --git a/src/test/java/org/bukkit/registry/PerRegistryTest.java b/src/test/java/org/bukkit/registry/PerRegistryTest.java
index 18859eea522ff26cbefb5bbc5065b5369ed6c189..319e000519fd719cea0e6daf2ba9cfa67e6958a3 100644
--- a/src/test/java/org/bukkit/registry/PerRegistryTest.java
+++ b/src/test/java/org/bukkit/registry/PerRegistryTest.java
@@ -49,19 +49,22 @@ public class PerRegistryTest {
 
     @ParameterizedTest
     @MethodSource("data")
-    public void testGet(Registry<?> registry) {
+    public <T extends Keyed> void testGet(Registry<T> registry) { // Paper - improve Registry
         registry.forEach(element -> {
+            NamespacedKey key = registry.getKey(element); // Paper - improve Registry
+            assertNotNull(key); // Paper - improve Registry
             // Values in the registry should be referentially equal to what is returned with #get()
             // This ensures that new instances are not created each time #get() is invoked
-            assertSame(element, registry.get(element.getKey()));
+            assertSame(element, registry.get(key)); // Paper - improve Registry
         });
     }
 
     @ParameterizedTest
     @MethodSource("data")
-    public void testMatch(Registry<?> registry) {
+    public <T extends Keyed> void testMatch(Registry<T> registry) { // Paper - improve Registry
         registry.forEach(element -> {
-            NamespacedKey key = element.getKey();
+            NamespacedKey key = registry.getKey(element); // Paper - improve Registry
+            assertNotNull(key); // Paper - improve Registry
 
             this.assertSameMatchWithKeyMessage(registry, element, key.toString()); // namespace:key
             this.assertSameMatchWithKeyMessage(registry, element, key.getKey()); // key
@@ -72,7 +75,7 @@ public class PerRegistryTest {
         });
     }
 
-    private void assertSameMatchWithKeyMessage(Registry<?> registry, Keyed element, String key) {
+    private <T extends Keyed> void assertSameMatchWithKeyMessage(Registry<T> registry, T element, String key) { // Paper - improve Registry
         assertSame(element, registry.match(key), key);
     }