Fixed redirect modifiers not always being ran

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

@@ -1,9 +1,11 @@
 package com.mojang.brigadier;
 
+import com.google.common.collect.Iterators;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
 import com.mojang.brigadier.builder.LiteralArgumentBuilder;
+import com.mojang.brigadier.context.CommandContext;
 import com.mojang.brigadier.context.CommandContextBuilder;
 import com.mojang.brigadier.exceptions.CommandException;
 import com.mojang.brigadier.exceptions.SimpleCommandExceptionType;
@@ -79,11 +81,12 @@ public class CommandDispatcher<S> {
 
         while (!contexts.isEmpty()) {
             final CommandContextBuilder<S> context = contexts.removeLast();
-            if (context.getChild() != null) {
+            final CommandContextBuilder<S> child = context.getChild();
-                if (!context.getNodes().isEmpty()) {
+            if (child != null) {
-                    final Function<S, Collection<S>> modifier = context.getNodes().keySet().iterator().next().getRedirectModifier();
+                if (!child.getNodes().isEmpty()) {
-                    for (final S source : modifier.apply(context.getSource())) {
+                    final Function<CommandContext<S>, Collection<S>> modifier = Iterators.getLast(context.getNodes().keySet().iterator()).getRedirectModifier();
-                        contexts.add(context.getChild().copy().withSource(source));
+                    for (final S source : modifier.apply(context.build())) {
+                        contexts.add(child.copy().withSource(source));
                     }
                 }
             } else if (context.getCommand() != null) {
@@ -151,6 +154,10 @@ public class CommandDispatcher<S> {
             }
         }
 
+        if (parentContext != null) {
+            parentContext.withChild(contextBuilder);
+        }
+
         return new ParseResults<>(rootContext, reader, errors);
     }
 

src/main/java/com/mojang/brigadier/builder/ArgumentBuilder.java

@@ -1,6 +1,7 @@
 package com.mojang.brigadier.builder;
 
 import com.mojang.brigadier.Command;
+import com.mojang.brigadier.context.CommandContext;
 import com.mojang.brigadier.tree.CommandNode;
 import com.mojang.brigadier.tree.RootCommandNode;
 
@@ -14,7 +15,7 @@ public abstract class ArgumentBuilder<S, T extends ArgumentBuilder<S, T>> {
     private Command<S> command;
     private Predicate<S> requirement = s -> true;
     private CommandNode<S> target;
-    private Function<S, Collection<S>> modifier = Collections::singleton;
+    private Function<CommandContext<S>, Collection<S>> modifier = s -> Collections.singleton(s.getSource());
 
     protected abstract T getThis();
 
@@ -48,7 +49,11 @@ public abstract class ArgumentBuilder<S, T extends ArgumentBuilder<S, T>> {
         return requirement;
     }
 
-    public T redirect(final CommandNode<S> target, final Function<S, Collection<S>> modifier) {
+    public T redirect(final CommandNode<S> target) {
+        return redirect(target, modifier);
+    }
+
+    public T redirect(final CommandNode<S> target, final Function<CommandContext<S>, Collection<S>> modifier) {
         if (!arguments.getChildren().isEmpty()) {
             throw new IllegalStateException("Cannot redirect a node with children");
         }
@@ -61,7 +66,7 @@ public abstract class ArgumentBuilder<S, T extends ArgumentBuilder<S, T>> {
         return target;
     }
 
-    public Function<S, Collection<S>> getRedirectModifier() {
+    public Function<CommandContext<S>, Collection<S>> getRedirectModifier() {
         return modifier;
     }
 

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

@@ -4,6 +4,7 @@ import com.mojang.brigadier.Command;
 import com.mojang.brigadier.StringReader;
 import com.mojang.brigadier.arguments.ArgumentType;
 import com.mojang.brigadier.builder.RequiredArgumentBuilder;
+import com.mojang.brigadier.context.CommandContext;
 import com.mojang.brigadier.context.CommandContextBuilder;
 import com.mojang.brigadier.context.ParsedArgument;
 import com.mojang.brigadier.exceptions.CommandException;
@@ -20,7 +21,7 @@ public class ArgumentCommandNode<S, T> extends CommandNode<S> {
     private final String name;
     private final ArgumentType<T> type;
 
-    public ArgumentCommandNode(final String name, final ArgumentType<T> type, final Command<S> command, final Predicate<S> requirement, final CommandNode<S> redirect, final Function<S, Collection<S>> modifier) {
+    public ArgumentCommandNode(final String name, final ArgumentType<T> type, final Command<S> command, final Predicate<S> requirement, final CommandNode<S> redirect, final Function<CommandContext<S>, Collection<S>> modifier) {
         super(command, requirement, redirect, modifier);
         this.name = name;
         this.type = type;

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

@@ -5,6 +5,7 @@ import com.google.common.collect.Maps;
 import com.mojang.brigadier.Command;
 import com.mojang.brigadier.StringReader;
 import com.mojang.brigadier.builder.ArgumentBuilder;
+import com.mojang.brigadier.context.CommandContext;
 import com.mojang.brigadier.context.CommandContextBuilder;
 import com.mojang.brigadier.exceptions.CommandException;
 
@@ -20,10 +21,10 @@ public abstract class CommandNode<S> implements Comparable<CommandNode<S>> {
     private Map<Object, CommandNode<S>> children = Maps.newLinkedHashMap();
     private final Predicate<S> requirement;
     private final CommandNode<S> redirect;
-    private final Function<S, Collection<S>> modifier;
+    private final Function<CommandContext<S>, Collection<S>> modifier;
     private Command<S> command;
 
-    protected CommandNode(final Command<S> command, final Predicate<S> requirement, final CommandNode<S> redirect, final Function<S, Collection<S>> modifier) {
+    protected CommandNode(final Command<S> command, final Predicate<S> requirement, final CommandNode<S> redirect, final Function<CommandContext<S>, Collection<S>> modifier) {
         this.command = command;
         this.requirement = requirement;
         this.redirect = redirect;
@@ -42,7 +43,7 @@ public abstract class CommandNode<S> implements Comparable<CommandNode<S>> {
         return redirect;
     }
 
-    public Function<S, Collection<S>> getRedirectModifier() {
+    public Function<CommandContext<S>, Collection<S>> getRedirectModifier() {
         return modifier;
     }
 

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

@@ -3,6 +3,7 @@ package com.mojang.brigadier.tree;
 import com.mojang.brigadier.Command;
 import com.mojang.brigadier.StringReader;
 import com.mojang.brigadier.builder.LiteralArgumentBuilder;
+import com.mojang.brigadier.context.CommandContext;
 import com.mojang.brigadier.context.CommandContextBuilder;
 import com.mojang.brigadier.exceptions.CommandException;
 import com.mojang.brigadier.exceptions.ParameterizedCommandExceptionType;
@@ -17,7 +18,7 @@ public class LiteralCommandNode<S> extends CommandNode<S> {
 
     private final String literal;
 
-    public LiteralCommandNode(final String literal, final Command<S> command, final Predicate<S> requirement, final CommandNode<S> redirect, final Function<S, Collection<S>> modifier) {
+    public LiteralCommandNode(final String literal, final Command<S> command, final Predicate<S> requirement, final CommandNode<S> redirect, final Function<CommandContext<S>, Collection<S>> modifier) {
         super(command, requirement, redirect, modifier);
         this.literal = literal;
     }

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

@@ -11,7 +11,7 @@ import java.util.Set;
 
 public class RootCommandNode<S> extends CommandNode<S> {
     public RootCommandNode() {
-        super(null, c -> true, null, Collections::singleton);
+        super(null, c -> true, null, s -> Collections.singleton(s.getSource()));
     }
 
     @Override

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

@@ -200,11 +200,11 @@ public class CommandDispatcherTest {
     @SuppressWarnings("unchecked")
     @Test
     public void testExecuteRedirectedMultipleTimes() throws Exception {
-        final Function<Object, Collection<Object>> modifier = mock(Function.class);
+        final Function<CommandContext<Object>, Collection<Object>> modifier = mock(Function.class);
         final Object source1 = new Object();
         final Object source2 = new Object();
 
-        when(modifier.apply(source)).thenReturn(Lists.newArrayList(source1, source2));
+        when(modifier.apply(argThat(hasProperty("source", is(source))))).thenReturn(Lists.newArrayList(source1, source2));
 
         subject.register(literal("actual").executes(command));
         subject.register(literal("redirected").redirect(subject.getRoot(), modifier));