From b4492e9bc5a4f2001f2803abcee11abc4e0b6b3c Mon Sep 17 00:00:00 2001
From: Jason Penilla <11360596+jpenilla@users.noreply.github.com>
Date: Tue, 31 Aug 2021 13:51:04 -0700
Subject: [PATCH] Fix Brigadier suggestions for KeyedWorldArgument

---
 .../bukkit/BukkitBrigadierMapper.java         | 51 ++++++++++++++-----
 .../paper/PaperBrigadierMapper.java           |  2 +-
 .../examples/bukkit/ExamplePlugin.java        | 10 ++++
 3 files changed, 48 insertions(+), 15 deletions(-)

diff --git a/cloud-minecraft/cloud-bukkit/src/main/java/cloud/commandframework/bukkit/BukkitBrigadierMapper.java b/cloud-minecraft/cloud-bukkit/src/main/java/cloud/commandframework/bukkit/BukkitBrigadierMapper.java
index cc969868..5e359383 100644
--- a/cloud-minecraft/cloud-bukkit/src/main/java/cloud/commandframework/bukkit/BukkitBrigadierMapper.java
+++ b/cloud-minecraft/cloud-bukkit/src/main/java/cloud/commandframework/bukkit/BukkitBrigadierMapper.java
@@ -169,6 +169,24 @@ public final class BukkitBrigadierMapper<C> {
     public <T extends ArgumentParser<C, ?>> void mapSimpleNMS(
             final @NonNull TypeToken<T> type,
             final @NonNull String argumentId
+    ) {
+        this.mapSimpleNMS(type, argumentId, false);
+    }
+
+    /**
+     * Attempt to register a mapping between a cloud argument parser type and an NMS brigadier argument type which
+     * has a no-args constructor.
+     *
+     * @param type       Type to map
+     * @param <T>        argument parser type
+     * @param argumentId network id of argument type
+     * @param useCloudSuggestions whether to use cloud suggestions
+     * @since 1.6.0
+     */
+    public <T extends ArgumentParser<C, ?>> void mapSimpleNMS(
+            final @NonNull TypeToken<T> type,
+            final @NonNull String argumentId,
+            final boolean useCloudSuggestions
     ) {
         final Constructor<?> constructor;
         try {
@@ -182,21 +200,26 @@ public final class BukkitBrigadierMapper<C> {
             );
             return;
         }
-        this.brigadierManager.registerMapping(type, builder -> builder.to(argument -> {
-            try {
-                return (ArgumentType<?>) constructor.newInstance();
-            } catch (final ReflectiveOperationException e) {
-                this.commandManager.getOwningPlugin().getLogger().log(
-                        Level.WARNING,
-                        String.format(
-                                "Failed to create instance of brigadier argument type '%s'.",
-                                GenericTypeReflector.erase(type.getType()).getCanonicalName()
-                        ),
-                        e
-                );
-                return fallbackType();
+        this.brigadierManager.registerMapping(type, builder -> {
+            builder.to(argument -> {
+                try {
+                    return (ArgumentType<?>) constructor.newInstance();
+                } catch (final ReflectiveOperationException e) {
+                    this.commandManager.getOwningPlugin().getLogger().log(
+                            Level.WARNING,
+                            String.format(
+                                    "Failed to create instance of brigadier argument type '%s'.",
+                                    GenericTypeReflector.erase(type.getType()).getCanonicalName()
+                            ),
+                            e
+                    );
+                    return fallbackType();
+                }
+            });
+            if (useCloudSuggestions) {
+                builder.cloudSuggestions();
             }
-        }));
+        });
     }
 
     /**
diff --git a/cloud-minecraft/cloud-paper/src/main/java/cloud/commandframework/paper/PaperBrigadierMapper.java b/cloud-minecraft/cloud-paper/src/main/java/cloud/commandframework/paper/PaperBrigadierMapper.java
index 2e45a55e..42604d42 100644
--- a/cloud-minecraft/cloud-paper/src/main/java/cloud/commandframework/paper/PaperBrigadierMapper.java
+++ b/cloud-minecraft/cloud-paper/src/main/java/cloud/commandframework/paper/PaperBrigadierMapper.java
@@ -49,7 +49,7 @@ final class PaperBrigadierMapper<C> {
         final Class<?> keyed = CraftBukkitReflection.findClass("org.bukkit.Keyed");
         if (keyed != null && keyed.isAssignableFrom(World.class)) {
             mapper.mapSimpleNMS(new TypeToken<KeyedWorldArgument.Parser<C>>() {
-            }, "resource_location");
+            }, "resource_location", true);
         }
     }
 
diff --git a/examples/example-bukkit/src/main/java/cloud/commandframework/examples/bukkit/ExamplePlugin.java b/examples/example-bukkit/src/main/java/cloud/commandframework/examples/bukkit/ExamplePlugin.java
index 6658aa26..f8293baf 100644
--- a/examples/example-bukkit/src/main/java/cloud/commandframework/examples/bukkit/ExamplePlugin.java
+++ b/examples/example-bukkit/src/main/java/cloud/commandframework/examples/bukkit/ExamplePlugin.java
@@ -64,6 +64,7 @@ import cloud.commandframework.minecraft.extras.MinecraftHelp;
 import cloud.commandframework.minecraft.extras.RichDescription;
 import cloud.commandframework.minecraft.extras.TextColorArgument;
 import cloud.commandframework.paper.PaperCommandManager;
+import cloud.commandframework.paper.argument.KeyedWorldArgument;
 import cloud.commandframework.permission.PredicatePermission;
 import cloud.commandframework.tasks.TaskConsumer;
 import cloud.commandframework.types.tuples.Pair;
@@ -478,6 +479,15 @@ public final class ExamplePlugin extends JavaPlugin {
                     final ItemStack stack = ctx.get("itemstack");
                     ((Player) ctx.getSender()).getInventory().addItem(stack);
                 }));
+
+        this.manager.command(builder.literal("keyed_world")
+                .argument(KeyedWorldArgument.of("world"))
+                .senderType(Player.class)
+                .handler(ctx -> {
+                    final World world = ctx.get("world");
+                    final Player sender = (Player) ctx.getSender();
+                    this.getServer().getScheduler().runTask(this, () -> sender.teleport(world.getSpawnLocation()));
+                }));
     }
 
     @CommandMethod("example help [query]")