Context-aware CommandNode#canUse method and Fix childless redirects

issue
https://github.com/Mojang/brigadier/issues/46
fix
https://github.com/PaperMC/velocity-brigadier/pull/1
https://github.com/PaperMC/velocity-brigadier/pull/2

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

@@ -332,19 +332,33 @@ public class CommandDispatcher<S> {
                     final CommandContextBuilder<S> childContext = new CommandContextBuilder<>(this, source, child.getRedirect(), reader.getCursor());
                     final ParseResults<S> parse = parseNodes(child.getRedirect(), reader, childContext);
                     context.withChild(parse.getContext());
-                    return new ParseResults<>(context, parse.getReader(), parse.getExceptions());
+                    final ParseResults<S> redirect = new ParseResults<>(context, parse.getReader(), parse.getExceptions());
+                    if (child.canUse(redirect)) {
+                        return redirect;
+                    }
                 } else {
                     final ParseResults<S> parse = parseNodes(child, reader, context);
+                    if (!child.canUse(parse)) {
+                        continue;
+                    }
                     if (potentials == null) {
                         potentials = new ArrayList<>(1);
                     }
                     potentials.add(parse);
                 }
             } else {
+                final CommandNode<S> redirect = child.getRedirect();
+                if (redirect != null && redirect.getCommand() != null) {
+                    context.withCommand(redirect.getCommand());
+                }
+                final ParseResults<S> parse = new ParseResults<>(context, reader, Collections.emptyMap());
+                if (!child.canUse(parse)) {
+                    continue;
+                }
                 if (potentials == null) {
                     potentials = new ArrayList<>(1);
                 }
-                potentials.add(new ParseResults<>(context, reader, Collections.emptyMap()));
+                potentials.add(parse);
             }
         }
 

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

@@ -4,6 +4,7 @@
 package com.mojang.brigadier.builder;
 
 import com.mojang.brigadier.Command;
+import com.mojang.brigadier.ParseResults;
 import com.mojang.brigadier.RedirectModifier;
 import com.mojang.brigadier.SingleRedirectModifier;
 import com.mojang.brigadier.tree.CommandNode;
@@ -17,6 +18,7 @@ public abstract class ArgumentBuilder<S, T extends ArgumentBuilder<S, T>> {
     private final RootCommandNode<S> arguments = new RootCommandNode<>();
     private Command<S> command;
     private Predicate<S> requirement = s -> true;
+    private Predicate<ParseResults<S>> contextRequirement = parse -> true;
     private CommandNode<S> target;
     private RedirectModifier<S> modifier = null;
     private boolean forks;
@@ -61,6 +63,15 @@ public abstract class ArgumentBuilder<S, T extends ArgumentBuilder<S, T>> {
         return requirement;
     }
 
+    public T requiresWithContext(final Predicate<ParseResults<S>> requirement) {
+        this.contextRequirement = requirement;
+        return getThis();
+    }
+
+    public Predicate<ParseResults<S>> getContextRequirement() {
+        return contextRequirement;
+    }
+
     public T redirect(final CommandNode<S> target) {
         return forward(target, null, false);
     }

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

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

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

@@ -45,8 +45,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(), getRedirectModifier(), isFork(), getSuggestionsProvider());
+        final ArgumentCommandNode<S, T> result = new ArgumentCommandNode<>(getName(), getType(), getCommand(), getRequirement(), getContextRequirement(), getRedirect(), getRedirectModifier(), isFork(), getSuggestionsProvider());
-
         for (final CommandNode<S> argument : getArguments()) {
             result.addChild(argument);
         }

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

@@ -158,6 +158,10 @@ public class CommandContext<S> {
         return nodes;
     }
 
+    public Map<String, ParsedArgument<S, ?>> getArguments() {
+        return arguments;
+    }
+
     public boolean hasNodes() {
         return !nodes.isEmpty();
     }

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

@@ -4,6 +4,7 @@
 package com.mojang.brigadier.tree;
 
 import com.mojang.brigadier.Command;
+import com.mojang.brigadier.ParseResults;
 import com.mojang.brigadier.RedirectModifier;
 import com.mojang.brigadier.StringReader;
 import com.mojang.brigadier.arguments.ArgumentType;
