Add utilities for running commands in stages (#140)

* Add more comprehensive test for context propagation in redirect

* Add tests for forking and exceptional execution

* Add test for redirect+execute nodes

* Cleanup some warnings from CommandDispatcherTest

* Remove leftover check (was needed when redirect was not cached)

* Refactor command execution loop and change "no command" behavior to not
depend on runtime

* Add utilities for running command chains manually

* Bump version

gradle.properties

@@ -1 +1 @@
-majorMinor: 1.1
+majorMinor: 1.2

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

@@ -6,6 +6,7 @@ package com.mojang.brigadier;
 import com.mojang.brigadier.builder.LiteralArgumentBuilder;
 import com.mojang.brigadier.context.CommandContext;
 import com.mojang.brigadier.context.CommandContextBuilder;
+import com.mojang.brigadier.context.ContextChain;
 import com.mojang.brigadier.context.SuggestionContext;
 import com.mojang.brigadier.exceptions.CommandSyntaxException;
 import com.mojang.brigadier.suggestion.Suggestions;
@@ -22,6 +23,7 @@ import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.function.Predicate;
@@ -214,76 +216,16 @@ public class CommandDispatcher<S> {
             }
         }
 
-        int result = 0;
-        int successfulForks = 0;
-        boolean forked = false;
-        boolean foundCommand = false;
         final String command = parse.getReader().getString();
         final CommandContext<S> original = parse.getContext().build(command);
-        List<CommandContext<S>> contexts = Collections.singletonList(original);
-        ArrayList<CommandContext<S>> next = null;
 
-        while (contexts != null) {
+        final Optional<ContextChain<S>> flatContext = ContextChain.tryFlatten(original);
-            final int size = contexts.size();
+        if (!flatContext.isPresent()) {
-            for (int i = 0; i < size; i++) {
-                final CommandContext<S> context = contexts.get(i);
-                final CommandContext<S> child = context.getChild();
-                if (child != null) {
-                    forked |= context.isForked();
-                    if (child.hasNodes()) {
-                        final RedirectModifier<S> modifier = context.getRedirectModifier();
-                        if (modifier == null) {
-                            if (next == null) {
-                                next = new ArrayList<>(1);
-                            }
-                            next.add(child.copyFor(context.getSource()));
-                        } else {
-                            try {
-                                final Collection<S> results = modifier.apply(context);
-                                if (!results.isEmpty()) {
-                                    if (next == null) {
-                                        next = new ArrayList<>(results.size());
-                                    }
-                                    for (final S source : results) {
-                                        next.add(child.copyFor(source));
-                                    }
-                                } else {
-                                    foundCommand = true;
-                                }
-                            } catch (final CommandSyntaxException ex) {
-                                consumer.onCommandComplete(context, false, 0);
-                                if (!forked) {
-                                    throw ex;
-                                }
-                            }
-                        }
-                    }
-                } else if (context.getCommand() != null) {
-                    foundCommand = true;
-                    try {
-                        final int value = context.getCommand().run(context);
-                        result += value;
-                        consumer.onCommandComplete(context, true, value);
-                        successfulForks++;
-                    } catch (final CommandSyntaxException ex) {
-                        consumer.onCommandComplete(context, false, 0);
-                        if (!forked) {
-                            throw ex;
-                        }
-                    }
-                }
-            }
-
-            contexts = next;
-            next = null;
-        }
-
-        if (!foundCommand) {
             consumer.onCommandComplete(original, false, 0);
             throw CommandSyntaxException.BUILT_IN_EXCEPTIONS.dispatcherUnknownCommand().createWithContext(parse.getReader());
         }
 
-        return forked ? successfulForks : result;
+        return flatContext.get().executeAll(original.getSource(), consumer);
     }
 
     /**

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

@@ -28,13 +28,30 @@ public class CommandContext<S> {
 
     private final S source;
     private final String input;
+    /**
+     * Executable part of command. Will be run only when context is last in chain.
+     */
     private final Command<S> command;
     private final Map<String, ParsedArgument<S, ?>> arguments;
     private final CommandNode<S> rootNode;
     private final List<ParsedCommandNode<S>> nodes;
     private final StringRange range;
     private final CommandContext<S> child;
