Fix partial command suggestions when using a suggestion provider

Signed-off-by: Irmo van den Berge <irmo.vandenberge@ziggo.nl>

cloud-core/src/main/java/cloud/commandframework/CommandTree.java

@@ -449,87 +449,7 @@ public final class CommandTree<C> {
         }
         final List<Node<CommandArgument<C, ?>>> children = root.getChildren();
         if (children.size() == 1 && !(children.get(0).getValue() instanceof StaticArgument)) {
-            // The value has to be a variable
+            return this.suggestionsForDynamicArgument(commandContext, commandQueue, children.get(0));
-            final Node<CommandArgument<C, ?>> child = children.get(0);
-
-            /* When we get in here, we need to treat compound arguments a little differently */
-            if (child.getValue() instanceof CompoundArgument) {
-                @SuppressWarnings("unchecked") final CompoundArgument<?, C, ?> compoundArgument = (CompoundArgument<?, C, ?>) child
-                        .getValue();
-                /* See how many arguments it requires */
-                final int requiredArguments = compoundArgument.getParserTuple().getSize();
-                /* Figure out whether we even need to care about this */
-                if (commandQueue.size() <= requiredArguments) {
-                    /* Attempt to pop as many arguments from the stack as possible */
-                    for (int i = 0; i < requiredArguments - 1 && commandQueue.size() > 1; i++) {
-                        commandQueue.remove();
-                        commandContext.store("__parsing_argument__", i + 2);
-                    }
-                }
-            } else if (child.getValue() instanceof FlagArgument) {
-                /* Remove all but last */
-                while (commandQueue.size() > 1) {
-                    commandContext.store(FlagArgument.FLAG_META, commandQueue.remove());
-                }
-            } else if (child.getValue() != null
-                    && GenericTypeReflector.erase(child.getValue().getValueType().getType()).isArray()) {
-                while (commandQueue.size() > 1) {
-                    commandQueue.remove();
-                }
-            } else if (child.getValue() != null
-                    && commandQueue.size() <= child.getValue().getParser().getRequestedArgumentCount()) {
-                for (int i = 0; i < child.getValue().getParser().getRequestedArgumentCount() - 1
-                        && commandQueue.size() > 1; i++) {
-                    commandContext.store(
-                            String.format("%s_%d", child.getValue().getName(), i),
-                            commandQueue.remove()
-                    );
-                }
-            }
-
-            if (child.getValue() != null) {
-                if (commandQueue.isEmpty()) {
-                    return Collections.emptyList();
-                } else if (child.isLeaf() && commandQueue.size() < 2) {
-                    commandContext.setCurrentArgument(child.getValue());
-                    return child.getValue().getSuggestionsProvider().apply(commandContext, commandQueue.peek());
-                } else if (child.isLeaf()) {
-                    if (child.getValue() instanceof CompoundArgument) {
-                        final String last = ((LinkedList<String>) commandQueue).getLast();
-                        commandContext.setCurrentArgument(child.getValue());
-                        return child.getValue().getSuggestionsProvider().apply(commandContext, last);
-                    }
-                    return Collections.emptyList();
-                } else if (commandQueue.peek().isEmpty()) {
-                    commandContext.setCurrentArgument(child.getValue());
-                    return child.getValue().getSuggestionsProvider().apply(commandContext, commandQueue.remove());
-                }
-
-                // START: Preprocessing
-                final ArgumentParseResult<Boolean> preParseResult = child.getValue().preprocess(
-                        commandContext,
-                        commandQueue
-                );
-                if (preParseResult.getFailure().isPresent() || !preParseResult.getParsedValue().orElse(false)) {
-                    final String value = commandQueue.peek() == null ? "" : commandQueue.peek();
-                    commandContext.setCurrentArgument(child.getValue());
-                    return child.getValue().getSuggestionsProvider().apply(commandContext, value);
-                }
-                // END: Preprocessing
-
-                // START: Parsing
-                commandContext.setCurrentArgument(child.getValue());
-                final ArgumentParseResult<?> result = child.getValue().getParser().parse(commandContext, commandQueue);
-                if (result.getParsedValue().isPresent()) {
-                    commandContext.store(child.getValue().getName(), result.getParsedValue().get());
-                    return this.getSuggestions(commandContext, commandQueue, child);
-                } else if (result.getFailure().isPresent()) {
-                    final String value = commandQueue.peek() == null ? "" : commandQueue.peek();
-                    commandContext.setCurrentArgument(child.getValue());
-                    return child.getValue().getSuggestionsProvider().apply(commandContext, value);
-                }
-                // END: Parsing
-            }
         }
         /* There are 0 or more static arguments as children. No variable child arguments are present */
         if (children.isEmpty() || commandQueue.isEmpty()) {
@@ -573,6 +493,99 @@ public final class CommandTree<C> {
         }
     }
 
