Add registry-related argument types (#10770)

* Add registry-related argument types

* fix tests

patches/api/Brigadier-based-command-API.patch

@@ -3,6 +3,7 @@ From: Owen1212055 <23108066+Owen1212055@users.noreply.github.com>
 Date: Mon, 1 Aug 2022 22:50:29 -0400
 Subject: [PATCH] Brigadier based command API
 
+Co-authored-by: Jake Potrebic <jake.m.potrebic@gmail.com>
 
 diff --git a/build.gradle.kts b/build.gradle.kts
 index 0000000000000000000000000000000000000000..0000000000000000000000000000000000000000 100644
@@ -953,6 +954,8 @@ index 0000000000000000000000000000000000000000..00000000000000000000000000000000
 +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..00000000000000000000000000000000
 +        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..00000000000000000000000000000000
 +        @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..0000000000000000000000000000000000000000
+--- /dev/null
++++ b/src/main/java/io/papermc/paper/command/brigadier/argument/RegistryArgumentExtractor.java
+@@ -0,0 +0,0 @@
++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..0000000000000000000000000000000000000000
@@ -1442,6 +1509,8 @@ index 0000000000000000000000000000000000000000..00000000000000000000000000000000
 +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..00000000000000000000000000000000
 +    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

patches/api/Test-changes.patch

@@ -132,8 +132,8 @@ index 0000000000000000000000000000000000000000..00000000000000000000000000000000
                  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;

patches/server/Add-RegistryAccess-for-managing-Registries.patch

@@ -23,6 +23,7 @@ index 0000000000000000000000000000000000000000..00000000000000000000000000000000
 +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..00000000000000000000000000000000
 +        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() {
 +    }
 +}

patches/server/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 <jake.m.potrebic@gmail.com>
+
 diff --git a/src/main/java/com/mojang/brigadier/CommandDispatcher.java b/src/main/java/com/mojang/brigadier/CommandDispatcher.java
 index 0000000000000000000000000000000000000000..0000000000000000000000000000000000000000 100644
 --- a/src/main/java/com/mojang/brigadier/CommandDispatcher.java
@@ -1096,6 +1098,10 @@ index 0000000000000000000000000000000000000000..00000000000000000000000000000000
 +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..00000000000000000000000000000000
 +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..00000000000000000000000000000000
 +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..00000000000000000000000000000000
 +        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);
 +    }

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 -> {