Improved the speed of redirected commands

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

@@ -1,7 +1,6 @@
 package com.mojang.brigadier;
 
 import com.google.common.collect.Iterables;
-import com.google.common.collect.Iterators;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
@@ -92,29 +91,33 @@ public class CommandDispatcher<S> {
         int successfulForks = 0;
         boolean forked = false;
         boolean foundCommand = false;
-        final Deque<CommandContextBuilder<S>> contexts = new ArrayDeque<>();
+        final Deque<CommandContext<S>> contexts = new ArrayDeque<>(8);
-        contexts.add(parse.getContext());
+        final String command = parse.getReader().getString();
+        contexts.add(parse.getContext().build(command));
 
         while (!contexts.isEmpty()) {
-            final CommandContextBuilder<S> builder = contexts.removeFirst();
+            final CommandContext<S> context = contexts.removeFirst();
-            final CommandContextBuilder<S> child = builder.getChild();
+            final CommandContext<S> child = context.getChild();
-            final CommandContext<S> context = builder.build(parse.getReader().getString());
             if (child != null) {
                 if (!child.getNodes().isEmpty()) {
                     foundCommand = true;
-                    final RedirectModifier<S> modifier = Iterators.getLast(builder.getNodes().keySet().iterator()).getRedirectModifier();
+                    final RedirectModifier<S> modifier = context.getRedirectModifier();
-                    final Collection<S> results = modifier.apply(context);
+                    if (modifier == null) {
-                    if (results.size() > 1) {
+                        contexts.add(child);
-                        forked = true;
+                    } else {
-                    }
+                        final Collection<S> results = modifier.apply(context);
-                    for (final S source : results) {
+                        if (results.size() > 1) {
-                        contexts.add(child.copy().withSource(source));
+                            forked = true;
+                        }
+                        for (final S source : results) {
+                            contexts.add(child.copyFor(source));
+                        }
                     }
                 }
-            } else if (builder.getCommand() != null) {
+            } else if (context.getCommand() != null) {
                 foundCommand = true;
                 try {
-                    final int value = builder.getCommand().run(context);
+                    final int value = context.getCommand().run(context);
                     result += value;
                     consumer.onCommandComplete(context, true, value);
                     successfulForks++;
@@ -128,7 +131,7 @@ public class CommandDispatcher<S> {
         }
 
         if (!foundCommand) {
-            consumer.onCommandComplete(parse.getContext().build(parse.getReader().getString()), false, 0);
+            consumer.onCommandComplete(parse.getContext().build(command), false, 0);
             throw ERROR_UNKNOWN_COMMAND.createWithContext(parse.getReader());
         }
 

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

@@ -6,7 +6,6 @@ import com.mojang.brigadier.tree.CommandNode;
 import com.mojang.brigadier.tree.RootCommandNode;
 
 import java.util.Collection;
-import java.util.Collections;
 import java.util.function.Predicate;
 
 public abstract class ArgumentBuilder<S, T extends ArgumentBuilder<S, T>> {
@@ -14,7 +13,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 RedirectModifier<S> modifier = s -> Collections.singleton(s.getSource());
+    private RedirectModifier<S> modifier = null;
 
     protected abstract T getThis();
 
@@ -57,7 +56,7 @@ public abstract class ArgumentBuilder<S, T extends ArgumentBuilder<S, T>> {
     }
 
     public T redirect(final CommandNode<S> target) {
-        return redirect(target, modifier);
+        return redirect(target, null);
     }
 
     public T redirect(final CommandNode<S> target, final RedirectModifier<S> modifier) {

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

@@ -3,6 +3,7 @@ package com.mojang.brigadier.context;
 import com.google.common.collect.Iterables;
 import com.google.common.primitives.Primitives;
 import com.mojang.brigadier.Command;
+import com.mojang.brigadier.RedirectModifier;
 import com.mojang.brigadier.tree.CommandNode;
 
 import java.util.Map;
@@ -15,8 +16,9 @@ public class CommandContext<S> {
     private final Map<CommandNode<S>, StringRange> nodes;
     private final StringRange range;
     private final CommandContext<S> child;
+    private final RedirectModifier<S> modifier;
 
-    public CommandContext(final S source, final String input, final Map<String, ParsedArgument<S, ?>> arguments, final Command<S> command, final Map<CommandNode<S>, StringRange> nodes, final StringRange range, final CommandContext<S> child) {
+    public CommandContext(final S source, final String input, final Map<String, ParsedArgument<S, ?>> arguments, final Command<S> command, final Map<CommandNode<S>, StringRange> nodes, final StringRange range, final CommandContext<S> child, final RedirectModifier<S> modifier) {
         this.source = source;
         this.input = input;
         this.arguments = arguments;
@@ -24,6 +26,11 @@ public class CommandContext<S> {
         this.nodes = nodes;
         this.range = range;
         this.child = child;
+        this.modifier = modifier;
+    }
+
+    public CommandContext<S> copyFor(final S source) {
+        return new CommandContext<>(source, input, arguments, command, nodes, range, child, modifier);
     }
 
     public CommandContext<S> getChild() {
@@ -88,6 +95,10 @@ public class CommandContext<S> {
         return result;
     }
 
+    public RedirectModifier<S> getRedirectModifier() {
+        return modifier;
+    }
+
     public StringRange getRange() {
         return range;
     }

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

@@ -3,6 +3,7 @@ package com.mojang.brigadier.context;
 import com.google.common.collect.Maps;
 import com.mojang.brigadier.Command;
 import com.mojang.brigadier.CommandDispatcher;
+import com.mojang.brigadier.RedirectModifier;
 import com.mojang.brigadier.tree.CommandNode;
 
 import java.util.Map;
@@ -15,6 +16,7 @@ public class CommandContextBuilder<S> {
     private Command<S> command;
     private CommandContextBuilder<S> child;
     private StringRange range;
+    private RedirectModifier<S> modifier = null;
 
     public CommandContextBuilder(final CommandDispatcher<S> dispatcher, final S source, final int start) {
         this.dispatcher = dispatcher;
@@ -48,6 +50,7 @@ public class CommandContextBuilder<S> {
     public CommandContextBuilder<S> withNode(final CommandNode<S> node, final StringRange range) {
         nodes.put(node, range);
         this.range = StringRange.encompassing(this.range, range);
+        this.modifier = node.getRedirectModifier();
         return this;
     }
 
@@ -87,7 +90,7 @@ public class CommandContextBuilder<S> {
     }
 
     public CommandContext<S> build(final String input) {
-        return new CommandContext<>(source, input, arguments, command, nodes, range, child == null ? null : child.build(input));
+        return new CommandContext<>(source, input, arguments, command, nodes, range, child == null ? null : child.build(input), modifier);
     }
 
     public CommandDispatcher<S> getDispatcher() {

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

@@ -250,7 +250,7 @@ public class CommandDispatcherTest {
     @Test
     public void testExecuteRedirectedMultipleTimes() throws Exception {
         subject.register(literal("actual").executes(command));
-        subject.register(literal("redirected").redirect(subject.getRoot(), Collections::singleton));
+        subject.register(literal("redirected").redirect(subject.getRoot()));
 
         final String input = "redirected redirected actual";
         final ParseResults<Object> parse = subject.parse(input, source);

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

@@ -90,11 +90,11 @@ public class CommandDispatcherUsagesTest {
         );
         subject.register(
             literal("j")
-                .redirect(subject.getRoot(), Collections::singleton)
+                .redirect(subject.getRoot())
         );
         subject.register(
             literal("k")
-                .redirect(get("h"), Collections::singleton)
+                .redirect(get("h"))
         );
     }
 

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

@@ -3,9 +3,6 @@ package com.mojang.brigadier.builder;
 import com.mojang.brigadier.tree.CommandNode;
 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;
@@ -37,7 +34,7 @@ public class ArgumentBuilderTest {
     @Test
     public void testRedirect() throws Exception {
         final CommandNode<Object> target = mock(CommandNode.class);
-        builder.redirect(target, Collections::singleton);
+        builder.redirect(target);
         assertThat(builder.getRedirect(), is(target));
     }
 
@@ -45,13 +42,13 @@ public class ArgumentBuilderTest {
     public void testRedirect_withChild() throws Exception {
         final CommandNode<Object> target = mock(CommandNode.class);
         builder.then(literal("foo"));
-        builder.redirect(target, Collections::singleton);
+        builder.redirect(target);
     }
 
     @Test(expected = IllegalStateException.class)
     public void testThen_withRedirect() throws Exception {
         final CommandNode<Object> target = mock(CommandNode.class);
-        builder.redirect(target, Collections::singleton);
+        builder.redirect(target);
         builder.then(literal("foo"));
     }