Allow redirects to modify context, supporting forks in the command chain

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

@@ -4,7 +4,6 @@ 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;
@@ -12,10 +11,13 @@ import com.mojang.brigadier.tree.CommandNode;
 import com.mojang.brigadier.tree.LiteralCommandNode;
 import com.mojang.brigadier.tree.RootCommandNode;
 
+import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Deque;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
 
@@ -69,20 +71,41 @@ public class CommandDispatcher<S> {
                 throw ERROR_UNKNOWN_ARGUMENT.createWithContext(parse.getReader());
             }
         }
-        final CommandContext<S> context = parse.getContext().build();
+
-        final Command<S> command = context.getCommand();
+        int result = 0;
-        if (command == null) {
+        boolean foundCommand = false;
+        final Deque<CommandContextBuilder<S>> contexts = new ArrayDeque<>();
+        contexts.add(parse.getContext());
+
+        while (!contexts.isEmpty()) {
+            final CommandContextBuilder<S> context = contexts.removeLast();
+            if (context.getChild() != null) {
+                if (!context.getNodes().isEmpty()) {
+                    final Function<S, Collection<S>> modifier = context.getNodes().keySet().iterator().next().getRedirectModifier();
+                    for (final S source : modifier.apply(context.getSource())) {
+                        contexts.add(context.getChild().copy().withSource(source));
+                    }
+                }
+            } else if (context.getCommand() != null) {
+                foundCommand = true;
+                result += context.getCommand().run(context.build());
+            }
+        }
+
+        if (!foundCommand) {
             throw ERROR_UNKNOWN_COMMAND.createWithContext(parse.getReader());
         }
-        return command.run(context);
+
+        return result;
     }
 
     public ParseResults<S> parse(final String command, final S source) throws CommandException {
         final StringReader reader = new StringReader(command);
-        return parseNodes(root, reader, new CommandContextBuilder<>(this, source));
+        final CommandContextBuilder<S> context = new CommandContextBuilder<>(this, source);
+        return parseNodes(root, reader, context, context, null);
     }
 
