From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: Jake Potrebic <jake.m.potrebic@gmail.com>
Date: Wed, 2 Mar 2022 13:33:08 -0800
Subject: [PATCH] Add PaperRegistry

PaperRegistry is a server-backed impl of bukkit's Registry interface

diff --git a/src/main/java/io/papermc/paper/registry/PaperRegistry.java b/src/main/java/io/papermc/paper/registry/PaperRegistry.java
new file mode 100644
index 0000000000000000000000000000000000000000..0000000000000000000000000000000000000000
--- /dev/null
+++ b/src/main/java/io/papermc/paper/registry/PaperRegistry.java
@@ -0,0 +0,0 @@
+package io.papermc.paper.registry;
+
+import com.google.common.base.Preconditions;
+import com.google.common.base.Suppliers;
+import net.minecraft.core.Holder;
+import net.minecraft.core.Registry;
+import net.minecraft.core.RegistryAccess;
+import net.minecraft.resources.ResourceKey;
+import net.minecraft.resources.ResourceLocation;
+import net.minecraft.server.MinecraftServer;
+import org.bukkit.Keyed;
+import org.bukkit.NamespacedKey;
+import org.bukkit.craftbukkit.util.CraftNamespacedKey;
+import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.checkerframework.framework.qual.DefaultQualifier;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
+
+@DefaultQualifier(NonNull.class)
+public abstract class PaperRegistry<API extends Keyed, MINECRAFT> implements org.bukkit.Registry<API> {
+
+    @SuppressWarnings("FieldMayBeFinal") // non-final for testing
+    private static Supplier<RegistryAccess> REGISTRY_ACCESS = Suppliers.memoize(() -> MinecraftServer.getServer().registryAccess());
+    private static final Map<RegistryKey<?, ?>, PaperRegistry<?, ?>> INTERNAL_REGISTRIES = new HashMap<>();
+    public static final Map<RegistryKey<?, ?>, PaperRegistry<?, ?>> REGISTRIES = Collections.unmodifiableMap(INTERNAL_REGISTRIES);
+    private static final Map<Class<?>, PaperRegistry<?, ?>> REGISTRY_BY_API_CLASS = new HashMap<>();
+    private static final Map<ResourceKey<? extends Registry<?>>, PaperRegistry<?, ?>> REGISTRY_BY_RES_KEY = new HashMap<>();
+
+    private boolean registered;
+    private final RegistryKey<API, MINECRAFT> registryKey;
+    private final Supplier<Registry<MINECRAFT>> registry;
+    private final Map<NamespacedKey, API> cache = new ConcurrentHashMap<>();
+    private final Map<NamespacedKey, ResourceKey<MINECRAFT>> resourceKeyCache = new ConcurrentHashMap<>();
+
+    public PaperRegistry(RegistryKey<API, MINECRAFT> registryKey) {
+        this.registryKey = registryKey;
+        this.registry = Suppliers.memoize(() -> REGISTRY_ACCESS.get().registryOrThrow(this.registryKey.resourceKey()));
+    }
+
+    @Override
+    public @Nullable API get(NamespacedKey key) {
+        return this.cache.computeIfAbsent(key, k -> {
+            final @Nullable MINECRAFT nms = this.registry.get().get(CraftNamespacedKey.toMinecraft(k));
+            if (nms != null) {
+                return this.convertToApi(k, nms);
+            }
+            return null;
+        });
+    }
+
+    public abstract @Nullable API convertToApi(NamespacedKey key, MINECRAFT nms);
+
+    public API convertToApiOrThrow(ResourceLocation resourceLocation, MINECRAFT nms) {
+        return Objects.requireNonNull(this.convertToApi(resourceLocation, nms), resourceLocation + " has a null api representation");
+    }
+
+    public @Nullable API convertToApi(ResourceLocation resourceLocation, MINECRAFT nms) {
+        return this.convertToApi(CraftNamespacedKey.fromMinecraft(resourceLocation), nms);
+    }
+
+    public API convertToApiOrThrow(Holder<MINECRAFT> nmsHolder) {
+        return Objects.requireNonNull(this.convertToApi(nmsHolder), nmsHolder + " has a null api representation");
+    }
+
+    public @Nullable API convertToApi(Holder<MINECRAFT> nmsHolder) {
+        final Optional<ResourceKey<MINECRAFT>> key = nmsHolder.unwrapKey();
+        if (nmsHolder.isBound() && key.isPresent()) {
+            return this.convertToApi(key.get().location(), nmsHolder.value());
+        } else if (!nmsHolder.isBound() && key.isPresent()) {
+            return this.convertToApi(key.get().location(), this.registry.get().getOrThrow(key.get()));
+        } else if (nmsHolder.isBound() && key.isEmpty()) {
+            final @Nullable ResourceLocation loc = this.registry.get().getKey(nmsHolder.value());
+            if (loc != null) {
+                return this.convertToApi(loc, nmsHolder.value());
+            }
+        }
+        throw new IllegalStateException("Cannot convert " + nmsHolder + " to an API type in: " + this.registryKey);
+    }
+
+    public void convertToApi(Iterable<Holder<MINECRAFT>> holders, Consumer<API> apiConsumer, boolean throwOnNull) {
+        for (Holder<MINECRAFT> holder : holders) {
+            final @Nullable API api = this.convertToApi(holder);
+            if (api == null && throwOnNull) {
+                throw new NullPointerException(holder + " has a null api representation");
+            } else if (api != null) {
+                apiConsumer.accept(api);
+            }
+        }
+    }
+
+    public MINECRAFT getMinecraftValue(API apiValue) {
+        return this.registry.get().getOptional(CraftNamespacedKey.toMinecraft(apiValue.getKey())).orElseThrow();
+    }
+
+    public Holder<MINECRAFT> getMinecraftHolder(API apiValue) {
+        return this.registry.get().getHolderOrThrow(this.resourceKeyCache.computeIfAbsent(apiValue.getKey(), key -> ResourceKey.create(this.registryKey.resourceKey(), CraftNamespacedKey.toMinecraft(key))));
+    }
+
+    @Override
+    public Iterator<API> iterator() {
+        return this.registry.get().keySet().stream().map(key -> this.get(CraftNamespacedKey.fromMinecraft(key))).iterator();
+    }
+
+    public void clearCache() {
+        this.cache.clear();
+    }
+
+    public void register() {
+        if (this.registered) {
+            throw new IllegalStateException("Already registered: " + this.registryKey.apiClass());
+        }
+        INTERNAL_REGISTRIES.put(this.registryKey, this);
+        REGISTRY_BY_API_CLASS.put(this.registryKey.apiClass(), this);
+        REGISTRY_BY_RES_KEY.put(this.registryKey.resourceKey(), this);
+        this.registered = true;
+    }
+
+    @Override
+    public boolean equals(@Nullable Object o) {
+        if (this == o) return true;
+        if (o == null || !PaperRegistry.class.isAssignableFrom(o.getClass())) return false;
+        PaperRegistry<?, ?> that = (PaperRegistry<?, ?>) o;
+        return this.registryKey.equals(that.registryKey);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(this.registryKey);
+    }
+
+    protected static <T> Supplier<Registry<T>> registryFor(ResourceKey<? extends Registry<T>> registryKey) {
+        return Suppliers.memoize(() -> REGISTRY_ACCESS.get().registryOrThrow(registryKey));
+    }
+
+    public static void clearCaches() {
+        for (PaperRegistry<?, ?> registry : INTERNAL_REGISTRIES.values()) {
+            registry.clearCache();
+        }
+    }
+
+    @SuppressWarnings("unchecked")
+    public static <T extends Keyed> PaperRegistry<T, ?> getRegistry(Class<T> classOfT) {
+        Preconditions.checkArgument(REGISTRY_BY_API_CLASS.containsKey(classOfT), "No registry for that type");
+        return (PaperRegistry<T, ?>) REGISTRY_BY_API_CLASS.get(classOfT);
+    }
+
+    @SuppressWarnings("unchecked")
+    public static <T> PaperRegistry<?, T> getRegistry(ResourceKey<? extends Registry<T>> resourceKey) {
+        Preconditions.checkArgument(REGISTRY_BY_RES_KEY.containsKey(resourceKey));
+        return (PaperRegistry<?, T>) REGISTRY_BY_RES_KEY.get(resourceKey);
+    }
+
+    @SuppressWarnings("unchecked")
+    public static <A extends Keyed, M> PaperRegistry<A, M> getRegistry(RegistryKey<A, M> registryKey) {
+        Preconditions.checkArgument(INTERNAL_REGISTRIES.containsKey(registryKey));
+        return (PaperRegistry<A, M>) INTERNAL_REGISTRIES.get(registryKey);
+    }
+}
diff --git a/src/main/java/io/papermc/paper/registry/RegistryKey.java b/src/main/java/io/papermc/paper/registry/RegistryKey.java
new file mode 100644
index 0000000000000000000000000000000000000000..0000000000000000000000000000000000000000
--- /dev/null
+++ b/src/main/java/io/papermc/paper/registry/RegistryKey.java
@@ -0,0 +0,0 @@
+package io.papermc.paper.registry;
+
+import net.minecraft.core.Registry;
+import net.minecraft.resources.ResourceKey;
+import org.bukkit.Keyed;
+
+public record RegistryKey<API extends Keyed, MINECRAFT>(Class<API> apiClass, ResourceKey<? extends Registry<MINECRAFT>> resourceKey) {
+}
diff --git a/src/main/java/net/minecraft/server/MinecraftServer.java b/src/main/java/net/minecraft/server/MinecraftServer.java
index 0000000000000000000000000000000000000000..0000000000000000000000000000000000000000 100644
--- a/src/main/java/net/minecraft/server/MinecraftServer.java
+++ b/src/main/java/net/minecraft/server/MinecraftServer.java
@@ -0,0 +0,0 @@ public abstract class MinecraftServer extends ReentrantBlockableEventLoop<TickTa
 
             this.worldData.setDataConfiguration(worlddataconfiguration);
             this.resources.managers.updateRegistryTags(this.registryAccess());
