diff --git a/src/main/java/com/mojang/brigadier/CommandDispatcher.java b/src/main/java/com/mojang/brigadier/CommandDispatcher.java index 7286063..2efa85f 100644 --- a/src/main/java/com/mojang/brigadier/CommandDispatcher.java +++ b/src/main/java/com/mojang/brigadier/CommandDispatcher.java @@ -1,9 +1,11 @@ package com.mojang.brigadier; +import com.google.common.collect.Iterators; 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; @@ -79,11 +81,12 @@ public class CommandDispatcher { while (!contexts.isEmpty()) { final CommandContextBuilder context = contexts.removeLast(); - if (context.getChild() != null) { - if (!context.getNodes().isEmpty()) { - final Function> modifier = context.getNodes().keySet().iterator().next().getRedirectModifier(); - for (final S source : modifier.apply(context.getSource())) { - contexts.add(context.getChild().copy().withSource(source)); + final CommandContextBuilder child = context.getChild(); + if (child != null) { + if (!child.getNodes().isEmpty()) { + final Function, Collection> modifier = Iterators.getLast(context.getNodes().keySet().iterator()).getRedirectModifier(); + for (final S source : modifier.apply(context.build())) { + contexts.add(child.copy().withSource(source)); } } } else if (context.getCommand() != null) { @@ -151,6 +154,10 @@ public class CommandDispatcher { } } + if (parentContext != null) { + parentContext.withChild(contextBuilder); + } + return new ParseResults<>(rootContext, reader, errors); } diff --git a/src/main/java/com/mojang/brigadier/builder/ArgumentBuilder.java b/src/main/java/com/mojang/brigadier/builder/ArgumentBuilder.java index 799c2c3..34b0e38 100644 --- a/src/main/java/com/mojang/brigadier/builder/ArgumentBuilder.java +++ b/src/main/java/com/mojang/brigadier/builder/ArgumentBuilder.java @@ -1,6 +1,7 @@ package com.mojang.brigadier.builder; import com.mojang.brigadier.Command; +import com.mojang.brigadier.context.CommandContext; import com.mojang.brigadier.tree.CommandNode; import com.mojang.brigadier.tree.RootCommandNode; @@ -14,7 +15,7 @@ public abstract class ArgumentBuilder> { private Command command; private Predicate requirement = s -> true; private CommandNode target; - private Function> modifier = Collections::singleton; + private Function, Collection> modifier = s -> Collections.singleton(s.getSource()); protected abstract T getThis(); @@ -48,7 +49,11 @@ public abstract class ArgumentBuilder> { return requirement; } - public T redirect(final CommandNode target, final Function> modifier) { + public T redirect(final CommandNode target) { + return redirect(target, modifier); + } + + public T redirect(final CommandNode target, final Function, Collection> modifier) { if (!arguments.getChildren().isEmpty()) { throw new IllegalStateException("Cannot redirect a node with children"); } @@ -61,7 +66,7 @@ public abstract class ArgumentBuilder> { return target; } - public Function> getRedirectModifier() { + public Function, Collection> getRedirectModifier() { return modifier; } diff --git a/src/main/java/com/mojang/brigadier/tree/ArgumentCommandNode.java b/src/main/java/com/mojang/brigadier/tree/ArgumentCommandNode.java index 9574f0b..a0c1456 100644 --- a/src/main/java/com/mojang/brigadier/tree/ArgumentCommandNode.java +++ b/src/main/java/com/mojang/brigadier/tree/ArgumentCommandNode.java @@ -4,6 +4,7 @@ import com.mojang.brigadier.Command; import com.mojang.brigadier.StringReader; import com.mojang.brigadier.arguments.ArgumentType; import com.mojang.brigadier.builder.RequiredArgumentBuilder; +import com.mojang.brigadier.context.CommandContext; import com.mojang.brigadier.context.CommandContextBuilder; import com.mojang.brigadier.context.ParsedArgument; import com.mojang.brigadier.exceptions.CommandException; @@ -20,7 +21,7 @@ public class ArgumentCommandNode extends CommandNode { private final String name; private final ArgumentType type; - public ArgumentCommandNode(final String name, final ArgumentType type, final Command command, final Predicate requirement, final CommandNode redirect, final Function> modifier) { + public ArgumentCommandNode(final String name, final ArgumentType type, final Command command, final Predicate requirement, final CommandNode redirect, final Function, Collection> modifier) { super(command, requirement, redirect, modifier); this.name = name; this.type = type; diff --git a/src/main/java/com/mojang/brigadier/tree/CommandNode.java b/src/main/java/com/mojang/brigadier/tree/CommandNode.java index dab7240..42d6a83 100644 --- a/src/main/java/com/mojang/brigadier/tree/CommandNode.java +++ b/src/main/java/com/mojang/brigadier/tree/CommandNode.java @@ -5,6 +5,7 @@ import com.google.common.collect.Maps; import com.mojang.brigadier.Command; import com.mojang.brigadier.StringReader; import com.mojang.brigadier.builder.ArgumentBuilder; +import com.mojang.brigadier.context.CommandContext; import com.mojang.brigadier.context.CommandContextBuilder; import com.mojang.brigadier.exceptions.CommandException; @@ -20,10 +21,10 @@ public abstract class CommandNode implements Comparable> { private Map> children = Maps.newLinkedHashMap(); private final Predicate requirement; private final CommandNode redirect; - private final Function> modifier; + private final Function, Collection> modifier; private Command command; - protected CommandNode(final Command command, final Predicate requirement, final CommandNode redirect, final Function> modifier) { + protected CommandNode(final Command command, final Predicate requirement, final CommandNode redirect, final Function, Collection> modifier) { this.command = command; this.requirement = requirement; this.redirect = redirect; @@ -42,7 +43,7 @@ public abstract class CommandNode implements Comparable> { return redirect; } - public Function> getRedirectModifier() { + public Function, Collection> getRedirectModifier() { return modifier; } diff --git a/src/main/java/com/mojang/brigadier/tree/LiteralCommandNode.java b/src/main/java/com/mojang/brigadier/tree/LiteralCommandNode.java index 28e55e8..eda2b0c 100644 --- a/src/main/java/com/mojang/brigadier/tree/LiteralCommandNode.java +++ b/src/main/java/com/mojang/brigadier/tree/LiteralCommandNode.java @@ -3,6 +3,7 @@ package com.mojang.brigadier.tree; import com.mojang.brigadier.Command; import com.mojang.brigadier.StringReader; 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.ParameterizedCommandExceptionType; @@ -17,7 +18,7 @@ public class LiteralCommandNode extends CommandNode { private final String literal; - public LiteralCommandNode(final String literal, final Command command, final Predicate requirement, final CommandNode redirect, final Function> modifier) { + public LiteralCommandNode(final String literal, final Command command, final Predicate requirement, final CommandNode redirect, final Function, Collection> modifier) { super(command, requirement, redirect, modifier); this.literal = literal; } diff --git a/src/main/java/com/mojang/brigadier/tree/RootCommandNode.java b/src/main/java/com/mojang/brigadier/tree/RootCommandNode.java index cf43586..f1b6fe7 100644 --- a/src/main/java/com/mojang/brigadier/tree/RootCommandNode.java +++ b/src/main/java/com/mojang/brigadier/tree/RootCommandNode.java @@ -11,7 +11,7 @@ import java.util.Set; public class RootCommandNode extends CommandNode { public RootCommandNode() { - super(null, c -> true, null, Collections::singleton); + super(null, c -> true, null, s -> Collections.singleton(s.getSource())); } @Override diff --git a/src/test/java/com/mojang/brigadier/CommandDispatcherTest.java b/src/test/java/com/mojang/brigadier/CommandDispatcherTest.java index 21f013d..7d8d128 100644 --- a/src/test/java/com/mojang/brigadier/CommandDispatcherTest.java +++ b/src/test/java/com/mojang/brigadier/CommandDispatcherTest.java @@ -200,11 +200,11 @@ public class CommandDispatcherTest { @SuppressWarnings("unchecked") @Test public void testExecuteRedirectedMultipleTimes() throws Exception { - final Function> modifier = mock(Function.class); + final Function, Collection> modifier = mock(Function.class); final Object source1 = new Object(); final Object source2 = new Object(); - when(modifier.apply(source)).thenReturn(Lists.newArrayList(source1, source2)); + when(modifier.apply(argThat(hasProperty("source", is(source))))).thenReturn(Lists.newArrayList(source1, source2)); subject.register(literal("actual").executes(command)); subject.register(literal("redirected").redirect(subject.getRoot(), modifier));