Implement command suggestions for Forge and Fabric (#25)

6 files changed, 130 insertions, 30 deletions

diff --git a/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricClientSparkPlugin.java b/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricClientSparkPlugin.java
index d87d5be..a23dc48 100644
--- a/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricClientSparkPlugin.java
+++ b/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricClientSparkPlugin.java
@@ -20,7 +20,13 @@
 
 package me.lucko.spark.fabric.plugin;
 
+import com.mojang.brigadier.Command;
 import com.mojang.brigadier.CommandDispatcher;
+import com.mojang.brigadier.context.CommandContext;
+import com.mojang.brigadier.exceptions.CommandSyntaxException;
+import com.mojang.brigadier.suggestion.SuggestionProvider;
+import com.mojang.brigadier.suggestion.Suggestions;
+import com.mojang.brigadier.suggestion.SuggestionsBuilder;
 import me.lucko.spark.common.sampler.TickCounter;
 import me.lucko.spark.fabric.FabricCommandSender;
 import me.lucko.spark.fabric.FabricSparkGameHooks;
@@ -33,10 +39,11 @@ import net.minecraft.server.command.CommandOutput;
 import net.minecraft.server.command.CommandSource;
 
 import java.util.Arrays;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Stream;
 
-public class FabricClientSparkPlugin extends FabricSparkPlugin {
+public class FabricClientSparkPlugin extends FabricSparkPlugin implements SuggestionProvider<CommandSource> {
 
     public static void register(FabricSparkMod mod, MinecraftClient client) {
         FabricClientSparkPlugin plugin = new FabricClientSparkPlugin(mod, client);
@@ -67,7 +74,7 @@ public class FabricClientSparkPlugin extends FabricSparkPlugin {
             CommandDispatcher<CommandSource> dispatcher = connection.getCommandDispatcher();
             if (dispatcher != this.dispatcher) {
                 this.dispatcher = dispatcher;
-                registerCommands(this.dispatcher, c -> 0, "sparkc", "sparkclient");
+                registerCommands(this.dispatcher, c -> Command.SINGLE_SUCCESS, this, "sparkc", "sparkclient");
                 FabricSparkGameHooks.INSTANCE.setChatSendCallback(this::onClientChat);
             }
         } catch (Exception e) {
@@ -88,6 +95,24 @@ public class FabricClientSparkPlugin extends FabricSparkPlugin {
     }
 
     @Override
+    public CompletableFuture<Suggestions> getSuggestions(CommandContext<CommandSource> context, SuggestionsBuilder builder)
+            throws CommandSyntaxException {
+        String[] split = context.getInput().split(" ");
+        if (split.length == 0 || (!split[0].equals("/sparkc") && !split[0].equals("/sparkclient"))) {
+            return Suggestions.empty();
+        }
+
+        String[] args = Arrays.copyOfRange(split, 1, split.length);
+
+        return CompletableFuture.supplyAsync(() -> {
+            for (String each : this.platform.tabCompleteCommand(new FabricCommandSender(this.minecraft.player, this), args)) {
+                builder.suggest(each);
+            }
+            return builder.build();
+        });
+    }
+
+    @Override
     public boolean hasPermission(CommandOutput sender, String permission) {
         return true;
     }
diff --git a/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricServerSparkPlugin.java b/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricServerSparkPlugin.java
index 9729e55..50d6560 100644
--- a/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricServerSparkPlugin.java
+++ b/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricServerSparkPlugin.java
@@ -24,6 +24,9 @@ import com.mojang.brigadier.Command;
 import com.mojang.brigadier.CommandDispatcher;
 import com.mojang.brigadier.context.CommandContext;
 import com.mojang.brigadier.exceptions.CommandSyntaxException;
+import com.mojang.brigadier.suggestion.SuggestionProvider;
+import com.mojang.brigadier.suggestion.Suggestions;
+import com.mojang.brigadier.suggestion.SuggestionsBuilder;
 import me.lucko.spark.common.sampler.TickCounter;
 import me.lucko.spark.fabric.FabricCommandSender;
 import me.lucko.spark.fabric.FabricSparkMod;
@@ -32,17 +35,19 @@ import net.minecraft.entity.player.PlayerEntity;
 import net.minecraft.server.MinecraftServer;
 import net.minecraft.server.command.CommandOutput;
 import net.minecraft.server.command.ServerCommandSource;
+import net.minecraft.server.network.ServerPlayerEntity;
 
 import java.util.Arrays;
+import java.util.concurrent.CompletableFuture;
 import java.util.stream.Stream;
 
-public class FabricServerSparkPlugin extends FabricSparkPlugin implements Command<ServerCommandSource> {
+public class FabricServerSparkPlugin extends FabricSparkPlugin implements Command<ServerCommandSource>, SuggestionProvider<ServerCommandSource> {
 
     public static void register(FabricSparkMod mod, MinecraftServer server) {
         CommandDispatcher<ServerCommandSource> dispatcher = server.getCommandManager().getDispatcher();
 
         FabricServerSparkPlugin plugin = new FabricServerSparkPlugin(mod, server);
-        registerCommands(dispatcher, plugin, "spark");
+        registerCommands(dispatcher, plugin, plugin, "spark");
     }
 
     private final MinecraftServer server;
@@ -52,17 +57,39 @@ public class FabricServerSparkPlugin extends FabricSparkPlugin implements Comman
         this.server = server;
     }
 
-    @Override
-    public int run(CommandContext<ServerCommandSource> context) throws CommandSyntaxException {
+    private String /*Nullable*/ [] processArgs(CommandContext<ServerCommandSource> context) {
         String[] split = context.getInput().split(" ");
         if (split.length == 0 || !split[0].equals("/spark")) {
-            return 0;
+            return null;
         }
 
-        String[] args = Arrays.copyOfRange(split, 1, split.length);
+        return Arrays.copyOfRange(split, 1, split.length);
+    }
+
+    @Override
+    public int run(CommandContext<ServerCommandSource> context) throws CommandSyntaxException {
+        String[] args = processArgs(context);
+        if (args == null)
+            return 0;
 
         this.platform.executeCommand(new FabricCommandSender(context.getSource().getPlayer(), this), args);
-        return 1;
+        return Command.SINGLE_SUCCESS;
+    }
+
+    @Override
+    public CompletableFuture<Suggestions> getSuggestions(CommandContext<ServerCommandSource> context, SuggestionsBuilder builder)
+            throws CommandSyntaxException {
+        String[] args = processArgs(context);
+        if (args == null)
+            return Suggestions.empty();
+        ServerPlayerEntity player = context.getSource().getPlayer();
+
+        return CompletableFuture.supplyAsync(() -> {
+            for (String each : this.platform.tabCompleteCommand(new FabricCommandSender(player, this), args)) {
+                builder.suggest(each);
+            }
+            return builder.build();
+        });
     }
 
     @Override
diff --git a/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricSparkPlugin.java b/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricSparkPlugin.java
index 6672358..1f6f098 100644
--- a/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricSparkPlugin.java
+++ b/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricSparkPlugin.java
@@ -26,6 +26,7 @@ import com.mojang.brigadier.CommandDispatcher;
 import com.mojang.brigadier.arguments.StringArgumentType;
 import com.mojang.brigadier.builder.LiteralArgumentBuilder;
 import com.mojang.brigadier.builder.RequiredArgumentBuilder;
+import com.mojang.brigadier.suggestion.SuggestionProvider;
 import com.mojang.brigadier.tree.LiteralCommandNode;
 import me.lucko.spark.common.SparkPlatform;
 import me.lucko.spark.common.SparkPlugin;
@@ -39,7 +40,7 @@ import java.util.concurrent.ScheduledExecutorService;
 
 public abstract class FabricSparkPlugin implements SparkPlugin {
 
-    public static <T> void registerCommands(CommandDispatcher<T> dispatcher, Command<T> executor, String... aliases) {
+    public static <T> void registerCommands(CommandDispatcher<T> dispatcher, Command<T> executor, SuggestionProvider<T> suggestor, String... aliases) {
         if (aliases.length == 0) {
             return;
         }
@@ -48,6 +49,7 @@ public abstract class FabricSparkPlugin implements SparkPlugin {
         LiteralArgumentBuilder<T> command = LiteralArgumentBuilder.<T>literal(mainName)
                 .executes(executor)
                 .then(RequiredArgumentBuilder.<T, String>argument("args", StringArgumentType.greedyString())
+                        .suggests(suggestor)
                         .executes(executor)
                 );
 
diff --git a/spark-forge/src/main/java/me/lucko/spark/forge/plugin/ForgeClientSparkPlugin.java b/spark-forge/src/main/java/me/lucko/spark/forge/plugin/ForgeClientSparkPlugin.java
index 7de0038..ee32325 100644
--- a/spark-forge/src/main/java/me/lucko/spark/forge/plugin/ForgeClientSparkPlugin.java
+++ b/spark-forge/src/main/java/me/lucko/spark/forge/plugin/ForgeClientSparkPlugin.java
@@ -20,7 +20,13 @@
 
 package me.lucko.spark.forge.plugin;
 
+import com.mojang.brigadier.Command;
 import com.mojang.brigadier.CommandDispatcher;
+import com.mojang.brigadier.context.CommandContext;
+import com.mojang.brigadier.exceptions.CommandSyntaxException;
+import com.mojang.brigadier.suggestion.SuggestionProvider;
+import com.mojang.brigadier.suggestion.Suggestions;
+import com.mojang.brigadier.suggestion.SuggestionsBuilder;
 import me.lucko.spark.common.sampler.TickCounter;
 import me.lucko.spark.forge.ForgeCommandSender;
 import me.lucko.spark.forge.ForgeSparkMod;
@@ -36,20 +42,12 @@ import net.minecraftforge.event.TickEvent;
 import net.minecraftforge.eventbus.api.SubscribeEvent;
 import net.minecraftforge.fml.event.lifecycle.FMLClientSetupEvent;
 
-import java.lang.reflect.Field;
 import java.util.Arrays;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Stream;
 
-public class ForgeClientSparkPlugin extends ForgeSparkPlugin {
-
-    private static final Field COMMAND_DISPATCHER_FIELD;
-    static {
-        COMMAND_DISPATCHER_FIELD = Arrays.stream(ClientPlayNetHandler.class.getDeclaredFields())
-                .filter(f -> f.getType() == CommandDispatcher.class)
-                .findFirst().orElseThrow(() -> new RuntimeException("No field with CommandDispatcher type"));
-        COMMAND_DISPATCHER_FIELD.setAccessible(true);
-    }
+public class ForgeClientSparkPlugin extends ForgeSparkPlugin implements SuggestionProvider<ISuggestionProvider> {
 
     public static void register(ForgeSparkMod mod, FMLClientSetupEvent event) {
         Minecraft minecraft = event.getMinecraftSupplier().get();
@@ -80,10 +78,10 @@ public class ForgeClientSparkPlugin extends ForgeSparkPlugin {
         }
 
         try {
-            CommandDispatcher<ISuggestionProvider> dispatcher = (CommandDispatcher) COMMAND_DISPATCHER_FIELD.get(connection);
+            CommandDispatcher<ISuggestionProvider> dispatcher = connection.func_195515_i();
             if (dispatcher != this.dispatcher) {
                 this.dispatcher = dispatcher;
-                registerCommands(this.dispatcher, context -> 1, "sparkc", "sparkclient");
+                registerCommands(this.dispatcher, context -> Command.SINGLE_SUCCESS, this, "sparkc", "sparkclient");
             }
         } catch (Exception e) {
             e.printStackTrace();
@@ -105,6 +103,25 @@ public class ForgeClientSparkPlugin extends ForgeSparkPlugin {
     }
 
     @Override
+    public CompletableFuture<Suggestions> getSuggestions(CommandContext<ISuggestionProvider> context, SuggestionsBuilder builder)
+            throws CommandSyntaxException {
+        String chat = context.getInput();
+        String[] split = chat.split(" ");
+        if (split.length == 0 || (!split[0].equals("/sparkc") && !split[0].equals("/sparkclient"))) {
+            return Suggestions.empty();
+        }
+
+        String[] args = Arrays.copyOfRange(split, 1, split.length);
+
+        return CompletableFuture.supplyAsync(() -> {
+            for (String each : this.platform.tabCompleteCommand(new ForgeCommandSender(this.minecraft.player, this), args)) {
+                builder.suggest(each);
+            }
+            return builder.build();
+        });
+    }
+
+    @Override
     public boolean hasPermission(ICommandSource sender, String permission) {
         return true;
     }
diff --git a/spark-forge/src/main/java/me/lucko/spark/forge/plugin/ForgeServerSparkPlugin.java b/spark-forge/src/main/java/me/lucko/spark/forge/plugin/ForgeServerSparkPlugin.java
index 5eb61f8..7a91e5c 100644
--- a/spark-forge/src/main/java/me/lucko/spark/forge/plugin/ForgeServerSparkPlugin.java
+++ b/spark-forge/src/main/java/me/lucko/spark/forge/plugin/ForgeServerSparkPlugin.java
@@ -24,6 +24,9 @@ import com.mojang.brigadier.Command;
 import com.mojang.brigadier.CommandDispatcher;
 import com.mojang.brigadier.context.CommandContext;
 import com.mojang.brigadier.exceptions.CommandSyntaxException;
+import com.mojang.brigadier.suggestion.SuggestionProvider;
+import com.mojang.brigadier.suggestion.Suggestions;
+import com.mojang.brigadier.suggestion.SuggestionsBuilder;
 import me.lucko.spark.common.sampler.TickCounter;
 import me.lucko.spark.forge.ForgeCommandSender;
 import me.lucko.spark.forge.ForgeSparkMod;
@@ -31,6 +34,7 @@ import me.lucko.spark.forge.ForgeTickCounter;
 import net.minecraft.command.CommandSource;
 import net.minecraft.command.ICommandSource;
 import net.minecraft.entity.player.PlayerEntity;
+import net.minecraft.entity.player.ServerPlayerEntity;
 import net.minecraft.server.MinecraftServer;
 import net.minecraftforge.event.TickEvent;
 import net.minecraftforge.fml.event.server.FMLServerStartingEvent;
@@ -38,16 +42,17 @@ import net.minecraftforge.server.permission.DefaultPermissionLevel;
 import net.minecraftforge.server.permission.PermissionAPI;
 
 import java.util.Arrays;
+import java.util.concurrent.CompletableFuture;
 import java.util.stream.Stream;
 
-public class ForgeServerSparkPlugin extends ForgeSparkPlugin implements Command<CommandSource> {
+public class ForgeServerSparkPlugin extends ForgeSparkPlugin implements Command<CommandSource>, SuggestionProvider<CommandSource> {
 
     public static void register(ForgeSparkMod mod, FMLServerStartingEvent event) {
         MinecraftServer server = event.getServer();
         CommandDispatcher<CommandSource> dispatcher = event.getCommandDispatcher();
 
         ForgeServerSparkPlugin plugin = new ForgeServerSparkPlugin(mod, server);
-        registerCommands(dispatcher, plugin, "spark");
+        registerCommands(dispatcher, plugin, plugin, "spark");
         PermissionAPI.registerNode("spark", DefaultPermissionLevel.OP, "Access to the spark command");
     }
 
@@ -58,17 +63,39 @@ public class ForgeServerSparkPlugin extends ForgeSparkPlugin implements Command<
         this.server = server;
     }
 
-    @Override
-    public int run(CommandContext<CommandSource> context) throws CommandSyntaxException {
+    private String /* Nullable */[] processArgs(CommandContext<CommandSource> context) {
         String[] split = context.getInput().split(" ");
         if (split.length == 0 || !split[0].equals("/spark")) {
-            return 0;
+            return null;
         }
 
-        String[] args = Arrays.copyOfRange(split, 1, split.length);
+        return Arrays.copyOfRange(split, 1, split.length);
+    }
+
+    @Override
+    public int run(CommandContext<CommandSource> context) throws CommandSyntaxException {
+        String[] args = processArgs(context);
+        if (args == null)
+            return 0;
 
         this.platform.executeCommand(new ForgeCommandSender(context.getSource().asPlayer(), this), args);
-        return 1;
+        return Command.SINGLE_SUCCESS;
+    }
+
+    @Override
+    public CompletableFuture<Suggestions> getSuggestions(CommandContext<CommandSource> context, SuggestionsBuilder builder)
+            throws CommandSyntaxException {
+        String[] args = processArgs(context);
+        if (args == null)
+            return Suggestions.empty();
+
+        ServerPlayerEntity player = context.getSource().asPlayer();
+        return CompletableFuture.supplyAsync(() -> {
+            for (String each : this.platform.tabCompleteCommand(new ForgeCommandSender(player, this), args)) {
+                builder.suggest(each);
+            }
+            return builder.build();
+        });
     }
 
     @Override
diff --git a/spark-forge/src/main/java/me/lucko/spark/forge/plugin/ForgeSparkPlugin.java b/spark-forge/src/main/java/me/lucko/spark/forge/plugin/ForgeSparkPlugin.java
index 915ca6d..0f0ce11 100644
--- a/spark-forge/src/main/java/me/lucko/spark/forge/plugin/ForgeSparkPlugin.java
+++ b/spark-forge/src/main/java/me/lucko/spark/forge/plugin/ForgeSparkPlugin.java
@@ -26,6 +26,7 @@ import com.mojang.brigadier.CommandDispatcher;
 import com.mojang.brigadier.arguments.StringArgumentType;
 import com.mojang.brigadier.builder.LiteralArgumentBuilder;
 import com.mojang.brigadier.builder.RequiredArgumentBuilder;
+import com.mojang.brigadier.suggestion.SuggestionProvider;
 import com.mojang.brigadier.tree.LiteralCommandNode;
 import me.lucko.spark.common.SparkPlatform;
 import me.lucko.spark.common.SparkPlugin;
@@ -39,7 +40,7 @@ import java.util.concurrent.ScheduledExecutorService;
 
 public abstract class ForgeSparkPlugin implements SparkPlugin {
 
-    public static <T> void registerCommands(CommandDispatcher<T> dispatcher, Command<T> executor, String... aliases) {
+    public static <T> void registerCommands(CommandDispatcher<T> dispatcher, Command<T> executor, SuggestionProvider<T> suggestor, String... aliases) {
         if (aliases.length == 0) {
             return;
         }
@@ -48,6 +49,7 @@ public abstract class ForgeSparkPlugin implements SparkPlugin {
         LiteralArgumentBuilder<T> command = LiteralArgumentBuilder.<T>literal(mainName)
                 .executes(executor)
                 .then(RequiredArgumentBuilder.<T, String>argument("args", StringArgumentType.greedyString())
+                        .suggests(suggestor)
                         .executes(executor)
                 );