Add registry-related argument types (#10770)

* Add registry-related argument types

* fix tests

6 files changed, 155 insertions, 13 deletions

diff --git a/patches/api/0003-Test-changes.patch b/patches/api/0003-Test-changes.patch
index 6b6d57cb77..2aa6193edd 100644
--- a/patches/api/0003-Test-changes.patch
+++ b/patches/api/0003-Test-changes.patch
@@ -66,7 +66,7 @@ index 0000000000000000000000000000000000000000..77154095cfb8b259bdb318e8ff40cb6f
 +    }
 +}
 diff --git a/src/test/java/org/bukkit/AnnotationTest.java b/src/test/java/org/bukkit/AnnotationTest.java
-index 64e7aef6220097edefdff3b98a771b988365930d..abadff47166722fdc756afdbc6ac7242b6bd4fb0 100644
+index 64e7aef6220097edefdff3b98a771b988365930d..a899f63eb2ce58b3cf708e91819cbbdeffda5d9f 100644
 --- a/src/test/java/org/bukkit/AnnotationTest.java
 +++ b/src/test/java/org/bukkit/AnnotationTest.java
 @@ -29,7 +29,13 @@ public class AnnotationTest {
@@ -132,8 +132,8 @@ index 64e7aef6220097edefdff3b98a771b988365930d..abadff47166722fdc756afdbc6ac7242
                  for (int i = 0; i < paramTypes.length; i++) {
                      if (mustBeAnnotated(paramTypes[i]) ^ isWellAnnotated(method.invisibleParameterAnnotations == null ? null : method.invisibleParameterAnnotations[i])) {
 +                        // Paper start
-+                        if (method.invisibleTypeAnnotations != null) {
-+                            for (final org.objectweb.asm.tree.TypeAnnotationNode invisibleTypeAnnotation : method.invisibleTypeAnnotations) {
++                        if (method.invisibleTypeAnnotations != null || method.visibleTypeAnnotations != null) {
++                            for (final org.objectweb.asm.tree.TypeAnnotationNode invisibleTypeAnnotation : java.util.Objects.requireNonNullElse(method.invisibleTypeAnnotations, method.visibleTypeAnnotations)) {
 +                                final org.objectweb.asm.TypeReference ref = new org.objectweb.asm.TypeReference(invisibleTypeAnnotation.typeRef);
 +                                if (ref.getSort() == org.objectweb.asm.TypeReference.METHOD_FORMAL_PARAMETER && ref.getTypeParameterIndex() == i && java.util.Arrays.asList(ACCEPTED_ANNOTATIONS).contains(invisibleTypeAnnotation.desc)) {
 +                                    continue dancing;
diff --git a/patches/api/0476-Brigadier-based-command-API.patch b/patches/api/0476-Brigadier-based-command-API.patch
index 1f9ce43f71..b665703a8a 100644
--- a/patches/api/0476-Brigadier-based-command-API.patch
+++ b/patches/api/0476-Brigadier-based-command-API.patch
@@ -3,6 +3,7 @@ From: Owen1212055 <[email protected]>
 Date: Mon, 1 Aug 2022 22:50:29 -0400
 Subject: [PATCH] Brigadier based command API
 
+Co-authored-by: Jake Potrebic <[email protected]>
 
 diff --git a/build.gradle.kts b/build.gradle.kts
 index eecf458e1250ee9968630cf5c3c3287a1693e52e..fd39ed209b20c927054b8482c400beeeeab460a3 100644
@@ -938,10 +939,10 @@ index 0000000000000000000000000000000000000000..2db12952461c92a64505d6646f6f49f8
 +}
 diff --git a/src/main/java/io/papermc/paper/command/brigadier/argument/ArgumentTypes.java b/src/main/java/io/papermc/paper/command/brigadier/argument/ArgumentTypes.java
 new file mode 100644
-index 0000000000000000000000000000000000000000..6c5ffca60a499099fa552020d68060c20abc44b1
+index 0000000000000000000000000000000000000000..1d5c599d1b9c8bf07720e651bdbe9dadb1335b45
 --- /dev/null
 +++ b/src/main/java/io/papermc/paper/command/brigadier/argument/ArgumentTypes.java
-@@ -0,0 +1,324 @@
+@@ -0,0 +1,349 @@
 +package io.papermc.paper.command.brigadier.argument;
 +
 +import com.mojang.brigadier.arguments.ArgumentType;
@@ -953,6 +954,8 @@ index 0000000000000000000000000000000000000000..6c5ffca60a499099fa552020d68060c2
 +import io.papermc.paper.command.brigadier.argument.resolvers.selector.EntitySelectorArgumentResolver;
 +import io.papermc.paper.command.brigadier.argument.resolvers.selector.PlayerSelectorArgumentResolver;
 +import io.papermc.paper.entity.LookAnchor;
++import io.papermc.paper.registry.RegistryKey;
++import io.papermc.paper.registry.TypedKey;
 +import java.util.UUID;
 +import net.kyori.adventure.key.Key;
 +import net.kyori.adventure.text.Component;
@@ -1263,6 +1266,29 @@ index 0000000000000000000000000000000000000000..6c5ffca60a499099fa552020d68060c2
 +        return provider().templateRotation();
 +    }
 +
++    /**
++     * An argument for a resource in a {@link org.bukkit.Registry}.
++     *
++     * @param registryKey the registry's key
++     * @return argument
++     * @param <T> the registry value type
++     */
++    public static <T> @NotNull ArgumentType<T> resource(final @NotNull RegistryKey<T> registryKey) {
++        return provider().resource(registryKey);
++    }
++
++    /**
++     * An argument for a typed key for a {@link org.bukkit.Registry}.
++     *
++     * @param registryKey the registry's key
++     * @return argument
++     * @param <T> the registry value type
++     * @see RegistryArgumentExtractor#getTypedKey(com.mojang.brigadier.context.CommandContext, RegistryKey, String)
++     */
++    public static <T> @NotNull ArgumentType<TypedKey<T>> resourceKey(final @NotNull RegistryKey<T> registryKey) {
++        return provider().resourceKey(registryKey);
++    }
++
 +    private ArgumentTypes() {
 +    }
 +}
@@ -1378,6 +1404,47 @@ index 0000000000000000000000000000000000000000..02acac7f9186677d19c0a62095cc3012
 +        @NotNull T convert(@NotNull N nativeType) throws CommandSyntaxException;
 +    }
 +}
+diff --git a/src/main/java/io/papermc/paper/command/brigadier/argument/RegistryArgumentExtractor.java b/src/main/java/io/papermc/paper/command/brigadier/argument/RegistryArgumentExtractor.java
+new file mode 100644
+index 0000000000000000000000000000000000000000..475266144edf0f7dc4e7939abaf9e1705c4e6461
+--- /dev/null
++++ b/src/main/java/io/papermc/paper/command/brigadier/argument/RegistryArgumentExtractor.java
+@@ -0,0 +1,35 @@
++package io.papermc.paper.command.brigadier.argument;
++
++import com.mojang.brigadier.context.CommandContext;
++import io.papermc.paper.registry.RegistryKey;
++import io.papermc.paper.registry.TypedKey;
++import org.checkerframework.checker.nullness.qual.NonNull;
++
++/**
++ * Utilities for extracting registry-related arguments from a {@link CommandContext}.
++ */
++public final class RegistryArgumentExtractor {
++
++    /**
++     * Gets a typed key argument from a command context.
++     *
++     * @param context the command context
++     * @param registryKey the registry key for the typed key
++     * @param name the argument name
++     * @return the typed key argument
++     * @param <T> the value type
++     * @param <S> the sender type
++     * @throws IllegalArgumentException if the registry key doesn't match the typed key
++     */
++    @SuppressWarnings("unchecked")
++    public static <T, S> @NonNull TypedKey<T> getTypedKey(final @NonNull CommandContext<S> context, final @NonNull RegistryKey<T> registryKey, final @NonNull String name) {
++        final TypedKey<T> typedKey = context.getArgument(name, TypedKey.class);
++        if (typedKey.registryKey().equals(registryKey)) {
++            return typedKey;
++        }
++        throw new IllegalArgumentException(registryKey + " is not the correct registry for " + typedKey);
++    }
++
++    private RegistryArgumentExtractor() {
++    }
++}
 diff --git a/src/main/java/io/papermc/paper/command/brigadier/argument/SignedMessageResolver.java b/src/main/java/io/papermc/paper/command/brigadier/argument/SignedMessageResolver.java
 new file mode 100644
 index 0000000000000000000000000000000000000000..159b691e7a1a7066f3e706e80d75ca8f87a3a964
@@ -1427,10 +1494,10 @@ index 0000000000000000000000000000000000000000..159b691e7a1a7066f3e706e80d75ca8f
 +}
 diff --git a/src/main/java/io/papermc/paper/command/brigadier/argument/VanillaArgumentProvider.java b/src/main/java/io/papermc/paper/command/brigadier/argument/VanillaArgumentProvider.java
 new file mode 100644
-index 0000000000000000000000000000000000000000..da9afa07f919ab139645f06e23b308783d01357a
+index 0000000000000000000000000000000000000000..fbbbf324c002dddd868ba2fb56ddda92149ced3c
 --- /dev/null
 +++ b/src/main/java/io/papermc/paper/command/brigadier/argument/VanillaArgumentProvider.java
-@@ -0,0 +1,98 @@
+@@ -0,0 +1,104 @@
 +package io.papermc.paper.command.brigadier.argument;
 +
 +import com.mojang.brigadier.arguments.ArgumentType;
@@ -1442,6 +1509,8 @@ index 0000000000000000000000000000000000000000..da9afa07f919ab139645f06e23b30878
 +import io.papermc.paper.command.brigadier.argument.resolvers.selector.EntitySelectorArgumentResolver;
 +import io.papermc.paper.command.brigadier.argument.resolvers.selector.PlayerSelectorArgumentResolver;
 +import io.papermc.paper.entity.LookAnchor;
++import io.papermc.paper.registry.RegistryKey;
++import io.papermc.paper.registry.TypedKey;
 +import java.util.Optional;
 +import java.util.ServiceLoader;
 +import java.util.UUID;
@@ -1528,6 +1597,10 @@ index 0000000000000000000000000000000000000000..da9afa07f919ab139645f06e23b30878
 +    ArgumentType<Mirror> templateMirror();
 +
 +    ArgumentType<StructureRotation> templateRotation();
++
++    <T> ArgumentType<TypedKey<T>> resourceKey(RegistryKey<T> registryKey);
++
++    <T> ArgumentType<T> resource(RegistryKey<T> registryKey);
 +}
 diff --git a/src/main/java/io/papermc/paper/command/brigadier/argument/predicate/ItemStackPredicate.java b/src/main/java/io/papermc/paper/command/brigadier/argument/predicate/ItemStackPredicate.java
 new file mode 100644
diff --git a/patches/server/0481-Add-RegistryAccess-for-managing-Registries.patch b/patches/server/0481-Add-RegistryAccess-for-managing-Registries.patch
index cd3f10527d..c06dec8fbe 100644
--- a/patches/server/0481-Add-RegistryAccess-for-managing-Registries.patch
+++ b/patches/server/0481-Add-RegistryAccess-for-managing-Registries.patch
@@ -12,10 +12,10 @@ public net.minecraft.server.RegistryLayer STATIC_ACCESS
 
 diff --git a/src/main/java/io/papermc/paper/registry/PaperRegistries.java b/src/main/java/io/papermc/paper/registry/PaperRegistries.java
 new file mode 100644
-index 0000000000000000000000000000000000000000..4f4595356f2d17c261a84e13d37351f06433177b
+index 0000000000000000000000000000000000000000..5b6d0c5c788bfd158494a88665a2b9b8c45a9ffe
 --- /dev/null
 +++ b/src/main/java/io/papermc/paper/registry/PaperRegistries.java
-@@ -0,0 +1,108 @@
+@@ -0,0 +1,119 @@
 +package io.papermc.paper.registry;
 +
 +import io.papermc.paper.registry.entry.RegistryEntry;
@@ -23,6 +23,7 @@ index 0000000000000000000000000000000000000000..4f4595356f2d17c261a84e13d37351f0
 +import java.util.IdentityHashMap;
 +import java.util.List;
 +import java.util.Map;
++import java.util.Objects;
 +import net.minecraft.core.Registry;
 +import net.minecraft.core.registries.Registries;
 +import net.minecraft.resources.ResourceKey;
@@ -121,6 +122,16 @@ index 0000000000000000000000000000000000000000..4f4595356f2d17c261a84e13d37351f0
 +        return (RegistryEntry<M, T, R>) BY_REGISTRY_KEY.get(registryKey);
 +    }
 +