-    private ParseResults<S> parseNodes(final CommandNode<S> node, final StringReader reader, final CommandContextBuilder<S> contextBuilder) throws CommandException {
+    private ParseResults<S> parseNodes(final CommandNode<S> node, final StringReader reader, final CommandContextBuilder<S> contextBuilder, CommandContextBuilder<S> rootContext, final CommandContextBuilder<S> parentContext) throws CommandException {
         final S source = contextBuilder.getSource();
         final Map<CommandNode<S>, CommandException> errors = Maps.newHashMap();
 
@@ -105,20 +128,30 @@ public class CommandDispatcher<S> {
                 continue;
             }
 
+            if (rootContext == contextBuilder) {
+                rootContext = context;
+            }
+
+            if (parentContext != null) {
+                parentContext.withChild(context);
+            }
+
             context.withCommand(child.getCommand());
             if (reader.canRead()) {
                 reader.skip();
                 if (child.getRedirect() != null) {
-                    return parseNodes(child.getRedirect(), reader, context.redirect(child.getRedirect()));
+                    final CommandContextBuilder<S> childContext = new CommandContextBuilder<>(this, source);
+                    childContext.withNode(child.getRedirect(), "");
+                    return parseNodes(child.getRedirect(), reader, childContext, rootContext, context);
                 } else {
-                    return parseNodes(child, reader, context);
+                    return parseNodes(child, reader, context, rootContext, parentContext);
                 }
             } else {
-                return new ParseResults<>(context);
+                return new ParseResults<>(rootContext);
             }
         }
 
-        return new ParseResults<>(contextBuilder, reader, errors);
+        return new ParseResults<>(rootContext, reader, errors);
     }
 
     public String[] getAllUsage(final CommandNode<S> node, final S source) {

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

@@ -5,6 +5,8 @@ import com.mojang.brigadier.tree.CommandNode;
 import com.mojang.brigadier.tree.RootCommandNode;
 
 import java.util.Collection;
+import java.util.Collections;
+import java.util.function.Function;
 import java.util.function.Predicate;
 
 public abstract class ArgumentBuilder<S, T extends ArgumentBuilder<S, T>> {
@@ -12,6 +14,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;
 
     protected abstract T getThis();
 
@@ -45,11 +48,12 @@ public abstract class ArgumentBuilder<S, T extends ArgumentBuilder<S, T>> {
         return requirement;
     }
 
-    public T redirect(final CommandNode<S> target) {
+    public T redirect(final CommandNode<S> target, final Function<S, Collection<S>> modifier) {
         if (!arguments.getChildren().isEmpty()) {
             throw new IllegalStateException("Cannot redirect a node with children");
         }
         this.target = target;
+        this.modifier = modifier;
         return getThis();
     }
 
@@ -57,5 +61,9 @@ public abstract class ArgumentBuilder<S, T extends ArgumentBuilder<S, T>> {
         return target;
     }
 
+    public Function<S, Collection<S>> getRedirectModifier() {
+        return modifier;
+    }
+
     public abstract CommandNode<S> build();
 }

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

@@ -25,7 +25,7 @@ public class LiteralArgumentBuilder<S> extends ArgumentBuilder<S, LiteralArgumen
 
     @Override
     public LiteralCommandNode<S> build() {
-        final LiteralCommandNode<S> result = new LiteralCommandNode<>(getLiteral(), getCommand(), getRequirement(), getRedirect());
+        final LiteralCommandNode<S> result = new LiteralCommandNode<>(getLiteral(), getCommand(), getRequirement(), getRedirect(), getRedirectModifier());
 
         for (final CommandNode<S> argument : getArguments()) {
             result.addChild(argument);

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

@@ -31,7 +31,7 @@ public class RequiredArgumentBuilder<S, T> extends ArgumentBuilder<S, RequiredAr
     }
 
     public ArgumentCommandNode<S, T> build() {
-        final ArgumentCommandNode<S, T> result = new ArgumentCommandNode<>(getName(), getType(), getCommand(), getRequirement(), getRedirect());
+        final ArgumentCommandNode<S, T> result = new ArgumentCommandNode<>(getName(), getType(), getCommand(), getRequirement(), getRedirect(), getRedirectModifier());
 
         for (final CommandNode<S> argument : getArguments()) {
             result.addChild(argument);

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

@@ -13,19 +13,19 @@ public class CommandContext<S> {
     private final Map<String, ParsedArgument<S, ?>> arguments;
     private final Map<CommandNode<S>, String> nodes;
     private final String input;
-    private final CommandContext<S> parent;
+    private final CommandContext<S> child;
 
-    public CommandContext(final S source, final Map<String, ParsedArgument<S, ?>> arguments, final Command<S> command, final Map<CommandNode<S>, String> nodes, final String input, final CommandContext<S> parent) {
+    public CommandContext(final S source, final Map<String, ParsedArgument<S, ?>> arguments, final Command<S> command, final Map<CommandNode<S>, String> nodes, final String input, final CommandContext<S> child) {
         this.source = source;
         this.arguments = arguments;
         this.command = command;
         this.nodes = nodes;
         this.input = input;
-        this.parent = parent;
+        this.child = child;
     }
 
-    public CommandContext<S> getParent() {
+    public CommandContext<S> getChild() {
-        return parent;
+        return child;
     }
 
     public Command<S> getCommand() {
@@ -63,7 +63,7 @@ public class CommandContext<S> {
         if (!Iterables.elementsEqual(nodes.entrySet(), that.nodes.entrySet())) return false;
         if (command != null ? !command.equals(that.command) : that.command != null) return false;
         if (!source.equals(that.source)) return false;
-        if (parent != null ? !parent.equals(that.parent) : that.parent != null) return false;
+        if (child != null ? !child.equals(that.child) : that.child != null) return false;
 
         return true;
     }
@@ -74,7 +74,7 @@ public class CommandContext<S> {
         result = 31 * result + arguments.hashCode();
         result = 31 * result + (command != null ? command.hashCode() : 0);
         result = 31 * result + nodes.hashCode();
-        result = 31 * result + (parent != null ? parent.hashCode() : 0);
+        result = 31 * result + (child != null ? child.hashCode() : 0);
         return result;
     }
 

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

@@ -13,7 +13,7 @@ public class CommandContextBuilder<S> {
     private final CommandDispatcher<S> dispatcher;
     private S source;
     private Command<S> command;
-    private CommandContext<S> parent;
+    private CommandContextBuilder<S> child;
 
     public CommandContextBuilder(final CommandDispatcher<S> dispatcher, final S source) {
         this.dispatcher = dispatcher;
@@ -53,19 +53,21 @@ public class CommandContextBuilder<S> {
         copy.command = command;
         copy.arguments.putAll(arguments);
         copy.nodes.putAll(nodes);
-        copy.parent = parent;
+        copy.child = child;
         return copy;
     }
 
-    public CommandContextBuilder<S> redirect(final CommandNode<S> newRoot) {
+    public CommandContextBuilder<S> withChild(final CommandContextBuilder<S> child) {
-        final CommandContextBuilder<S> result = new CommandContextBuilder<>(dispatcher, source);
+        this.child = child;
-        result.withNode(newRoot, "");
+        return this;
-        result.parent = build();
-        return result;
     }
 
-    public CommandContext<S> getParent() {
+    public CommandContextBuilder<S> getChild() {
-        return parent;
+        return child;
+    }
+
+    public Command<S> getCommand() {
+        return command;
     }
 
     public String getInput() {
@@ -89,7 +91,7 @@ public class CommandContextBuilder<S> {
     }
 
     public CommandContext<S> build() {
-        return new CommandContext<>(source, arguments, command, nodes, getInput(), parent);
+        return new CommandContext<>(source, arguments, command, nodes, getInput(), child == null ? null : child.build());
     }
 
     public CommandDispatcher<S> getDispatcher() {

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

@@ -8,7 +8,9 @@ import com.mojang.brigadier.context.CommandContextBuilder;
 import com.mojang.brigadier.context.ParsedArgument;
 import com.mojang.brigadier.exceptions.CommandException;
 
+import java.util.Collection;
 import java.util.Set;
+import java.util.function.Function;
 import java.util.function.Predicate;
 
 public class ArgumentCommandNode<S, T> extends CommandNode<S> {
@@ -18,8 +20,8 @@ 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) {
+    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) {
-        super(command, requirement, redirect);
+        super(command, requirement, redirect, modifier);
         this.name = name;
         this.type = type;
     }
@@ -69,7 +71,7 @@ public class ArgumentCommandNode<S, T> extends CommandNode<S> {
     public RequiredArgumentBuilder<S, T> createBuilder() {
         final RequiredArgumentBuilder<S, T> builder = RequiredArgumentBuilder.argument(name, type);
         builder.requires(getRequirement());
-        builder.redirect(getRedirect());
+        builder.redirect(getRedirect(), getRedirectModifier());
         if (getCommand() != null) {
             builder.executes(getCommand());
         }

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

@@ -12,6 +12,7 @@ import java.util.Collection;
 import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
 
@@ -19,12 +20,14 @@ 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 Command<S> command;
 
-    protected CommandNode(final Command<S> command, final Predicate<S> requirement, final CommandNode<S> redirect) {
+    protected CommandNode(final Command<S> command, final Predicate<S> requirement, final CommandNode<S> redirect, final Function<S, Collection<S>> modifier) {
         this.command = command;
         this.requirement = requirement;
         this.redirect = redirect;
+        this.modifier = modifier;
     }
 
     public Command<S> getCommand() {
@@ -39,6 +42,10 @@ public abstract class CommandNode<S> implements Comparable<CommandNode<S>> {
         return redirect;
     }
 
+    public Function<S, Collection<S>> getRedirectModifier() {
+        return modifier;
+    }
+
     public boolean canUse(final S source) {
         return requirement.test(source);
     }

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

@@ -7,7 +7,9 @@ import com.mojang.brigadier.context.CommandContextBuilder;
 import com.mojang.brigadier.exceptions.CommandException;
 import com.mojang.brigadier.exceptions.ParameterizedCommandExceptionType;
 
+import java.util.Collection;
 import java.util.Set;
+import java.util.function.Function;
 import java.util.function.Predicate;
 
 public class LiteralCommandNode<S> extends CommandNode<S> {
@@ -15,8 +17,8 @@ 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) {
+    public LiteralCommandNode(final String literal, final Command<S> command, final Predicate<S> requirement, final CommandNode<S> redirect, final Function<S, Collection<S>> modifier) {
-        super(command, requirement, redirect);
+        super(command, requirement, redirect, modifier);
         this.literal = literal;
     }
 
@@ -78,7 +80,7 @@ public class LiteralCommandNode<S> extends CommandNode<S> {
     public LiteralArgumentBuilder<S> createBuilder() {
         final LiteralArgumentBuilder<S> builder = LiteralArgumentBuilder.literal(this.literal);
         builder.requires(getRequirement());
-        builder.redirect(getRedirect());
+        builder.redirect(getRedirect(), getRedirectModifier());
         if (getCommand() != null) {
             builder.executes(getCommand());
         }

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

@@ -5,11 +5,13 @@ import com.mojang.brigadier.builder.ArgumentBuilder;
 import com.mojang.brigadier.context.CommandContextBuilder;
 import com.mojang.brigadier.exceptions.CommandException;
 
+import java.util.Collection;
+import java.util.Collections;
 import java.util.Set;
 
 public class RootCommandNode<S> extends CommandNode<S> {
     public RootCommandNode() {
-        super(null, c -> true, null);
+        super(null, c -> true, null, Collections::singleton);
     }
 
     @Override

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

@@ -6,6 +6,8 @@ import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.runners.MockitoJUnitRunner;
 
+import java.util.Collections;
+
 import static com.mojang.brigadier.builder.LiteralArgumentBuilder.literal;
 import static org.hamcrest.Matchers.emptyArray;
 import static org.hamcrest.Matchers.equalTo;
@@ -43,7 +45,7 @@ public class CommandDispatcherCompletionsTest {
     public void testCommand_redirect() throws Exception {
         subject.register(literal("foo"));
         subject.register(literal("bar"));
-        subject.register(literal("redirect").redirect(subject.getRoot()));
+        subject.register(literal("redirect").redirect(subject.getRoot(), Collections::singleton));
         assertThat(subject.getCompletionSuggestions("redirect ", source), equalTo(new String[]{"bar", "foo", "redirect"}));
     }
 

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

@@ -1,6 +1,8 @@
 package com.mojang.brigadier;
 
+import com.google.common.collect.Lists;
 import com.mojang.brigadier.context.CommandContext;
+import com.mojang.brigadier.context.CommandContextBuilder;
 import com.mojang.brigadier.exceptions.CommandException;
 import com.mojang.brigadier.tree.LiteralCommandNode;
 import org.junit.Before;
@@ -9,17 +11,20 @@ import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.runners.MockitoJUnitRunner;
 
+import java.util.Collection;
 import java.util.Collections;
+import java.util.function.Function;
 
 import static com.mojang.brigadier.arguments.IntegerArgumentType.integer;
 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.hamcrest.Matchers.hasProperty;
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.notNullValue;
 import static org.junit.Assert.assertThat;
 import static org.junit.Assert.fail;
-import static org.mockito.Matchers.notNull;
+import static org.mockito.Matchers.argThat;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
@@ -172,26 +177,54 @@ public class CommandDispatcherTest {
     @Test
     public void testExecuteRedirected() throws Exception {
         subject.register(literal("actual").executes(command));
-        subject.register(literal("redirected").redirect(subject.getRoot()));
+        subject.register(literal("redirected").redirect(subject.getRoot(), Collections::singleton));
 
         final ParseResults<Object> parse = subject.parse("redirected redirected actual", source);
-        assertThat(parse.getContext().getInput(), equalTo("actual"));
+        assertThat(parse.getContext().getInput(), equalTo("redirected"));
-        assertThat(parse.getContext().getNodes().size(), is(2));
+        assertThat(parse.getContext().getNodes().size(), is(1));
 
-        final CommandContext<Object> parent1 = parse.getContext().getParent();
+        final CommandContextBuilder<Object> child1 = parse.getContext().getChild();
-        assertThat(parent1, is(notNullValue()));
+        assertThat(child1, is(notNullValue()));
-        assertThat(parent1.getInput(), equalTo("redirected"));
+        assertThat(child1.getInput(), equalTo("redirected"));
-        assertThat(parent1.getNodes().size(), is(2));
+        assertThat(child1.getNodes().size(), is(2));
 
-        final CommandContext<Object> parent2 = parent1.getParent();
+        final CommandContextBuilder<Object> child2 = child1.getChild();
-        assertThat(parent2, is(notNullValue()));
+        assertThat(child2, is(notNullValue()));
-        assertThat(parent2.getInput(), equalTo("redirected"));
+        assertThat(child2.getInput(), equalTo("actual"));
-        assertThat(parent2.getNodes().size(), is(1));
+        assertThat(child2.getNodes().size(), is(2));
 
         assertThat(subject.execute(parse), is(42));
         verify(command).run(any(CommandContext.class));
     }
 
+    @SuppressWarnings("unchecked")
+    @Test
+    public void testExecuteRedirectedMultipleTimes() throws Exception {
+        final Function<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));
+
+        subject.register(literal("actual").executes(command));
+        subject.register(literal("redirected").redirect(subject.getRoot(), modifier));
+
+        final ParseResults<Object> parse = subject.parse("redirected actual", source);
+        assertThat(parse.getContext().getInput(), equalTo("redirected"));
+        assertThat(parse.getContext().getNodes().size(), is(1));
+        assertThat(parse.getContext().getSource(), is(source));
+
+        final CommandContextBuilder<Object> parent = parse.getContext().getChild();
+        assertThat(parent, is(notNullValue()));
+        assertThat(parent.getInput(), equalTo("actual"));
+        assertThat(parent.getNodes().size(), is(2));
+        assertThat(parent.getSource(), is(source));
+
+        assertThat(subject.execute(parse), is(84));
+        verify(command).run(argThat(hasProperty("source", is(source1))));
+        verify(command).run(argThat(hasProperty("source", is(source2))));
+    }
+
     @Test
     public void testExecuteOrphanedSubcommand() throws Exception {
         subject.register(literal("foo").then(

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

@@ -10,6 +10,7 @@ import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.runners.MockitoJUnitRunner;
 
+import java.util.Collections;
 import java.util.Map;
 
 import static com.mojang.brigadier.builder.LiteralArgumentBuilder.literal;
@@ -89,11 +90,11 @@ public class CommandDispatcherUsagesTest {
         );
         subject.register(
             literal("j")
-                .redirect(subject.getRoot())
+                .redirect(subject.getRoot(), Collections::singleton)
         );
         subject.register(
             literal("k")
-                .redirect(get("h"))
+                .redirect(get("h"), Collections::singleton)
         );
     }
 

src/test/java/com/mojang/brigadier/builder/ArgumentBuilderTest.java

@@ -5,6 +5,8 @@ import org.junit.Before;
 import org.junit.Test;
 import org.mockito.Mockito;
 
+import java.util.Collections;
+
 import static com.mojang.brigadier.arguments.IntegerArgumentType.integer;
 import static com.mojang.brigadier.builder.LiteralArgumentBuilder.literal;
 import static com.mojang.brigadier.builder.RequiredArgumentBuilder.argument;
@@ -35,7 +37,7 @@ public class ArgumentBuilderTest {
     @Test
     public void testRedirect() throws Exception {
         final CommandNode<Object> target = mock(CommandNode.class);
-        builder.redirect(target);
+        builder.redirect(target, Collections::singleton);
         assertThat(builder.getRedirect(), is(target));
     }
 
@@ -43,13 +45,13 @@ public class ArgumentBuilderTest {
     public void testRedirect_withChild() throws Exception {
         final CommandNode<Object> target = mock(CommandNode.class);
         builder.then(literal("foo"));
-        builder.redirect(target);
+        builder.redirect(target, Collections::singleton);
     }
 
     @Test(expected = IllegalStateException.class)
     public void testThen_withRedirect() throws Exception {
         final CommandNode<Object> target = mock(CommandNode.class);
-        builder.redirect(target);
+        builder.redirect(target, Collections::singleton);
         builder.then(literal("foo"));
     }