Fixed issues with suggesting redirects

build.gradle

@@ -3,7 +3,7 @@ import groovy.io.FileType
 apply plugin: 'java-library'
 apply plugin: 'maven'
 
-version = '0.1.10'
+version = '0.1.11'
 group = 'com.mojang'
 
 task wrapper(type: Wrapper) {

src/main/java/com/mojang/brigadier/CommandDispatcher.java

@@ -177,11 +177,11 @@ public class CommandDispatcher<S> {
             }
 
             context.withCommand(child.getCommand());
-            if (reader.canRead(2)) {
+            if (reader.canRead(child.getRedirect() == null ? 2 : 1)) {
                 reader.skip();
                 if (child.getRedirect() != null) {
                     final CommandContextBuilder<S> childContext = new CommandContextBuilder<>(this, source, reader.getCursor());
-                    childContext.withNode(child.getRedirect(), new StringRange(reader.getCursor(), reader.getCursor()));
+                    childContext.withNode(child.getRedirect(), new StringRange(cursor, reader.getCursor() - 1));
                     final ParseResults<S> parse = parseNodes(child.getRedirect(), reader, childContext);
                     context.withChild(parse.getContext());
                     return new ParseResults<>(context, parse.getReader(), parse.getExceptions());
@@ -311,7 +311,8 @@ public class CommandDispatcher<S> {
     }
 
     public CompletableFuture<CommandSuggestions> getCompletionSuggestions(final ParseResults<S> parse) {
-        final CommandContextBuilder<S> context = parse.getContext();
+        final CommandContextBuilder<S> rootContext = parse.getContext();
+        final CommandContextBuilder<S> context = rootContext.getLastChild();
         final CommandNode<S> parent;
         final int start;
 
@@ -326,6 +327,10 @@ public class CommandDispatcher<S> {
             final Map.Entry<CommandNode<S>, StringRange> entry = Iterables.get(context.getNodes().entrySet(), context.getNodes().size() - 2);
             parent = entry.getKey();
             start = entry.getValue().getEnd() + 1;
+        } else if (rootContext != context && context.getNodes().size() > 0) {
+            final Map.Entry<CommandNode<S>, StringRange> entry = Iterables.getLast(context.getNodes().entrySet());
+            parent = entry.getKey();
+            start = entry.getValue().getEnd() + 1;
         } else {
             parent = root;
             start = 0;

src/main/java/com/mojang/brigadier/context/CommandContext.java

@@ -30,6 +30,14 @@ public class CommandContext<S> {
         return child;
     }
 
+    public CommandContext<S> getLastChild() {
+        CommandContext<S> result = this;
+        while (result.getChild() != null) {
+            result = result.getChild();
+        }
+        return result;
+    }
+
     public Command<S> getCommand() {
         return command;
     }

src/main/java/com/mojang/brigadier/context/CommandContextBuilder.java

@@ -70,6 +70,14 @@ public class CommandContextBuilder<S> {
         return child;
     }
 
+    public CommandContextBuilder<S> getLastChild() {
+        CommandContextBuilder<S> result = this;
+        while (result.getChild() != null) {
+            result = result.getChild();
+        }
+        return result;
+    }
+
     public Command<S> getCommand() {
         return command;
     }

src/main/java/com/mojang/brigadier/tree/LiteralCommandNode.java

@@ -50,9 +50,9 @@ public class LiteralCommandNode<S> extends CommandNode<S> {
     }
 
     @Override
-    public CompletableFuture<Collection<String>> listSuggestions(CommandContext<S> context, final String command) {
+    public CompletableFuture<Collection<String>> listSuggestions(final CommandContext<S> context, final String command) {
         if (literal.toLowerCase().startsWith(command.toLowerCase())) {
-            return CompletableFuture.completedFuture(Collections.singleton(literal));
+            return CompletableFuture.completedFuture(Collections.singleton(literal + " "));
         } else {
             return CompletableFuture.completedFuture(Collections.emptyList());
         }

src/test/java/com/mojang/brigadier/CommandDispatcherTest.java

@@ -259,12 +259,12 @@ public class CommandDispatcherTest {
 
         final CommandContextBuilder<Object> child1 = parse.getContext().getChild();
         assertThat(child1, is(notNullValue()));
-        assertThat(child1.getRange().get(input), equalTo("redirected"));
+        assertThat(child1.getRange().get(input), equalTo("redirected redirected"));
         assertThat(child1.getNodes().size(), is(2));
 
         final CommandContextBuilder<Object> child2 = child1.getChild();
         assertThat(child2, is(notNullValue()));
-        assertThat(child2.getRange().get(input), equalTo("actual"));
+        assertThat(child2.getRange().get(input), equalTo("redirected actual"));
         assertThat(child2.getNodes().size(), is(2));
 
         assertThat(subject.execute(parse), is(42));
@@ -291,7 +291,7 @@ public class CommandDispatcherTest {
 
         final CommandContextBuilder<Object> parent = parse.getContext().getChild();
         assertThat(parent, is(notNullValue()));
-        assertThat(parent.getRange().get(input), equalTo("actual"));
+        assertThat(parent.getRange().get(input), equalTo("redirected actual"));
         assertThat(parent.getNodes().size(), is(2));
         assertThat(parent.getSource(), is(source));
 

src/test/java/com/mojang/brigadier/CommandSuggestionsTest.java

@@ -2,13 +2,17 @@ package com.mojang.brigadier;
 
 import com.google.common.collect.Lists;
 import com.mojang.brigadier.context.StringRange;
+import com.mojang.brigadier.tree.LiteralCommandNode;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.runners.MockitoJUnitRunner;
 
+import static com.mojang.brigadier.arguments.IntegerArgumentType.integer;
+import static com.mojang.brigadier.arguments.StringArgumentType.word;
 import static com.mojang.brigadier.builder.LiteralArgumentBuilder.literal;
+import static com.mojang.brigadier.builder.RequiredArgumentBuilder.argument;
 import static org.hamcrest.Matchers.equalTo;
 import static org.junit.Assert.assertThat;
 
@@ -71,9 +75,116 @@ public class CommandSuggestionsTest {
                 .then(literal("baz"))
         );
 
-        final CommandSuggestions result = subject.getCompletionSuggestions(subject.parse("parent b", source)).join();
+        final ParseResults<Object> parse = subject.parse("parent b", source);
+        final CommandSuggestions result = subject.getCompletionSuggestions(parse).join();
 
         assertThat(result.getRange(), equalTo(new StringRange(7, 8)));
         assertThat(result.getSuggestions(), equalTo(Lists.newArrayList("bar", "baz")));
     }
+
+    @Test
+    public void getCompletionSuggestions_redirect() throws Exception {
+        final LiteralCommandNode<Object> actual = subject.register(literal("actual").then(literal("sub")));
+        subject.register(literal("redirect").redirect(actual));
+
+        final ParseResults<Object> parse = subject.parse("redirect ", source);
+        final CommandSuggestions result = subject.getCompletionSuggestions(parse).join();
+
+        assertThat(result.getRange(), equalTo(new StringRange(9, 9)));
+        assertThat(result.getSuggestions(), equalTo(Lists.newArrayList("sub")));
+    }
+
+    @Test
+    public void getCompletionSuggestions_redirectPartial() throws Exception {
+        final LiteralCommandNode<Object> actual = subject.register(literal("actual").then(literal("sub")));
+        subject.register(literal("redirect").redirect(actual));
+
+        final ParseResults<Object> parse = subject.parse("redirect s", source);
+        final CommandSuggestions result = subject.getCompletionSuggestions(parse).join();
+
+        assertThat(result.getRange(), equalTo(new StringRange(9, 10)));
+        assertThat(result.getSuggestions(), equalTo(Lists.newArrayList("sub")));
+    }
+
+    @Test
+    public void getCompletionSuggestions_redirect_lots() throws Exception {
+        final LiteralCommandNode<Object> loop = subject.register(literal("redirect"));
+        subject.register(
+            literal("redirect")
+                .then(
+                    literal("loop")
+                        .then(
+                            argument("loop", integer())
+                                .redirect(loop)
+                        )
+                )
+        );
+
+        final CommandSuggestions result = subject.getCompletionSuggestions(subject.parse("redirect loop 1 loop 02 loop 003 ", source)).join();
+
+        assertThat(result.getRange(), equalTo(new StringRange(33, 33)));
+        assertThat(result.getSuggestions(), equalTo(Lists.newArrayList("loop")));
+    }
+
+    @Test
+    public void getCompletionSuggestions_execute_simulation() throws Exception {
+        final LiteralCommandNode<Object> execute = subject.register(literal("execute"));
+        subject.register(
+            literal("execute")
+                .then(
+                    literal("as")
+                        .then(
+                            argument("name", word())
+                                .redirect(execute)
+                        )
+                )
+                .then(
+                    literal("store")
+                        .then(
+                            argument("name", word())
+                                .redirect(execute)
+                        )
+                )
+                .then(
+                    literal("run")
+                        .executes(c -> 0)
+                )
+        );
+
+        final ParseResults<Object> parse = subject.parse("execute as Dinnerbone as", source);
+        final CommandSuggestions result = subject.getCompletionSuggestions(parse).join();
+
+        assertThat(result.getRange(), equalTo(new StringRange(22, 24)));
+        assertThat(result.getSuggestions(), equalTo(Lists.newArrayList("as")));
+    }
+
+    @Test
+    public void getCompletionSuggestions_execute_simulation_partial() throws Exception {
+        final LiteralCommandNode<Object> execute = subject.register(literal("execute"));
+        subject.register(
+            literal("execute")
+                .then(
+                    literal("as")
+                        .then(literal("bar").redirect(execute))
+                        .then(literal("baz").redirect(execute))
+                )
+                .then(
+                    literal("store")
+                        .then(
+                            argument("name", word())
+                                .redirect(execute)
+                        )
+                )
+                .then(
+                    literal("run")
+                        .executes(c -> 0)
+                )
+        );
+
+        final ParseResults<Object> parse = subject.parse("execute as bar as ", source);
+        final CommandSuggestions result = subject.getCompletionSuggestions(parse).join();
+
+        assertThat(result.getRange(), equalTo(new StringRange(18, 18)));
+        assertThat(result.getSuggestions(), equalTo(Lists.newArrayList("bar", "baz")));
+    }
 }