Add utilities for running commands in stages

gradle.properties

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

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

@@ -6,6 +6,7 @@
 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.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 int execute(final ParseResults<S> parse) throws CommandSyntaxException {
             }
         }
 
-        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 int size = contexts.size();
-            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) {
+        final Optional<ContextChain<S>> flatContext = ContextChain.tryFlatten(original);
+        if (!flatContext.isPresent()) {
             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,
+    }
+}