+    /**
+     * Modifier of source. Will be run only when context has children (i.e. is not last in chain).
+     */
     private final RedirectModifier<S> modifier;
+    /**
+     * Special modifier for running this context and children.
+     * Only relevant if it's not last in chain.
+     * <br/>
+     *
+     * Effects:
+     * <ul>
+     *     <li>Exceptions from {@link #command} or {@link #modifier} will be ignored</li>
+     *     <li>Result of command will be number of elements run by element in chain (instead of sum of {@link #command} results</li>
+     * </ul>
+     */
     private final boolean forks;
 
     public CommandContext(final S source, final String input, final Map<String, ParsedArgument<S, ?>> arguments, final Command<S> command, final CommandNode<S> rootNode, final List<ParsedCommandNode<S>> nodes, final StringRange range, final CommandContext<S> child, final RedirectModifier<S> modifier, boolean forks) {

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

@@ -0,0 +1,142 @@
+package com.mojang.brigadier.context;
+
+import com.mojang.brigadier.RedirectModifier;
+import com.mojang.brigadier.ResultConsumer;
+import com.mojang.brigadier.exceptions.CommandSyntaxException;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.Optional;
+
+public class ContextChain<S> {
+    // TODO ideally those two would have separate types, but modifiers and executables expect full context
+    private final List<CommandContext<S>> modifiers;
+    private final CommandContext<S> executable;
+
+    private ContextChain<S> nextStageCache = null;
+
+    public ContextChain(final List<CommandContext<S>> modifiers, final CommandContext<S> executable) {
+        if (executable.getCommand() == null) {
+            throw new IllegalArgumentException("Last command in chain must be executable");
+        }
+        this.modifiers = modifiers;
+        this.executable = executable;
+    }
+
+    public static <S> Optional<ContextChain<S>> tryFlatten(final CommandContext<S> rootContext) {
+        final List<CommandContext<S>> modifiers = new ArrayList<>();
+
+        CommandContext<S> current = rootContext;
+
+        while (true) {
+            final CommandContext<S> child = current.getChild();
+            if (child == null) {
+                // Last entry must be executable command
+                if (current.getCommand() == null) {
+                    return Optional.empty();
+                }
+
+                return Optional.of(new ContextChain<>(modifiers, current));
+            }
+
+            modifiers.add(current);
+            current = child;
+        }
+    }
+
+    public static <S> Collection<S> runModifier(final CommandContext<S> modifier, final S source, final ResultConsumer<S> resultConsumer, final boolean forkedMode) throws CommandSyntaxException {
+        final RedirectModifier<S> sourceModifier = modifier.getRedirectModifier();
+
+        // Note: source currently in context is irrelevant at this point, since we might have updated it in one of earlier stages
+        if (sourceModifier == null) {
+            // Simple redirect, just propagate source to next node
+            return Collections.singleton(source);
+        }
+
+        final CommandContext<S> contextToUse = modifier.copyFor(source);
+        try {
+            return sourceModifier.apply(contextToUse);
+        } catch (final CommandSyntaxException ex) {
+            resultConsumer.onCommandComplete(contextToUse, false, 0);
+            if (forkedMode) {
+                return Collections.emptyList();
+            }
+            throw ex;
+        }
+    }
+
+    public static <S> int runExecutable(final CommandContext<S> executable, final S source, final ResultConsumer<S> resultConsumer, final boolean forkedMode) throws CommandSyntaxException {
+        final CommandContext<S> contextToUse = executable.copyFor(source);
+        try {
+            final int result = executable.getCommand().run(contextToUse);
+            resultConsumer.onCommandComplete(contextToUse, true, result);
+            return forkedMode ? 1 : result;
+        } catch (final CommandSyntaxException ex) {
+            resultConsumer.onCommandComplete(contextToUse, false, 0);
+            if (forkedMode) {
+                return 0;
+            }
+            throw ex;
+        }
+    }
+
+    public int executeAll(final S source, final ResultConsumer<S> resultConsumer) throws CommandSyntaxException {
+        if (modifiers.isEmpty()) {
+            // Fast path - just a single stage
+            return runExecutable(executable, source, resultConsumer, false);
+        }
+
+        boolean forkedMode = false;
+        List<S> currentSources = Collections.singletonList(source);
+
+        for (final CommandContext<S> modifier : modifiers) {
+            forkedMode |= modifier.isForked();
+
+            List<S> nextSources = new ArrayList<>();
+            for (final S sourceToRun : currentSources) {
+                nextSources.addAll(runModifier(modifier, sourceToRun, resultConsumer, forkedMode));
+            }
+            if (nextSources.isEmpty()) {
+                return 0;
+            }
+            currentSources = nextSources;
+        }
+
+        int result = 0;
+        for (final S executionSource : currentSources) {
+            result += runExecutable(executable, executionSource, resultConsumer, forkedMode);
+        }
+
+        return result;
+    }
+
+    public Stage getStage() {
+        return modifiers.isEmpty() ? Stage.EXECUTE : Stage.MODIFY;
+    }
+
+    public CommandContext<S> getTopContext() {
+        if (modifiers.isEmpty()) {
+            return executable;
+        }
+        return modifiers.get(0);
+    }
+
+    public ContextChain<S> nextStage() {
+        final int modifierCount = modifiers.size();
+        if (modifierCount == 0) {
+            return null;
+        }
+
+        if (nextStageCache == null) {
+            nextStageCache = new ContextChain<>(modifiers.subList(1, modifierCount), executable);
+        }
+        return nextStageCache;
+    }
+
+    public enum Stage {
+        MODIFY,
+        EXECUTE,
+    }
+}

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

@@ -5,19 +5,25 @@ package com.mojang.brigadier;
 
 import com.google.common.collect.Lists;
 import com.mojang.brigadier.arguments.IntegerArgumentType;
+import com.mojang.brigadier.builder.LiteralArgumentBuilder;
+import com.mojang.brigadier.builder.RequiredArgumentBuilder;
 import com.mojang.brigadier.context.CommandContext;
 import com.mojang.brigadier.context.CommandContextBuilder;
 import com.mojang.brigadier.exceptions.CommandSyntaxException;
 import com.mojang.brigadier.tree.LiteralCommandNode;
+import com.mojang.brigadier.tree.RootCommandNode;
+import org.hamcrest.CustomMatcher;
+import org.hamcrest.Matcher;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.runners.MockitoJUnitRunner;
 
+import java.util.Arrays;
 import java.util.Collections;
-import java.util.concurrent.atomic.AtomicBoolean;
 
+import static com.mojang.brigadier.arguments.IntegerArgumentType.getInteger;
 import static com.mojang.brigadier.arguments.IntegerArgumentType.integer;
 import static com.mojang.brigadier.builder.LiteralArgumentBuilder.literal;
 import static com.mojang.brigadier.builder.RequiredArgumentBuilder.argument;
@@ -29,11 +35,14 @@ import static org.hamcrest.Matchers.nullValue;
 import static org.junit.Assert.assertThat;
 import static org.junit.Assert.fail;
 import static org.mockito.Matchers.argThat;
+import static org.mockito.Matchers.eq;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+import static org.mockito.Mockito.verifyZeroInteractions;
 import static org.mockito.Mockito.when;
 
 @RunWith(MockitoJUnitRunner.class)
@@ -43,6 +52,8 @@ public class CommandDispatcherTest {
     private Command<Object> command;
     @Mock
     private Object source;
+    @Mock
+    private ResultConsumer<Object> consumer;
 
     @Before
     public void setUp() throws Exception {
@@ -86,7 +97,7 @@ public class CommandDispatcherTest {
     }
 
     @Test
-    public void testExecuteUnknownCommand() throws Exception {
+    public void testExecuteUnknownCommand() {
         subject.register(literal("bar"));
         subject.register(literal("baz"));
 
@@ -100,7 +111,7 @@ public class CommandDispatcherTest {
     }
 
     @Test
-    public void testExecuteImpermissibleCommand() throws Exception {
+    public void testExecuteImpermissibleCommand() {
         subject.register(literal("foo").requires(s -> false));
 
         try {
@@ -113,7 +124,7 @@ public class CommandDispatcherTest {
     }
 
     @Test
-    public void testExecuteEmptyCommand() throws Exception {
+    public void testExecuteEmptyCommand() {
         subject.register(literal(""));
 
         try {
@@ -126,7 +137,7 @@ public class CommandDispatcherTest {
     }
 
     @Test
-    public void testExecuteUnknownSubcommand() throws Exception {
+    public void testExecuteUnknownSubcommand() {
         subject.register(literal("foo").executes(command));
 
         try {
@@ -139,7 +150,7 @@ public class CommandDispatcherTest {
     }
 
     @Test
-    public void testExecuteIncorrectLiteral() throws Exception {
+    public void testExecuteIncorrectLiteral() {
         subject.register(literal("foo").executes(command).then(literal("bar")));
 
         try {
@@ -152,7 +163,7 @@ public class CommandDispatcherTest {
     }
 
     @Test
-    public void testExecuteAmbiguousIncorrectArgument() throws Exception {
+    public void testExecuteAmbiguousIncorrectArgument() {
         subject.register(
             literal("foo").executes(command)
                 .then(literal("bar"))
@@ -186,9 +197,8 @@ public class CommandDispatcherTest {
         verify(subCommand).run(any(CommandContext.class));
     }
 
-    @SuppressWarnings("unchecked")
     @Test
-    public void testParseIncompleteLiteral() throws Exception {
+    public void testParseIncompleteLiteral() {
         subject.register(literal("foo").then(literal("bar").executes(command)));
 
         final ParseResults<Object> parse = subject.parse("foo ", source);
@@ -196,9 +206,8 @@ public class CommandDispatcherTest {
         assertThat(parse.getContext().getNodes().size(), is(1));
     }
 
-    @SuppressWarnings("unchecked")
     @Test
-    public void testParseIncompleteArgument() throws Exception {
+    public void testParseIncompleteArgument() {
         subject.register(literal("foo").then(argument("bar", integer()).executes(command)));
 
         final ParseResults<Object> parse = subject.parse("foo ", source);
@@ -295,6 +304,49 @@ public class CommandDispatcherTest {
         verify(command).run(any(CommandContext.class));
     }
 
+    @Test
+    public void testCorrectExecuteContextAfterRedirect() throws Exception {
+        final CommandDispatcher<Integer> subject = new CommandDispatcher<>();
+
+        final RootCommandNode<Integer> root = subject.getRoot();
+        final LiteralArgumentBuilder<Integer> add = literal("add");
+        final LiteralArgumentBuilder<Integer> blank = literal("blank");
+        final RequiredArgumentBuilder<Integer, Integer> addArg = argument("value", integer());
+        final LiteralArgumentBuilder<Integer> run = literal("run");
+
+        subject.register(add.then(addArg.redirect(root, c -> c.getSource() + getInteger(c, "value"))));
+        subject.register(blank.redirect(root));
+        subject.register(run.executes(CommandContext::getSource));
+
+        assertThat(subject.execute("run", 0), is(0));
+        assertThat(subject.execute("run", 1), is(1));
+
+        assertThat(subject.execute("add 5 run", 1), is(1 + 5));
+        assertThat(subject.execute("add 5 add 6 run", 2), is(2 + 5 + 6));
+        assertThat(subject.execute("add 5 blank run", 1), is(1 + 5));
+        assertThat(subject.execute("blank add 5 run", 1), is(1 + 5));
+        assertThat(subject.execute("add 5 blank add 6 run", 2), is(2 + 5 + 6));
+        assertThat(subject.execute("add 5 blank blank add 6 run", 2), is(2 + 5 + 6));
+    }
+
+    @Test
+    public void testSharedRedirectAndExecuteNodes() throws CommandSyntaxException {
+        final CommandDispatcher<Integer> subject = new CommandDispatcher<>();
+
+        final RootCommandNode<Integer> root = subject.getRoot();
+        final LiteralArgumentBuilder<Integer> add = literal("add");
+        final RequiredArgumentBuilder<Integer, Integer> addArg = argument("value", integer());
+
+        subject.register(add.then(
+            addArg
+                .redirect(root, c -> c.getSource() + getInteger(c, "value"))
+                .executes(CommandContext::getSource)
+        ));
+
+        assertThat(subject.execute("add 5", 1), is(1));
+        assertThat(subject.execute("add 5 add 6", 1), is(1 + 5));
+    }
+
     @SuppressWarnings("unchecked")
     @Test
     public void testExecuteRedirected() throws Exception {
@@ -336,9 +388,9 @@ public class CommandDispatcherTest {
             .then(literal("bar")
                 .then(argument("value", integer()).executes(context -> IntegerArgumentType.getInteger(context, "value"))))
             .then(literal("awa").executes(context -> 2)));
-        final LiteralCommandNode<Object> baz = subject.register(literal("baz").redirect(foo));
+        subject.register(literal("baz").redirect(foo));
         try {
-            int result = subject.execute("baz bar", source);
+            subject.execute("baz bar", source);
             fail("Should have thrown an exception");
         } catch (CommandSyntaxException e) {
             assertThat(e.getType(), is(CommandSyntaxException.BUILT_IN_EXCEPTIONS.dispatcherUnknownCommand()));
@@ -346,23 +398,19 @@ public class CommandDispatcherTest {
     }
 
     @Test
-    public void testRedirectModifierEmptyResult() {
+    public void testRedirectModifierEmptyResult() throws CommandSyntaxException {
         final LiteralCommandNode<Object> foo = subject.register(literal("foo")
             .then(literal("bar")
                 .then(argument("value", integer()).executes(context -> IntegerArgumentType.getInteger(context, "value"))))
             .then(literal("awa").executes(context -> 2)));
         final RedirectModifier<Object> emptyModifier = context -> Collections.emptyList();
-        final LiteralCommandNode<Object> baz = subject.register(literal("baz").fork(foo, emptyModifier));
+        subject.register(literal("baz").fork(foo, emptyModifier));
-        try {
         int result = subject.execute("baz bar 100", source);
         assertThat(result, is(0)); // No commands executed, so result is 0
-        } catch (CommandSyntaxException e) {
-            fail("Should not throw an exception");
-        }
     }
 
     @Test
-    public void testExecuteOrphanedSubcommand() throws Exception {
+    public void testExecuteOrphanedSubcommand() {
         subject.register(literal("foo").then(
             argument("bar", integer())
         ).executes(command));
@@ -376,6 +424,7 @@ public class CommandDispatcherTest {
         }
     }
 
+    @SuppressWarnings("unchecked")
     @Test
     public void testExecute_invalidOther() throws Exception {
         final Command<Object> wrongCommand = mock(Command.class);
@@ -388,7 +437,7 @@ public class CommandDispatcherTest {
     }
 
     @Test
-    public void parse_noSpaceSeparator() throws Exception {
+    public void parse_noSpaceSeparator() {
         subject.register(literal("foo").then(argument("bar", integer()).executes(command)));
 
         try {
@@ -401,7 +450,7 @@ public class CommandDispatcherTest {
     }
 
     @Test
-    public void testExecuteInvalidSubcommand() throws Exception {
+    public void testExecuteInvalidSubcommand() {
         subject.register(literal("foo").then(
             argument("bar", integer())
         ).executes(command));
@@ -435,4 +484,168 @@ public class CommandDispatcherTest {
     public void testFindNodeDoesntExist() {
         assertThat(subject.findNode(Lists.newArrayList("foo", "bar")), is(nullValue()));
     }
+
+    @Test
+    public void testResultConsumerInNonErrorRun() throws CommandSyntaxException {
+        subject.setConsumer(consumer);
+
+        subject.register(literal("foo").executes(command));
+        when(command.run(any())).thenReturn(5);
+
+        assertThat(subject.execute("foo", source), is(5));
+        verify(consumer).onCommandComplete(any(), eq(true), eq(5));
+        verifyNoMoreInteractions(consumer);
+    }
+
+    @Test
+    public void testResultConsumerInForkedNonErrorRun() throws CommandSyntaxException {
+        subject.setConsumer(consumer);
+
+        subject.register(literal("foo").executes(c -> (Integer)(c.getSource())));
+        final Object[] contexts = new Object[] {9, 10, 11};
+
+        subject.register(literal("repeat").fork(subject.getRoot(), context -> Arrays.asList(contexts)));
+
+        assertThat(subject.execute("repeat foo", source), is(contexts.length));
+        verify(consumer).onCommandComplete(argThat(contextSourceMatches(contexts[0])), eq(true), eq(9));
+        verify(consumer).onCommandComplete(argThat(contextSourceMatches(contexts[1])), eq(true), eq(10));
+        verify(consumer).onCommandComplete(argThat(contextSourceMatches(contexts[2])), eq(true), eq(11));
+        verifyNoMoreInteractions(consumer);
+    }
+
+    @Test
+    public void testExceptionInNonForkedCommand() throws CommandSyntaxException {
+        subject.setConsumer(consumer);
+        subject.register(literal("crash").executes(command));
+        final CommandSyntaxException exception = CommandSyntaxException.BUILT_IN_EXCEPTIONS.readerExpectedBool().create();
+        when(command.run(any())).thenThrow(exception);
+
+        try {
+            subject.execute("crash", source);
+            fail();
+        } catch (final CommandSyntaxException ex) {
+            assertThat(ex, is(exception));
+        }
+
+        verify(consumer).onCommandComplete(any(), eq(false), eq(0));
+        verifyNoMoreInteractions(consumer);
+    }
+
+    @Test
+    public void testExceptionInNonForkedRedirectedCommand() throws CommandSyntaxException {
+        subject.setConsumer(consumer);
+        subject.register(literal("crash").executes(command));
+        subject.register(literal("redirect").redirect(subject.getRoot()));
+
+        final CommandSyntaxException exception = CommandSyntaxException.BUILT_IN_EXCEPTIONS.readerExpectedBool().create();
+        when(command.run(any())).thenThrow(exception);
+
+        try {
+            subject.execute("redirect crash", source);
+            fail();
+        } catch (final CommandSyntaxException ex) {
+            assertThat(ex, is(exception));
+        }
+
+        verify(consumer).onCommandComplete(any(), eq(false), eq(0));
+        verifyNoMoreInteractions(consumer);
+    }
+
+    @Test
+    public void testExceptionInForkedRedirectedCommand() throws CommandSyntaxException {
+        subject.setConsumer(consumer);
+        subject.register(literal("crash").executes(command));
+        subject.register(literal("redirect").fork(subject.getRoot(), Collections::singleton));
+
+        final CommandSyntaxException exception = CommandSyntaxException.BUILT_IN_EXCEPTIONS.readerExpectedBool().create();
+        when(command.run(any())).thenThrow(exception);
+
+        assertThat(subject.execute("redirect crash", source), is(0));
+        verify(consumer).onCommandComplete(any(), eq(false), eq(0));
+        verifyNoMoreInteractions(consumer);
+    }
+
+    @Test
+    public void testExceptionInNonForkedRedirect() throws CommandSyntaxException {
+        final CommandSyntaxException exception = CommandSyntaxException.BUILT_IN_EXCEPTIONS.readerExpectedBool().create();
+
+        subject.setConsumer(consumer);
+        subject.register(literal("noop").executes(command));
+        subject.register(literal("redirect").redirect(subject.getRoot(), context -> {
+            throw exception;
+        }));
+
+        when(command.run(any())).thenReturn(3);
+
+        try {
+            subject.execute("redirect noop", source);
+            fail();
+        } catch (final CommandSyntaxException ex) {
+            assertThat(ex, is(exception));
+        }
+
+        verifyZeroInteractions(command);
+        verify(consumer).onCommandComplete(any(), eq(false), eq(0));
+        verifyNoMoreInteractions(consumer);
+    }
+
+    @Test
+    public void testExceptionInForkedRedirect() throws CommandSyntaxException {
+        final CommandSyntaxException exception = CommandSyntaxException.BUILT_IN_EXCEPTIONS.readerExpectedBool().create();
+
+        subject.setConsumer(consumer);
+        subject.register(literal("noop").executes(command));
+        subject.register(literal("redirect").fork(subject.getRoot(), context -> {
+            throw exception;
+        }));
+
+        when(command.run(any())).thenReturn(3);
+
+
+        assertThat(subject.execute("redirect noop", source), is(0));
+
+        verifyZeroInteractions(command);
+        verify(consumer).onCommandComplete(any(), eq(false), eq(0));
+        verifyNoMoreInteractions(consumer);
+    }
+
+    @Test
+    public void testPartialExceptionInForkedRedirect() throws CommandSyntaxException {
+        final CommandSyntaxException exception = CommandSyntaxException.BUILT_IN_EXCEPTIONS.readerExpectedBool().create();
+        final Object otherSource = new Object();
+        final Object rejectedSource = new Object();
+
+        subject.setConsumer(consumer);
+        subject.register(literal("run").executes(command));
+        subject.register(literal("split").fork(subject.getRoot(), context -> Arrays.asList(source, rejectedSource, otherSource)));
+        subject.register(literal("filter").fork(subject.getRoot(), context -> {
+            final Object currentSource = context.getSource();
+            if (currentSource == rejectedSource) {
+                throw exception;
+            }
+            return Collections.singleton(currentSource);
+        }));
+
+        when(command.run(any())).thenReturn(3);
+
+        assertThat(subject.execute("split filter run", source), is(2));
+
+        verify(command).run(argThat(contextSourceMatches(source)));
+        verify(command).run(argThat(contextSourceMatches(otherSource)));
+        verifyNoMoreInteractions(command);
+
+        verify(consumer).onCommandComplete(argThat(contextSourceMatches(rejectedSource)), eq(false), eq(0));
+        verify(consumer).onCommandComplete(argThat(contextSourceMatches(source)), eq(true), eq(3));
+        verify(consumer).onCommandComplete(argThat(contextSourceMatches(otherSource)), eq(true), eq(3));
+        verifyNoMoreInteractions(consumer);
+    }
+
+    public static Matcher<CommandContext<Object>> contextSourceMatches(final Object source) {
+        return new CustomMatcher<CommandContext<Object>>("source " + source) {
+            @Override
+            public boolean matches(Object object) {
+                return (object instanceof CommandContext) && ((CommandContext<?>) object).getSource() == source;
+            }
+        };
+    }
 }

src/test/java/com/mojang/brigadier/context/ContextChainTest.java

@@ -0,0 +1,137 @@
+package com.mojang.brigadier.context;
+
+import com.mojang.brigadier.Command;
+import com.mojang.brigadier.CommandDispatcher;
+import com.mojang.brigadier.CommandDispatcherTest;
+import com.mojang.brigadier.ParseResults;
+import com.mojang.brigadier.ResultConsumer;
+import com.mojang.brigadier.exceptions.CommandSyntaxException;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.runners.MockitoJUnitRunner;
+
+import java.util.Optional;
+
+import static com.mojang.brigadier.builder.LiteralArgumentBuilder.literal;
+import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.notNullValue;
+import static org.hamcrest.Matchers.nullValue;
+import static org.junit.Assert.assertThat;
+import static org.mockito.Matchers.any;
+import static org.mockito.Matchers.argThat;
+import static org.mockito.Matchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+import static org.mockito.Mockito.when;
+
+@RunWith(MockitoJUnitRunner.class)
+public class ContextChainTest {
+
+    @Test
+    @SuppressWarnings("unchecked")
+    public void testExecuteAllForSingleCommand() throws CommandSyntaxException {
+        final ResultConsumer<Object> consumer = mock(ResultConsumer.class);
+        final Command<Object> command = mock(Command.class);
+
+        when(command.run(any())).thenReturn(4);
+
+        final CommandDispatcher<Object> dispatcher = new CommandDispatcher<>();
+        dispatcher.register(literal("foo").executes(command));
+        final Object source = "compile_source";
+
+        final ParseResults<Object> result = dispatcher.parse("foo", source);
+        final CommandContext<Object> topContext = result.getContext().build("foo");
+        final ContextChain<Object> chain = ContextChain.tryFlatten(topContext).orElseThrow(AssertionError::new);
+
+        final Object runtimeSource = "runtime_source";
+        assertThat(chain.executeAll(runtimeSource, consumer), is(4));
+
+        verify(command).run(argThat(CommandDispatcherTest.contextSourceMatches(runtimeSource)));
+
+        verify(consumer).onCommandComplete(argThat(CommandDispatcherTest.contextSourceMatches(runtimeSource)), eq(true), eq(4));
+        verifyNoMoreInteractions(consumer);
+    }
+
+    @Test
+    @SuppressWarnings("unchecked")
+    public void testExecuteAllForRedirectedCommand() throws CommandSyntaxException {
+        final ResultConsumer<Object> consumer = mock(ResultConsumer.class);
+        final Command<Object> command = mock(Command.class);
+
+        when(command.run(any())).thenReturn(4);
+
+        final Object redirectedSource = "redirected_source";
+
+        final CommandDispatcher<Object> dispatcher = new CommandDispatcher<>();
+        dispatcher.register(literal("foo").executes(command));
+        dispatcher.register(literal("bar").redirect(dispatcher.getRoot(), context -> redirectedSource));
+        final Object source = "compile_source";
+
+        final ParseResults<Object> result = dispatcher.parse("bar foo", source);
+        final CommandContext<Object> topContext = result.getContext().build("bar foo");
+        final ContextChain<Object> chain = ContextChain.tryFlatten(topContext).orElseThrow(AssertionError::new);
+
+        final Object runtimeSource = "runtime_source";
+        assertThat(chain.executeAll(runtimeSource, consumer), is(4));
+
+        verify(command).run(argThat(CommandDispatcherTest.contextSourceMatches(redirectedSource)));
+
+        verify(consumer).onCommandComplete(argThat(CommandDispatcherTest.contextSourceMatches(redirectedSource)), eq(true), eq(4));
+        verifyNoMoreInteractions(consumer);
+    }
+
+    @Test
+    public void testSingleStageExecution() {
+        final CommandDispatcher<Object> dispatcher = new CommandDispatcher<>();
+        dispatcher.register(literal("foo").executes(context -> 1));
+        final Object source = new Object();
+
+        final ParseResults<Object> result = dispatcher.parse("foo", source);
+        final CommandContext<Object> topContext = result.getContext().build("foo");
+        final ContextChain<Object> chain = ContextChain.tryFlatten(topContext).orElseThrow(AssertionError::new);
+
+        assertThat(chain.getStage(), is(ContextChain.Stage.EXECUTE));
+        assertThat(chain.getTopContext(), is(topContext));
+        assertThat(chain.nextStage(), nullValue());
+    }
+
+    @Test
+    public void testMultiStageExecution() {
+        final CommandDispatcher<Object> dispatcher = new CommandDispatcher<>();
+        dispatcher.register(literal("foo").executes(context -> 1));
+        dispatcher.register(literal("bar").redirect(dispatcher.getRoot()));
+        final Object source = new Object();
+
+        final ParseResults<Object> result = dispatcher.parse("bar bar foo", source);
+        final CommandContext<Object> topContext = result.getContext().build("bar bar foo");
+        final ContextChain<Object> stage0 = ContextChain.tryFlatten(topContext).orElseThrow(AssertionError::new);
+
+        assertThat(stage0.getStage(), is(ContextChain.Stage.MODIFY));
+        assertThat(stage0.getTopContext(), is(topContext));
+
+        final ContextChain<Object> stage1 = stage0.nextStage();
+        assertThat(stage1, notNullValue());
+        assertThat(stage1.getStage(), is(ContextChain.Stage.MODIFY));
+        assertThat(stage1.getTopContext(), is(topContext.getChild()));
+
+        final ContextChain<Object> stage2 = stage1.nextStage();
+        assertThat(stage2, notNullValue());
+        assertThat(stage2.getStage(), is(ContextChain.Stage.EXECUTE));
+        assertThat(stage2.getTopContext(), is(topContext.getChild().getChild()));
+
+        assertThat(stage2.nextStage(), nullValue());
+    }
+
+    @Test
+    public void testMissingExecute() {
+        final CommandDispatcher<Object> dispatcher = new CommandDispatcher<>();
+        dispatcher.register(literal("foo").executes(context -> 1));
+        dispatcher.register(literal("bar").redirect(dispatcher.getRoot()));
+
+        final Object source = new Object();
+        final ParseResults<Object> result = dispatcher.parse("bar bar", source);
+        final CommandContext<Object> topContext = result.getContext().build("bar bar");
+        assertThat(ContextChain.tryFlatten(topContext), is(Optional.empty()));
+    }
+}