++    @SuppressWarnings("unchecked")
++    public static <M, T> RegistryKey<T> fromNms(final ResourceKey<? extends Registry<M>> registryResourceKey) {
++        return (RegistryKey<T>) Objects.requireNonNull(BY_RESOURCE_KEY.get(registryResourceKey), registryResourceKey + " doesn't have an api RegistryKey").apiKey();
++    }
++
++    @SuppressWarnings("unchecked")
++    public static <M, T> ResourceKey<? extends Registry<M>> toNms(final RegistryKey<T> registryKey) {
++        return (ResourceKey<? extends Registry<M>>) Objects.requireNonNull(BY_REGISTRY_KEY.get(registryKey), registryKey + " doesn't have an mc registry ResourceKey").mcKey();
++    }
++
 +    private PaperRegistries() {
 +    }
 +}
diff --git a/patches/server/0482-Add-StructuresLocateEvent.patch b/patches/server/0482-Add-StructuresLocateEvent.patch
index 82cdf051db..856b096609 100644
--- a/patches/server/0482-Add-StructuresLocateEvent.patch
+++ b/patches/server/0482-Add-StructuresLocateEvent.patch
@@ -6,10 +6,10 @@ Subject: [PATCH] Add StructuresLocateEvent
 Co-authored-by: Jake Potrebic <[email protected]>
 
 diff --git a/src/main/java/io/papermc/paper/registry/PaperRegistries.java b/src/main/java/io/papermc/paper/registry/PaperRegistries.java