@@ -28,6 +29,13 @@ public class ArgumentCommandNode<S, T> extends CommandNode<S> {
     private final ArgumentType<T> type;
     private final SuggestionProvider<S> customSuggestions;
 
+    public ArgumentCommandNode(final String name, final ArgumentType<T> type, final Command<S> command, final Predicate<S> requirement, final Predicate<ParseResults<S>> contextRequirement, final CommandNode<S> redirect, final RedirectModifier<S> modifier, final boolean forks, final SuggestionProvider<S> customSuggestions) {
+        super(command, requirement, contextRequirement, redirect, modifier, forks);
+        this.name = name;
+        this.type = type;
+        this.customSuggestions = customSuggestions;
+    }
+
     public ArgumentCommandNode(final String name, final ArgumentType<T> type, final Command<S> command, final Predicate<S> requirement, final CommandNode<S> redirect, final RedirectModifier<S> modifier, final boolean forks, final SuggestionProvider<S> customSuggestions) {
         super(command, requirement, redirect, modifier, forks);
         this.name = name;

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

@@ -5,6 +5,7 @@ package com.mojang.brigadier.tree;
 
 import com.mojang.brigadier.AmbiguityConsumer;
 import com.mojang.brigadier.Command;
+import com.mojang.brigadier.ParseResults;
 import com.mojang.brigadier.RedirectModifier;
 import com.mojang.brigadier.StringReader;
 import com.mojang.brigadier.builder.ArgumentBuilder;
@@ -28,6 +29,7 @@ public abstract class CommandNode<S> implements Comparable<CommandNode<S>> {
     private final Map<String, LiteralCommandNode<S>> literals = new LinkedHashMap<>();
     private final Map<String, ArgumentCommandNode<S, ?>> arguments = new LinkedHashMap<>();
     private final Predicate<S> requirement;
+    private final Predicate<ParseResults<S>> contextRequirement;
     private final CommandNode<S> redirect;
     private final RedirectModifier<S> modifier;
     private final boolean forks;
@@ -36,6 +38,16 @@ public abstract class CommandNode<S> implements Comparable<CommandNode<S>> {
     protected CommandNode(final Command<S> command, final Predicate<S> requirement, final CommandNode<S> redirect, final RedirectModifier<S> modifier, final boolean forks) {
         this.command = command;
         this.requirement = requirement;
+        this.contextRequirement = parse -> true;
+        this.redirect = redirect;
+        this.modifier = modifier;
+        this.forks = forks;
+    }
+
+    protected CommandNode(final Command<S> command, final Predicate<S> requirement, final Predicate<ParseResults<S>> contextRequirement, final CommandNode<S> redirect, final RedirectModifier<S> modifier, final boolean forks) {
+        this.command = command;
+        this.requirement = requirement;
+        this.contextRequirement = contextRequirement;
         this.redirect = redirect;
         this.modifier = modifier;
         this.forks = forks;
@@ -65,6 +77,10 @@ public abstract class CommandNode<S> implements Comparable<CommandNode<S>> {
         return requirement.test(source);
     }
 
+    public boolean canUse(final ParseResults<S> parse) {
+        return contextRequirement.test(parse);
+    }
+
     public void addChild(final CommandNode<S> node) {
         if (node instanceof RootCommandNode) {
             throw new UnsupportedOperationException("Cannot add a RootCommandNode as a child to any other CommandNode");
@@ -138,6 +154,10 @@ public abstract class CommandNode<S> implements Comparable<CommandNode<S>> {
         return requirement;
     }
 
+    public Predicate<ParseResults<S>> getContextRequirement() {
+        return contextRequirement;
+    }
+
     public abstract String getName();
 
     public abstract String getUsageText();

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

@@ -4,6 +4,7 @@
 package com.mojang.brigadier.tree;
 
 import com.mojang.brigadier.Command;
+import com.mojang.brigadier.ParseResults;
 import com.mojang.brigadier.RedirectModifier;
 import com.mojang.brigadier.StringReader;
 import com.mojang.brigadier.builder.LiteralArgumentBuilder;
@@ -30,6 +31,12 @@ public class LiteralCommandNode<S> extends CommandNode<S> {
         this.literalLowerCase = literal.toLowerCase(Locale.ROOT);
     }
 
+    public LiteralCommandNode(final String literal, final Command<S> command, final Predicate<S> requirement, final Predicate<ParseResults<S>> contextRequirement, final CommandNode<S> redirect, final RedirectModifier<S> modifier, final boolean forks) {
+        super(command, requirement, contextRequirement, redirect, modifier, forks);
+        this.literal = literal;
+        this.literalLowerCase = literal.toLowerCase(Locale.ROOT);
+    }
+
     public String getLiteral() {
         return literal;
     }

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

@@ -10,6 +10,7 @@ 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.CommandNode;
 import com.mojang.brigadier.tree.LiteralCommandNode;
 import com.mojang.brigadier.tree.RootCommandNode;
 import org.hamcrest.CustomMatcher;
@@ -27,23 +28,13 @@ 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;
-import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.*;
-import static org.hamcrest.Matchers.hasProperty;
-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.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.*;
-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)
 public class CommandDispatcherTest {
@@ -215,6 +206,16 @@ public class CommandDispatcherTest {
         assertThat(parse.getContext().getNodes().size(), is(1));
     }
 
+    @Test
+    public void testParseChildlessRedirect() throws Exception {
+        final CommandNode<Object> target = subject.register(literal("foo").executes(command));
+        final CommandNode<Object> redirect = subject.register(literal("redirect").redirect(target));
+
+        final ParseResults<Object> parse = subject.parse("redirect", source);
+        assertThat(parse.getContext().getCommand(), equalTo(target.getCommand()));
+        assertThat(parse.getContext().getNodes().get(0).getNode(), equalTo(redirect));
+    }
+
     @SuppressWarnings("unchecked")
     @Test
     public void testExecuteAmbiguiousParentSubcommand() throws Exception {