+    private @NonNull List<String> suggestionsForDynamicArgument(
+            final @NonNull CommandContext<C> commandContext,
+            final @NonNull Queue<@NonNull String> commandQueue,
+            final @NonNull Node<CommandArgument<C, ?>> child
+    ) {
+        /* When we get in here, we need to treat compound arguments a little differently */
+        if (child.getValue() instanceof CompoundArgument) {
+            @SuppressWarnings("unchecked") final CompoundArgument<?, C, ?> compoundArgument = (CompoundArgument<?, C, ?>) child
+                    .getValue();
+            /* See how many arguments it requires */
+            final int requiredArguments = compoundArgument.getParserTuple().getSize();
+            /* Figure out whether we even need to care about this */
+            if (commandQueue.size() <= requiredArguments) {
+                /* Attempt to pop as many arguments from the stack as possible */
+                for (int i = 0; i < requiredArguments - 1 && commandQueue.size() > 1; i++) {
+                    commandQueue.remove();
+                    commandContext.store("__parsing_argument__", i + 2);
+                }
+            }
+        } else if (child.getValue() instanceof FlagArgument) {
+            /* Remove all but last */
+            while (commandQueue.size() > 1) {
+                commandContext.store(FlagArgument.FLAG_META, commandQueue.remove());
+            }
+        } else if (child.getValue() != null
+                && GenericTypeReflector.erase(child.getValue().getValueType().getType()).isArray()) {
+            while (commandQueue.size() > 1) {
+                commandQueue.remove();
+            }
+        } else if (child.getValue() != null
+                && commandQueue.size() <= child.getValue().getParser().getRequestedArgumentCount()) {
+            for (int i = 0; i < child.getValue().getParser().getRequestedArgumentCount() - 1
+                    && commandQueue.size() > 1; i++) {
+                commandContext.store(
+                        String.format("%s_%d", child.getValue().getName(), i),
+                        commandQueue.remove()
+                );
+            }
+        }
+
+        if (child.getValue() != null) {
+            if (commandQueue.isEmpty()) {
+                return Collections.emptyList();
+            } else if (child.isLeaf() && commandQueue.size() < 2) {
+                commandContext.setCurrentArgument(child.getValue());
+                return child.getValue().getSuggestionsProvider().apply(commandContext, commandQueue.peek());
+            } else if (child.isLeaf()) {
+                if (child.getValue() instanceof CompoundArgument) {
+                    final String last = ((LinkedList<String>) commandQueue).getLast();
+                    commandContext.setCurrentArgument(child.getValue());
+                    return child.getValue().getSuggestionsProvider().apply(commandContext, last);
+                }
+                return Collections.emptyList();
+            } else if (commandQueue.peek().isEmpty()) {
+                commandContext.setCurrentArgument(child.getValue());
+                return child.getValue().getSuggestionsProvider().apply(commandContext, commandQueue.remove());
+            }
+
+            // Store original input command queue before the parsers below modify it
+            final Queue<String> commandQueueOriginal = new LinkedList<String>(commandQueue);
+
+            // START: Preprocessing
+            final ArgumentParseResult<Boolean> preParseResult = child.getValue().preprocess(
+                    commandContext,
+                    commandQueue
+            );
+            final boolean preParseSuccess = !preParseResult.getFailure().isPresent()
+                    && preParseResult.getParsedValue().orElse(false);
+            // END: Preprocessing
+
+            if (preParseSuccess) {
+                // START: Parsing
+                commandContext.setCurrentArgument(child.getValue());
+                final ArgumentParseResult<?> result = child.getValue().getParser().parse(commandContext, commandQueue);
+                if (result.getParsedValue().isPresent() && !commandQueue.isEmpty()) {
+                    commandContext.store(child.getValue().getName(), result.getParsedValue().get());
+                    return this.getSuggestions(commandContext, commandQueue, child);
+                }
+                // END: Parsing
+            }
+
+            // Restore original command input queue
+            commandQueue.clear();
+            commandQueue.addAll(commandQueueOriginal);
+
+            // Fallback: use suggestion provider of argument
+            commandContext.setCurrentArgument(child.getValue());
+            return child.getValue().getSuggestionsProvider().apply(commandContext, stringOrEmpty(commandQueue.peek()));
+        }
+
+        return Collections.emptyList();
+    }
+
     private @NonNull String stringOrEmpty(final @Nullable String string) {
         if (string == null) {
             return "";

cloud-core/src/main/java/cloud/commandframework/arguments/standard/StringArgument.java

@@ -306,14 +306,6 @@ public final class StringArgument<C> extends CommandArgument<C, String> {
             }
 
             if (this.stringMode == StringMode.SINGLE) {
-                if (commandContext.isSuggestions()) {
-                    final List<String> suggestions = this.suggestionsProvider.apply(commandContext, inputQueue.peek());
-                    if (!suggestions.isEmpty() && !suggestions.contains(input)) {
-                        return ArgumentParseResult.failure(new IllegalArgumentException(
-                                String.format("'%s' is not one of: %s", input, String.join(", ", suggestions))
-                        ));
-                    }
-                }
                 inputQueue.remove();
                 return ArgumentParseResult.success(input);
             } else if (this.stringMode == StringMode.QUOTED) {

cloud-core/src/test/java/cloud/commandframework/CommandSuggestionsTest.java

@@ -90,6 +90,13 @@ public class CommandSuggestionsTest {
 
         manager.command(manager.commandBuilder("numberswithmin")
                 .argument(IntegerArgument.<TestCommandSender>newBuilder("num").withMin(5).withMax(100)));
+
+        manager.command(manager.commandBuilder("partial")
+                .argument(StringArgument.<TestCommandSender>newBuilder("arg").withSuggestionsProvider((contect, input) -> {
+                    return Arrays.asList("hi", "hey", "heya", "hai", "hello");
+                }))
+                .literal("literal")
+                .build());
     }
 
     @Test
@@ -121,7 +128,7 @@ public class CommandSuggestionsTest {
         Assertions.assertTrue(suggestions.isEmpty());
         final String input2 = "test var one";
         final List<String> suggestions2 = manager.suggest(new TestCommandSender(), input2);
-        Assertions.assertEquals(Collections.emptyList(), suggestions2);
+        Assertions.assertEquals(Collections.singletonList("one"), suggestions2);
         final String input3 = "test var one f";
         final List<String> suggestions3 = manager.suggest(new TestCommandSender(), input3);
         Assertions.assertEquals(Collections.singletonList("foo"), suggestions3);
@@ -228,6 +235,47 @@ public class CommandSuggestionsTest {
         Assertions.assertEquals(Collections.emptyList(), suggestions3);
     }
 
+    @Test
+    void testStringArgumentWithSuggestionProvider() {
+        /*
+         * [/partial] - should not match anything
+         * [/partial ] - should show all possible suggestions unsorted
+         * [/partial h] - should show all starting with 'h' (which is all) unsorted
+         * [/partial he] - should show only those starting with he, unsorted
+         * [/partial hey] - should show 'hey' and 'heya' (matches exactly and starts with)
+         * [/partial hi] - should show only 'hi', it is the only one that matches exactly
+         * [/partial b] - should show no suggestions, none match
+         * [/partial hello ] - should show the literal following the argument (suggested)
+         * [/partial bonjour ] - should show the literal following the argument (not suggested)
+         */
+        final String input = "partial";
+        final List<String> suggestions = manager.suggest(new TestCommandSender(), input);
+        Assertions.assertEquals(Collections.emptyList(), suggestions);
+        final String input2 = "partial ";
+        final List<String> suggestions2 = manager.suggest(new TestCommandSender(), input2);
+        Assertions.assertEquals(Arrays.asList("hi", "hey", "heya", "hai", "hello"), suggestions2);
+        final String input3 = "partial h";
+        final List<String> suggestions3 = manager.suggest(new TestCommandSender(), input3);
+        Assertions.assertEquals(Arrays.asList("hi", "hey", "heya", "hai", "hello"), suggestions3);
+        final String input4 = "partial he";
+        final List<String> suggestions4 = manager.suggest(new TestCommandSender(), input4);
+        Assertions.assertEquals(Arrays.asList("hey", "heya", "hello"), suggestions4);
+        final String input5 = "partial hey";
+        final List<String> suggestions5 = manager.suggest(new TestCommandSender(), input5);
+        Assertions.assertEquals(Arrays.asList("hey", "heya"), suggestions5);
+        final String input6 = "partial hi";
+        final List<String> suggestions6 = manager.suggest(new TestCommandSender(), input6);
+        Assertions.assertEquals(Collections.singletonList("hi"), suggestions6);
+        final String input7 = "partial b";
+        final List<String> suggestions7 = manager.suggest(new TestCommandSender(), input7);
+        Assertions.assertEquals(Collections.emptyList(), suggestions7);
+        final String input8 = "partial hello ";
+        final List<String> suggestions8 = manager.suggest(new TestCommandSender(), input8);
+        Assertions.assertEquals(Collections.singletonList("literal"), suggestions8);
+        final String input9 = "partial bonjour ";
+        final List<String> suggestions9 = manager.suggest(new TestCommandSender(), input9);
+        Assertions.assertEquals(Collections.singletonList("literal"), suggestions9);
+    }
 
     public enum TestEnum {
         FOO,