-index 4f4595356f2d17c261a84e13d37351f06433177b..3aa5aec0b17cbdff922009b940d49bb06945c2c6 100644
+index 5b6d0c5c788bfd158494a88665a2b9b8c45a9ffe..51979b3c3f1f3a3c63e0559c70bed9193fd35dbb 100644
 --- a/src/main/java/io/papermc/paper/registry/PaperRegistries.java
 +++ b/src/main/java/io/papermc/paper/registry/PaperRegistries.java
-@@ -45,6 +45,12 @@ import static io.papermc.paper.registry.entry.RegistryEntry.entry;
+@@ -46,6 +46,12 @@ import static io.papermc.paper.registry.entry.RegistryEntry.entry;
  @DefaultQualifier(NonNull.class)
  public final class PaperRegistries {
  
diff --git a/patches/server/1043-Brigadier-based-command-API.patch b/patches/server/1043-Brigadier-based-command-API.patch
index 2dd9a7550d..a318e2e061 100644
--- a/patches/server/1043-Brigadier-based-command-API.patch
+++ b/patches/server/1043-Brigadier-based-command-API.patch
@@ -9,6 +9,8 @@ public net.minecraft.commands.arguments.DimensionArgument ERROR_INVALID_VALUE
 public net.minecraft.server.ReloadableServerResources registryLookup
 public net.minecraft.server.ReloadableServerResources
 
+Co-authored-by: Jake Potrebic <[email protected]>
+
 diff --git a/src/main/java/com/mojang/brigadier/CommandDispatcher.java b/src/main/java/com/mojang/brigadier/CommandDispatcher.java
 index 4b4f812eb13d5f03bcf3f8724d8aa8dbbc724e8b..a4d5d7017e0be79844b996de85a63cad5f8488bc 100644
 --- a/src/main/java/com/mojang/brigadier/CommandDispatcher.java
@@ -1068,10 +1070,10 @@ index 0000000000000000000000000000000000000000..72966584089d3fee9778f572727c9b7f
 +}
 diff --git a/src/main/java/io/papermc/paper/command/brigadier/argument/VanillaArgumentProviderImpl.java b/src/main/java/io/papermc/paper/command/brigadier/argument/VanillaArgumentProviderImpl.java
 new file mode 100644
