Fix incorrect context being used for suggestions (MC-278771)

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

@@ -538,7 +538,7 @@ public CompletableFuture<Suggestions> getCompletionSuggestions(final ParseResult
         for (final CommandNode<S> node : parent.getChildren()) {
             CompletableFuture<Suggestions> future = Suggestions.empty();
             try {
-                future = node.listSuggestions(context.build(truncatedInput), new SuggestionsBuilder(truncatedInput, truncatedInputLowerCase, start));
+                future = node.listSuggestions(nodeBeforeCursor.context.build(truncatedInput), new SuggestionsBuilder(truncatedInput, truncatedInputLowerCase, start));
             } catch (final CommandSyntaxException ignored) {
             }
             futures[i++] = future;

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

@@ -122,23 +122,23 @@ public SuggestionContext<S> findSuggestionContext(final int cursor) {
                     return child.findSuggestionContext(cursor);
                 } else if (!nodes.isEmpty()) {
                     final ParsedCommandNode<S> last = nodes.get(nodes.size() - 1);
-                    return new SuggestionContext<>(last.getNode(), last.getRange().getEnd() + 1);
+                    return new SuggestionContext<>(this, last.getNode(), last.getRange().getEnd() + 1);
                 } else {
-                    return new SuggestionContext<>(rootNode, range.getStart());
+                    return new SuggestionContext<>(this, rootNode, range.getStart());
                 }
             } else {
                 CommandNode<S> prev = rootNode;
                 for (final ParsedCommandNode<S> node : nodes) {
                     final StringRange nodeRange = node.getRange();
                     if (nodeRange.getStart() <= cursor && cursor <= nodeRange.getEnd()) {
-                        return new SuggestionContext<>(prev, nodeRange.getStart());
+                        return new SuggestionContext<>(this, prev, nodeRange.getStart());
                     }
                     prev = node.getNode();
                 }
                 if (prev == null) {
                     throw new IllegalStateException("Can't find node before cursor");
                 }
-                return new SuggestionContext<>(prev, range.getStart());
+                return new SuggestionContext<>(this, prev, range.getStart());
             }
         }
         throw new IllegalStateException("Can't find node before cursor");

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

@@ -6,10 +6,12 @@
 import com.mojang.brigadier.tree.CommandNode;
 
 public class SuggestionContext<S> {
+    public final CommandContextBuilder<S> context;
     public final CommandNode<S> parent;
     public final int startPos;
 
-    public SuggestionContext(CommandNode<S> parent, int startPos) {
+    public SuggestionContext(CommandContextBuilder<S> context, CommandNode<S> parent, int startPos) {
+        this.context = context;
         this.parent = parent;
         this.startPos = startPos;
     }

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

@@ -4,8 +4,10 @@
 package com.mojang.brigadier;
 
 import com.google.common.collect.Lists;
+import com.mojang.brigadier.arguments.StringArgumentType;
 import com.mojang.brigadier.context.StringRange;
 import com.mojang.brigadier.suggestion.Suggestion;
+import com.mojang.brigadier.suggestion.SuggestionProvider;
 import com.mojang.brigadier.suggestion.Suggestions;
 import com.mojang.brigadier.tree.LiteralCommandNode;
 import org.junit.Before;
@@ -253,6 +255,28 @@ public void getCompletionSuggestions_redirect_lots() throws Exception {
         assertThat(result.getList(), equalTo(Lists.newArrayList(new Suggestion(StringRange.at(33), "loop"))));
     }
 
+    @Test
+    public void getCompletionSuggestions_redirect_contextualArgument() {
+        final LiteralCommandNode<Object> actual = subject.register(
+            literal("actual")
+                .then(argument("arg_one", word())
+                    .then(argument("arg_two", word())
+                        .suggests((context, builder) -> {
+                            final String argOne = StringArgumentType.getString(context, "arg_one");
+                            builder.suggest("contextual_" + argOne);
+                            return builder.buildFuture();
+                        })
+                    )
+                )
+        );
+        subject.register(literal("redirect").redirect(actual));
+
+        final Suggestions result = subject.getCompletionSuggestions(subject.parse("redirect first ", source)).join();
+
+        assertThat(result.getRange(), equalTo(StringRange.at(15)));
+        assertThat(result.getList(), equalTo(Lists.newArrayList(new Suggestion(StringRange.at(15), "contextual_first"))));
+    }
+
     @Test
     public void getCompletionSuggestions_execute_simulation() throws Exception {
         final LiteralCommandNode<Object> execute = subject.register(literal("execute"));