+            io.papermc.paper.registry.PaperRegistry.clearCaches(); // Paper
             new io.papermc.paper.event.server.ServerResourcesReloadedEvent(cause).callEvent(); // Paper
             // Paper start
             if (Thread.currentThread() != this.serverThread) {
diff --git a/src/main/java/org/bukkit/craftbukkit/util/CraftMagicNumbers.java b/src/main/java/org/bukkit/craftbukkit/util/CraftMagicNumbers.java
index 0000000000000000000000000000000000000000..0000000000000000000000000000000000000000 100644
--- a/src/main/java/org/bukkit/craftbukkit/util/CraftMagicNumbers.java
+++ b/src/main/java/org/bukkit/craftbukkit/util/CraftMagicNumbers.java
@@ -0,0 +0,0 @@ public final class CraftMagicNumbers implements UnsafeValues {
     public int nextEntityId() {
         return net.minecraft.world.entity.Entity.nextEntityId();
     }
+
+    @Override
+    public <T extends org.bukkit.Keyed> Registry<T> registryFor(Class<T> classOfT) {
+        return io.papermc.paper.registry.PaperRegistry.getRegistry(classOfT);
+    }
     // Paper end
 
     /**
diff --git a/src/test/java/org/bukkit/support/AbstractTestingBase.java b/src/test/java/org/bukkit/support/AbstractTestingBase.java
index 0000000000000000000000000000000000000000..0000000000000000000000000000000000000000 100644
--- a/src/test/java/org/bukkit/support/AbstractTestingBase.java
+++ b/src/test/java/org/bukkit/support/AbstractTestingBase.java
@@ -0,0 +0,0 @@ public abstract class AbstractTestingBase {
         LayeredRegistryAccess<RegistryLayer> layers = RegistryLayer.createRegistryAccess();
         layers = WorldLoader.loadAndReplaceLayer(resourceManager, layers, RegistryLayer.WORLDGEN, RegistryDataLoader.WORLDGEN_REGISTRIES);
         REGISTRY_CUSTOM = layers.compositeAccess().freeze();
+        // Paper start
+        try {
+            java.lang.reflect.Field field = io.papermc.paper.registry.PaperRegistry.class.getDeclaredField("REGISTRY_ACCESS");
+            field.trySetAccessible();
+            field.set(null, com.google.common.base.Suppliers.ofInstance(REGISTRY_CUSTOM));
+        } catch (ReflectiveOperationException ex) {
+            throw new IllegalStateException("Could not reflectively set RegistryAccess in PaperRegistry", ex);
+        }
+        // Paper end
         // Register vanilla pack
         DATA_PACK = ReloadableServerResources.loadResources(resourceManager, REGISTRY_CUSTOM, FeatureFlags.REGISTRY.allFlags(), Commands.CommandSelection.DEDICATED, 0, MoreExecutors.directExecutor(), MoreExecutors.directExecutor()).join();
         // Bind tags