-index 0000000000000000000000000000000000000000..5f0b46cb2b3f53a81a949fa64690133baa0a304b
+index 0000000000000000000000000000000000000000..93edb22c8500e79f86b101ef38955bca45a8d3a9
 --- /dev/null
 +++ b/src/main/java/io/papermc/paper/command/brigadier/argument/VanillaArgumentProviderImpl.java
-@@ -0,0 +1,321 @@
+@@ -0,0 +1,354 @@
 +package io.papermc.paper.command.brigadier.argument;
 +
 +import com.destroystokyo.paper.profile.CraftPlayerProfile;
@@ -1096,6 +1098,10 @@ index 0000000000000000000000000000000000000000..5f0b46cb2b3f53a81a949fa64690133b
 +import io.papermc.paper.command.brigadier.argument.resolvers.selector.PlayerSelectorArgumentResolver;
 +import io.papermc.paper.entity.LookAnchor;
 +import io.papermc.paper.math.Position;
++import io.papermc.paper.registry.PaperRegistries;
++import io.papermc.paper.registry.RegistryAccess;
++import io.papermc.paper.registry.RegistryKey;
++import io.papermc.paper.registry.TypedKey;
 +import java.util.Collection;
 +import java.util.Collections;
 +import java.util.List;
@@ -1119,6 +1125,8 @@ index 0000000000000000000000000000000000000000..5f0b46cb2b3f53a81a949fa64690133b
 +import net.minecraft.commands.arguments.MessageArgument;
 +import net.minecraft.commands.arguments.ObjectiveCriteriaArgument;
 +import net.minecraft.commands.arguments.RangeArgument;
