Add registry-related argument types (#10770)

* Add registry-related argument types

* fix tests

1 files changed, 77 insertions, 4 deletions

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