++import net.minecraft.commands.arguments.ResourceArgument;
++import net.minecraft.commands.arguments.ResourceKeyArgument;
 +import net.minecraft.commands.arguments.ResourceLocationArgument;
 +import net.minecraft.commands.arguments.ScoreboardSlotArgument;
 +import net.minecraft.commands.arguments.StyleArgument;
@@ -1139,6 +1147,7 @@ index 0000000000000000000000000000000000000000..5f0b46cb2b3f53a81a949fa64690133b
 +import net.minecraft.world.level.Level;
 +import org.bukkit.GameMode;
 +import org.bukkit.HeightMap;
++import org.bukkit.Keyed;
 +import org.bukkit.NamespacedKey;
 +import org.bukkit.World;
 +import org.bukkit.block.BlockState;
@@ -1349,6 +1358,32 @@ index 0000000000000000000000000000000000000000..5f0b46cb2b3f53a81a949fa64690133b
 +        return this.wrap(TemplateRotationArgument.templateRotation(), mirror -> StructureRotation.valueOf(mirror.name()));
 +    }
 +
++    @Override
++    public <T> ArgumentType<TypedKey<T>> resourceKey(final RegistryKey<T> registryKey) {
++        return this.wrap(
++            ResourceKeyArgument.key(PaperRegistries.toNms(registryKey)),
++            nmsRegistryKey -> TypedKey.create(registryKey, CraftNamespacedKey.fromMinecraft(nmsRegistryKey.location()))
++        );
++    }
++
++    @Override
++    public <T> ArgumentType<T> resource(final RegistryKey<T> registryKey) {
++        return this.resourceRaw(registryKey);
++    }
++
++    @SuppressWarnings({"unchecked", "rawtypes", "UnnecessaryLocalVariable"})
++    private <T, K extends Keyed> ArgumentType<T> resourceRaw(final RegistryKey registryKeyRaw) { // TODO remove Keyed
++        final RegistryKey<K> registryKey = registryKeyRaw;
++        return (ArgumentType<T>) this.wrap(
++            ResourceArgument.resource(PaperCommands.INSTANCE.getBuildContext(), PaperRegistries.toNms(registryKey)),
++            resource -> requireNonNull(
++                RegistryAccess.registryAccess()
++                    .getRegistry(registryKey)
++                    .get(CraftNamespacedKey.fromMinecraft(resource.key().location()))
++            )
++        );
++    }
++
 +    private <T> ArgumentType<T> wrap(final ArgumentType<T> base) {
 +        return this.wrap(base, identity -> identity);
 +    }
diff --git a/test-plugin/src/main/java/io/papermc/testplugin/brigtests/Registration.java b/test-plugin/src/main/java/io/papermc/testplugin/brigtests/Registration.java
index cd24899f34..67b0ed983d 100644
--- a/test-plugin/src/main/java/io/papermc/testplugin/brigtests/Registration.java
+++ b/test-plugin/src/main/java/io/papermc/testplugin/brigtests/Registration.java
@@ -5,10 +5,13 @@ import io.papermc.paper.command.brigadier.BasicCommand;
 import io.papermc.paper.command.brigadier.CommandSourceStack;
 import io.papermc.paper.command.brigadier.Commands;
 import io.papermc.paper.command.brigadier.argument.ArgumentTypes;
+import io.papermc.paper.command.brigadier.argument.RegistryArgumentExtractor;
 import io.papermc.paper.command.brigadier.argument.range.DoubleRangeProvider;
 import io.papermc.paper.plugin.bootstrap.BootstrapContext;
 import io.papermc.paper.plugin.lifecycle.event.LifecycleEventManager;
 import io.papermc.paper.plugin.lifecycle.event.types.LifecycleEvents;
+import io.papermc.paper.registry.RegistryKey;
+import io.papermc.paper.registry.TypedKey;
 import io.papermc.testplugin.brigtests.example.ExampleAdminCommand;
 import io.papermc.testplugin.brigtests.example.MaterialArgumentType;
 import java.util.Arrays;
@@ -18,6 +21,7 @@ import java.util.List;
 import org.bukkit.Material;
 import org.bukkit.command.CommandSender;
 import org.bukkit.command.defaults.BukkitCommand;
+import org.bukkit.enchantments.Enchantment;
 import org.bukkit.plugin.Plugin;
 import org.bukkit.plugin.java.JavaPlugin;
 import org.jetbrains.annotations.NotNull;
@@ -36,6 +40,25 @@ public final class Registration {
         final LifecycleEventManager<Plugin> lifecycleManager = plugin.getLifecycleManager();
         lifecycleManager.registerEventHandler(LifecycleEvents.COMMANDS, event -> {
             final Commands commands = event.registrar();
+            commands.register(Commands.literal("ench")
+                .then(
+                    Commands.argument("name", ArgumentTypes.resource(RegistryKey.ENCHANTMENT))
+                        .executes(ctx -> {
+                            ctx.getSource().getSender().sendPlainMessage(ctx.getArgument("name", Enchantment.class).toString());
+                            return Command.SINGLE_SUCCESS;
+                        })
+                ).build()
+            );
+            commands.register(Commands.literal("ench-key")
+                .then(
+                    Commands.argument("key", ArgumentTypes.resourceKey(RegistryKey.ENCHANTMENT))
+                        .executes(ctx -> {
+                            final TypedKey<Enchantment> key = RegistryArgumentExtractor.getTypedKey(ctx, RegistryKey.ENCHANTMENT, "key");
+                            ctx.getSource().getSender().sendPlainMessage(key.toString());
+                            return Command.SINGLE_SUCCESS;
+                        })
+                ).build()
+            );
             // ensure plugin commands override
             commands.register(Commands.literal("tag")
                     .executes(ctx -> {