Added suffix to integer, and pass context around some more

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

@@ -1,7 +1,5 @@
 package com.mojang.brigadier;
 
-import com.google.common.collect.ComparisonChain;
-import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
@@ -17,8 +15,6 @@ import com.mojang.brigadier.tree.RootCommandNode;
 
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.Collections;
-import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.function.Predicate;
@@ -51,6 +47,10 @@ public class CommandDispatcher<S> {
 
     public int execute(String input, S source) throws CommandException {
         final ParseResults<S> parse = parse(input, source);
+        return execute(parse);
+    }
+
+    public int execute(ParseResults<S> parse) throws CommandException {
         if (parse.getRemaining().length() > 0) {
             if (parse.getExceptions().size() == 1) {
                 throw parse.getExceptions().values().iterator().next();
@@ -69,7 +69,7 @@ public class CommandDispatcher<S> {
     }
 
     public ParseResults<S> parse(String command, S source) throws CommandException {
-        return parseNodes(root, command, new CommandContextBuilder<>(source));
+        return parseNodes(root, command, new CommandContextBuilder<>(this, source));
     }
 
     private ParseResults<S> parseNodes(CommandNode<S> node, String command, CommandContextBuilder<S> contextBuilder) throws CommandException {
@@ -190,16 +190,16 @@ public class CommandDispatcher<S> {
             if (!child.canUse(source)) {
                 continue;
             }
+            CommandContextBuilder<S> context = contextBuilder.copy();
             try {
-                CommandContextBuilder<S> context = contextBuilder.copy();
                 String remaining = child.parse(command, context);
                 if (remaining.isEmpty()) {
-                    child.listSuggestions(command, result);
+                    child.listSuggestions(command, result, context);
                 } else {
                     return findSuggestions(child, remaining.substring(1), context, result);
                 }
             } catch (CommandException e) {
-                child.listSuggestions(command, result);
+                child.listSuggestions(command, result, context);
             }
         }
 
@@ -207,7 +207,7 @@ public class CommandDispatcher<S> {
     }
 
     public String[] getCompletionSuggestions(String command, S source) {
-        final Set<String> nodes = findSuggestions(root, command, new CommandContextBuilder<>(source), Sets.newLinkedHashSet());
+        final Set<String> nodes = findSuggestions(root, command, new CommandContextBuilder<>(this, source), Sets.newLinkedHashSet());
 
         return nodes.toArray(new String[nodes.size()]);
     }

src/main/java/com/mojang/brigadier/arguments/ArgumentType.java

@@ -1,12 +1,13 @@
 package com.mojang.brigadier.arguments;
 
+import com.mojang.brigadier.context.CommandContextBuilder;
 import com.mojang.brigadier.context.ParsedArgument;
 import com.mojang.brigadier.exceptions.CommandException;
 
 import java.util.Set;
 
 public interface ArgumentType<T> {
-    <S> ParsedArgument<S, T> parse(String command) throws CommandException;
+    <S> ParsedArgument<S, T> parse(String command, CommandContextBuilder<S> contextBuilder) throws CommandException;
 
-    void listSuggestions(String command, Set<String> output);
+    <S> void listSuggestions(String command, Set<String> output, CommandContextBuilder<S> contextBuilder);
 }

src/main/java/com/mojang/brigadier/arguments/CommandArgumentType.java

@@ -0,0 +1,30 @@
+package com.mojang.brigadier.arguments;
+
+import com.mojang.brigadier.ParseResults;
+import com.mojang.brigadier.context.CommandContextBuilder;
+import com.mojang.brigadier.context.FixedParsedArgument;
+import com.mojang.brigadier.context.ParsedArgument;
+import com.mojang.brigadier.exceptions.CommandException;
+
+import java.util.Arrays;
+import java.util.Set;
+
+public class CommandArgumentType<T> implements ArgumentType<ParseResults<T>> {
+    public static <S> CommandArgumentType<S> command() {
+        return new CommandArgumentType<S>();
+    }
+
+    @Override
+    public <S> ParsedArgument<S, ParseResults<T>> parse(String command, CommandContextBuilder<S> contextBuilder) throws CommandException {
+        final ParseResults<S> parse = contextBuilder.getDispatcher().parse(command, contextBuilder.getSource());
+
+        //noinspection unchecked
+        return new FixedParsedArgument<>(command, (ParseResults<T>) parse);
+    }
+
+    @Override
+    public <S> void listSuggestions(String command, Set<String> output, CommandContextBuilder<S> contextBuilder) {
+        final String[] suggestions = contextBuilder.getDispatcher().getCompletionSuggestions(command, contextBuilder.getSource());
+        output.addAll(Arrays.asList(suggestions));
+    }
+}

src/main/java/com/mojang/brigadier/arguments/IntegerArgumentType.java

@@ -3,26 +3,29 @@ package com.mojang.brigadier.arguments;
 import com.google.common.base.Splitter;
 import com.mojang.brigadier.CommandDispatcher;
 import com.mojang.brigadier.context.CommandContext;
+import com.mojang.brigadier.context.CommandContextBuilder;
 import com.mojang.brigadier.context.FixedParsedArgument;
 import com.mojang.brigadier.context.ParsedArgument;
 import com.mojang.brigadier.exceptions.CommandException;
 import com.mojang.brigadier.exceptions.ParameterizedCommandExceptionType;
 
+import java.util.Objects;
 import java.util.Set;
 
 public class IntegerArgumentType implements ArgumentType<Integer> {
     public static final ParameterizedCommandExceptionType ERROR_NOT_A_NUMBER = new ParameterizedCommandExceptionType("argument.integer.invalid", "Expected an integer, found '${found}'", "found");
+    public static final ParameterizedCommandExceptionType ERROR_WRONG_SUFFIX = new ParameterizedCommandExceptionType("argument.integer.wrongsuffix", "Expected suffix '${suffix}'", "suffix");
     public static final ParameterizedCommandExceptionType ERROR_TOO_SMALL = new ParameterizedCommandExceptionType("argument.integer.low", "Integer must not be less than ${minimum}, found ${found}", "found", "minimum");
     public static final ParameterizedCommandExceptionType ERROR_TOO_BIG = new ParameterizedCommandExceptionType("argument.integer.big", "Integer must not be more than ${maximum}, found ${found}", "found", "maximum");
 
-    private static final Splitter SPLITTER = Splitter.on(CommandDispatcher.ARGUMENT_SEPARATOR).limit(2);
-
     private final int minimum;
     private final int maximum;
+    private final String suffix;
 
-    private IntegerArgumentType(int minimum, int maximum) {
+    private IntegerArgumentType(int minimum, int maximum, String suffix) {
         this.minimum = minimum;
         this.maximum = maximum;
+        this.suffix = suffix;
     }
 
     public static IntegerArgumentType integer() {
@@ -34,7 +37,11 @@ public class IntegerArgumentType implements ArgumentType<Integer> {
     }
 
     public static IntegerArgumentType integer(int min, int max) {
-        return new IntegerArgumentType(min, max);
+        return integer(min, max, "");
+    }
+
+    public static IntegerArgumentType integer(int min, int max, String suffix) {
+        return new IntegerArgumentType(min, max, suffix);
     }
 
     public static int getInteger(CommandContext<?> context, String name) {
@@ -42,11 +49,21 @@ public class IntegerArgumentType implements ArgumentType<Integer> {
     }
 
     @Override
-    public <S> ParsedArgument<S, Integer> parse(String command) throws CommandException {
+    public <S> ParsedArgument<S, Integer> parse(String command, CommandContextBuilder<S> contextBuilder) throws CommandException {
-        String raw = SPLITTER.split(command).iterator().next();
+        int end = command.indexOf(CommandDispatcher.ARGUMENT_SEPARATOR);
+        String raw = command;
+        if (end > -1) {
+            raw = command.substring(0, end);
+        }
+        String number = raw.substring(0, raw.length() - suffix.length());
+        String suffix = raw.substring(number.length());
+
+        if (!suffix.equals(this.suffix)) {
+            throw ERROR_WRONG_SUFFIX.create(this.suffix);
+        }
 
         try {
-            int value = Integer.parseInt(raw);
+            int value = Integer.parseInt(number);
 
             if (value < minimum) {
                 throw ERROR_TOO_SMALL.create(value, minimum);
@@ -62,7 +79,7 @@ public class IntegerArgumentType implements ArgumentType<Integer> {
     }
 
     @Override
-    public void listSuggestions(String command, Set<String> output) {
+    public <S> void listSuggestions(String command, Set<String> output, CommandContextBuilder<S> contextBuilder) {
     }
 
     @Override
@@ -71,7 +88,7 @@ public class IntegerArgumentType implements ArgumentType<Integer> {
         if (!(o instanceof IntegerArgumentType)) return false;
 
         IntegerArgumentType that = (IntegerArgumentType) o;
-        return maximum == that.maximum && minimum == that.minimum;
+        return maximum == that.maximum && minimum == that.minimum && Objects.equals(suffix, that.suffix);
     }
 
     @Override

src/main/java/com/mojang/brigadier/arguments/StringArgumentType.java

@@ -2,6 +2,7 @@ package com.mojang.brigadier.arguments;
 
 import com.mojang.brigadier.CommandDispatcher;
 import com.mojang.brigadier.context.CommandContext;
+import com.mojang.brigadier.context.CommandContextBuilder;
 import com.mojang.brigadier.context.FixedParsedArgument;
 import com.mojang.brigadier.context.ParsedArgument;
 import com.mojang.brigadier.exceptions.CommandException;
@@ -39,7 +40,7 @@ public class StringArgumentType implements ArgumentType<String> {
     }
 
     @Override
-    public <S> ParsedArgument<S, String> parse(String command) throws CommandException {
+    public <S> ParsedArgument<S, String> parse(String command, CommandContextBuilder<S> contextBuilder) throws CommandException {
         if (type == StringType.GREEDY_PHRASE) {
             return new FixedParsedArgument<>(command, command);
         } else if (type == StringType.SINGLE_WORLD) {
@@ -96,7 +97,7 @@ public class StringArgumentType implements ArgumentType<String> {
     }
 
     @Override
-    public void listSuggestions(String command, Set<String> output) {
+    public <S> void listSuggestions(String command, Set<String> output, CommandContextBuilder<S> contextBuilder) {
     }
 
     @Override

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

@@ -10,11 +10,13 @@ import java.util.Map;
 public class CommandContextBuilder<S> {
     private final Map<String, ParsedArgument<S, ?>> arguments = Maps.newHashMap();
     private final Map<CommandNode<S>, String> nodes = Maps.newLinkedHashMap();
+    private final CommandDispatcher<S> dispatcher;
     private final StringBuilder input = new StringBuilder();
     private S source;
     private Command<S> command;
 
-    public CommandContextBuilder(S source) {
+    public CommandContextBuilder(CommandDispatcher<S> dispatcher, S source) {
+        this.dispatcher = dispatcher;
         this.source = source;
     }
 
@@ -51,7 +53,7 @@ public class CommandContextBuilder<S> {
     }
 
     public CommandContextBuilder<S> copy() {
-        CommandContextBuilder<S> copy = new CommandContextBuilder<>(source);
+        CommandContextBuilder<S> copy = new CommandContextBuilder<>(dispatcher, source);
         copy.command = this.command;
         arguments.forEach((k, v) -> copy.arguments.put(k, v.copy()));
         copy.nodes.putAll(this.nodes);
@@ -70,4 +72,8 @@ public class CommandContextBuilder<S> {
     public CommandContext<S> build() {
         return new CommandContext<>(source, arguments, command, nodes, input.toString());
     }
+
+    public CommandDispatcher<S> getDispatcher() {
+        return dispatcher;
+    }
 }

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

@@ -43,7 +43,7 @@ public class ArgumentCommandNode<S, T> extends CommandNode<S> {
 
     @Override
     public String parse(String command, CommandContextBuilder<S> contextBuilder) throws CommandException {
-        ParsedArgument<S, T> parsed = type.parse(command);
+        ParsedArgument<S, T> parsed = type.parse(command, contextBuilder);
         int start = parsed.getRaw().length();
 
         contextBuilder.withArgument(name, parsed);
@@ -57,8 +57,8 @@ public class ArgumentCommandNode<S, T> extends CommandNode<S> {
     }
 
     @Override
-    public void listSuggestions(String command, Set<String> output) {
+    public void listSuggestions(String command, Set<String> output, CommandContextBuilder<S> contextBuilder) {
-        type.listSuggestions(command, output);
+        type.listSuggestions(command, output, contextBuilder);
     }
 
     @Override

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

@@ -76,7 +76,7 @@ public abstract class CommandNode<S> {
 
     public abstract String parse(String command, CommandContextBuilder<S> contextBuilder) throws CommandException;
 
-    public abstract void listSuggestions(String command, Set<String> output);
+    public abstract void listSuggestions(String command, Set<String> output, CommandContextBuilder<S> contextBuilder);
 
     public abstract ArgumentBuilder<S, ?> createBuilder();
 }

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

@@ -2,7 +2,6 @@ package com.mojang.brigadier.tree;
 
 import com.mojang.brigadier.Command;
 import com.mojang.brigadier.CommandDispatcher;
-import com.mojang.brigadier.builder.ArgumentBuilder;
 import com.mojang.brigadier.builder.LiteralArgumentBuilder;
 import com.mojang.brigadier.context.CommandContextBuilder;
 import com.mojang.brigadier.exceptions.CommandException;
@@ -44,7 +43,7 @@ public class LiteralCommandNode<S> extends CommandNode<S> {
     }
 
     @Override
-    public void listSuggestions(String command, Set<String> output) {
+    public void listSuggestions(String command, Set<String> output, CommandContextBuilder<S> contextBuilder) {
         if (literal.startsWith(command)) {
             output.add(literal);
         }

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

@@ -27,7 +27,7 @@ public class RootCommandNode<S> extends CommandNode<S> {
     }
 
     @Override
-    public void listSuggestions(String command, Set<String> output) {
+    public void listSuggestions(String command, Set<String> output, CommandContextBuilder<S> contextBuilder) {
     }
 
     @Override

src/test/java/com/mojang/brigadier/arguments/CommandArgumentTypeTest.java

@@ -0,0 +1,64 @@
+package com.mojang.brigadier.arguments;
+
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Sets;
+import com.mojang.brigadier.CommandDispatcher;
+import com.mojang.brigadier.ParseResults;
+import com.mojang.brigadier.context.CommandContextBuilder;
+import com.mojang.brigadier.context.ParsedArgument;
+import com.mojang.brigadier.exceptions.CommandException;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.runners.MockitoJUnitRunner;
+
+import java.util.Set;
+
+import static com.mojang.brigadier.arguments.CommandArgumentType.command;
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.is;
+import static org.junit.Assert.assertThat;
+import static org.junit.Assert.fail;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+@RunWith(MockitoJUnitRunner.class)
+public class CommandArgumentTypeTest {
+    @Mock
+    private Object source;
+    @Mock
+    private CommandDispatcher<Object> dispatcher;
+
+    @SuppressWarnings("unchecked")
+    @Test
+    public void testParse() throws Exception {
+        final ParseResults<Object> command = mock(ParseResults.class);
+        when(dispatcher.parse("hello world", source)).thenReturn(command);
+        final ParsedArgument<Object, ParseResults<Object>> argument = command().parse("hello world", new CommandContextBuilder<>(dispatcher, source));
+        assertThat(argument.getRaw(), equalTo("hello world"));
+        assertThat(argument.getResult(source), is(command));
+    }
+
+    @SuppressWarnings("unchecked")
+    @Test
+    public void testParse_fail() throws Exception {
+        final CommandException thrown = mock(CommandException.class);
+        when(dispatcher.parse("hello world", source)).thenThrow(thrown);
+        try {
+            command().parse("hello world", new CommandContextBuilder<>(dispatcher, source));
+            fail();
+        } catch (CommandException exception) {
+            assertThat(exception, is(thrown));
+        }
+    }
+
+    @Test
+    public void listSuggestions() throws Exception {
+        Set<String> output = Sets.newHashSet();
+        when(dispatcher.getCompletionSuggestions("foo bar baz", source)).thenReturn(new String[] {"a", "b"});
+        command().listSuggestions("foo bar baz", output, new CommandContextBuilder<>(dispatcher, source));
+        verify(dispatcher).getCompletionSuggestions("foo bar baz", source);
+        assertThat(output, equalTo(ImmutableSet.of("a", "b")));
+    }
+}

src/test/java/com/mojang/brigadier/arguments/IntegerArgumentTypeTest.java

@@ -3,6 +3,7 @@ package com.mojang.brigadier.arguments;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Sets;
 import com.google.common.testing.EqualsTester;
+import com.mojang.brigadier.CommandDispatcher;
 import com.mojang.brigadier.context.CommandContext;
 import com.mojang.brigadier.context.CommandContextBuilder;
 import com.mojang.brigadier.context.ParsedArgument;
@@ -11,6 +12,7 @@ import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
+import org.mockito.Mockito;
 import org.mockito.runners.MockitoJUnitRunner;
 
 import java.util.Set;
@@ -27,6 +29,8 @@ public class IntegerArgumentTypeTest {
     private IntegerArgumentType type;
     @Mock
     private Object source;
+    @Mock
+    private CommandDispatcher<Object> dispatcher;
 
     @Before
     public void setUp() throws Exception {
@@ -35,16 +39,57 @@ public class IntegerArgumentTypeTest {
 
     @Test
     public void testParse() throws Exception {
-        ParsedArgument<Object, Integer> result = type.parse("50");
+        ParsedArgument<Object, Integer> result = type.parse("50", new CommandContextBuilder<>(dispatcher, source));
 
         assertThat(result.getRaw(), is("50"));
         assertThat(result.getResult(source), is(50));
     }
 
+    @Test
+    public void testParse_suffix() throws Exception {
+        ParsedArgument<Object, Integer> result = integer(0, 100, "L").parse("50L", new CommandContextBuilder<>(dispatcher, source));
+
+        assertThat(result.getRaw(), is("50L"));
+        assertThat(result.getResult(source), is(50));
+    }
+
+    @Test
+    public void testParse_noSuffix() throws Exception {
+        try {
+            integer(0, 0, "L").parse("50", new CommandContextBuilder<>(dispatcher, source));
+            fail();
+        } catch (CommandException ex) {
+            assertThat(ex.getType(), is(IntegerArgumentType.ERROR_WRONG_SUFFIX));
+            assertThat(ex.getData(), is(ImmutableMap.<String, Object>of("suffix", "L")));
+        }
+    }
+
+    @Test
+    public void testParse_wrongSuffix() throws Exception {
+        try {
+            integer(0, 0, "L").parse("50B", new CommandContextBuilder<>(dispatcher, source));
+            fail();
+        } catch (CommandException ex) {
+            assertThat(ex.getType(), is(IntegerArgumentType.ERROR_WRONG_SUFFIX));
+            assertThat(ex.getData(), is(ImmutableMap.<String, Object>of("suffix", "L")));
+        }
+    }
+
+    @Test
+    public void testParse_unexpectedSuffix() throws Exception {
+        try {
+            type.parse("50L", new CommandContextBuilder<>(dispatcher, source));
+            fail();
+        } catch (CommandException ex) {
+            assertThat(ex.getType(), is(IntegerArgumentType.ERROR_NOT_A_NUMBER));
+            assertThat(ex.getData(), is(ImmutableMap.<String, Object>of("found", "50L")));
+        }
+    }
+
     @Test
     public void testParseInvalid() throws Exception {
         try {
-            type.parse("fifty");
+            type.parse("fifty", new CommandContextBuilder<>(dispatcher, source));
             fail();
         } catch (CommandException ex) {
             assertThat(ex.getType(), is(IntegerArgumentType.ERROR_NOT_A_NUMBER));
@@ -55,7 +100,7 @@ public class IntegerArgumentTypeTest {
     @Test
     public void testParseTooLow() throws Exception {
         try {
-            type.parse("-101");
+            type.parse("-101", new CommandContextBuilder<>(dispatcher, source));
             fail();
         } catch (CommandException ex) {
             assertThat(ex.getType(), is(IntegerArgumentType.ERROR_TOO_SMALL));
@@ -65,7 +110,7 @@ public class IntegerArgumentTypeTest {
 
     @Test
     public void testParseLowerLimit() throws Exception {
-        ParsedArgument<Object, Integer> result = type.parse("-100");
+        ParsedArgument<Object, Integer> result = type.parse("-100", new CommandContextBuilder<>(dispatcher, source));
 
         assertThat(result.getRaw(), is("-100"));
         assertThat(result.getResult(source), is(-100));
@@ -74,7 +119,7 @@ public class IntegerArgumentTypeTest {
     @Test
     public void testParseTooHigh() throws Exception {
         try {
-            type.parse("101");
+            type.parse("101", new CommandContextBuilder<>(dispatcher, source));
             fail();
         } catch (CommandException ex) {
             assertThat(ex.getType(), is(IntegerArgumentType.ERROR_TOO_BIG));
@@ -84,7 +129,7 @@ public class IntegerArgumentTypeTest {
 
     @Test
     public void testParseHigherLimit() throws Exception {
-        ParsedArgument<Object, Integer> result = type.parse("100");
+        ParsedArgument<Object, Integer> result = type.parse("100", new CommandContextBuilder<>(dispatcher, source));
 
         assertThat(result.getRaw(), is("100"));
         assertThat(result.getResult(source), is(100));
@@ -92,7 +137,7 @@ public class IntegerArgumentTypeTest {
 
     @Test
     public void testGetInteger() throws Exception {
-        CommandContext context = new CommandContextBuilder<>(new Object()).withArgument("foo", type.parse("100")).build();
+        CommandContext context = new CommandContextBuilder<>(dispatcher, new Object()).withArgument("foo", type.parse("100", new CommandContextBuilder<>(dispatcher, source))).build();
 
         assertThat(IntegerArgumentType.getInteger(context, "foo"), is(100));
     }
@@ -100,7 +145,8 @@ public class IntegerArgumentTypeTest {
     @Test
     public void testSuggestions() throws Exception {
         Set<String> set = Sets.newHashSet();
-        type.listSuggestions("", set);
+        @SuppressWarnings("unchecked") final CommandContextBuilder<Object> context = Mockito.mock(CommandContextBuilder.class);
+        type.listSuggestions("", set, context);
         assertThat(set, is(empty()));
     }
 
@@ -111,6 +157,8 @@ public class IntegerArgumentTypeTest {
             .addEqualityGroup(integer(-100, 100), integer(-100, 100))
             .addEqualityGroup(integer(-100, 50), integer(-100, 50))
             .addEqualityGroup(integer(-50, 100), integer(-50, 100))
+            .addEqualityGroup(integer(-50, 100, "foo"), integer(-50, 100, "foo"))
+            .addEqualityGroup(integer(-50, 100, "bar"), integer(-50, 100, "bar"))
             .testEquals();
     }
 

src/test/java/com/mojang/brigadier/arguments/StringArgumentTypeTest.java

@@ -1,11 +1,14 @@
 package com.mojang.brigadier.arguments;
 
 import com.google.common.collect.Sets;
+import com.mojang.brigadier.CommandDispatcher;
+import com.mojang.brigadier.context.CommandContextBuilder;
 import com.mojang.brigadier.context.ParsedArgument;
 import com.mojang.brigadier.exceptions.CommandException;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
+import org.mockito.Mockito;
 import org.mockito.runners.MockitoJUnitRunner;
 
 import java.util.Collections;
@@ -32,11 +35,13 @@ public class StringArgumentTypeTest {
     private StringArgumentType type;
     @Mock
     private Object source;
+    @Mock
+    private CommandDispatcher<Object> dispatcher;
 
     @Test
     public void testParseWord() throws Exception {
         type = word();
-        ParsedArgument<Object, String> result = type.parse("hello world");
+        ParsedArgument<Object, String> result = type.parse("hello world", new CommandContextBuilder<>(dispatcher, source));
 
         assertThat(result.getRaw(), is("hello"));
         assertThat(result.getResult(source), is("hello"));
@@ -45,7 +50,7 @@ public class StringArgumentTypeTest {
     @Test
     public void testParseWord_empty() throws Exception {
         type = word();
-        ParsedArgument<Object, String> result = type.parse("");
+        ParsedArgument<Object, String> result = type.parse("", new CommandContextBuilder<>(dispatcher, source));
 
         assertThat(result.getRaw(), is(""));
         assertThat(result.getResult(source), is(""));
@@ -54,7 +59,7 @@ public class StringArgumentTypeTest {
     @Test
     public void testParseWord_simple() throws Exception {
         type = word();
-        ParsedArgument<Object, String> result = type.parse("hello");
+        ParsedArgument<Object, String> result = type.parse("hello", new CommandContextBuilder<>(dispatcher, source));
 
         assertThat(result.getRaw(), is("hello"));
         assertThat(result.getResult(source), is("hello"));
@@ -63,7 +68,7 @@ public class StringArgumentTypeTest {
     @Test
     public void testParseString() throws Exception {
         type = string();
-        ParsedArgument<Object, String> result = type.parse("hello world");
+        ParsedArgument<Object, String> result = type.parse("hello world", new CommandContextBuilder<>(dispatcher, source));
 
         assertThat(result.getRaw(), is("hello"));
         assertThat(result.getResult(source), is("hello"));
@@ -72,7 +77,7 @@ public class StringArgumentTypeTest {
     @Test
     public void testParseGreedyString() throws Exception {
         type = greedyString();
-        ParsedArgument<Object, String> result = type.parse("hello world");
+        ParsedArgument<Object, String> result = type.parse("hello world", new CommandContextBuilder<>(dispatcher, source));
 
         assertThat(result.getRaw(), is("hello world"));
         assertThat(result.getResult(source), is("hello world"));
@@ -81,7 +86,7 @@ public class StringArgumentTypeTest {
     @Test
     public void testParse() throws Exception {
         type = string();
-        ParsedArgument<Object, String> result = type.parse("hello");
+        ParsedArgument<Object, String> result = type.parse("hello", new CommandContextBuilder<>(dispatcher, source));
 
         assertThat(result.getRaw(), is("hello"));
         assertThat(result.getResult(source), is("hello"));
@@ -90,7 +95,7 @@ public class StringArgumentTypeTest {
     @Test
     public void testParseWordQuoted() throws Exception {
         type = word();
-        ParsedArgument<Object, String> result = type.parse("\"hello \\\" world\"");
+        ParsedArgument<Object, String> result = type.parse("\"hello \\\" world\"", new CommandContextBuilder<>(dispatcher, source));
 
         assertThat(result.getRaw(), is("\"hello"));
         assertThat(result.getResult(source), is("\"hello"));
@@ -99,7 +104,7 @@ public class StringArgumentTypeTest {
     @Test
     public void testParseQuoted() throws Exception {
         type = string();
-        ParsedArgument<Object, String> result = type.parse("\"hello \\\" world\"");
+        ParsedArgument<Object, String> result = type.parse("\"hello \\\" world\"", new CommandContextBuilder<>(dispatcher, source));
 
         assertThat(result.getRaw(), is("\"hello \\\" world\""));
         assertThat(result.getResult(source), is("hello \" world"));
@@ -108,7 +113,7 @@ public class StringArgumentTypeTest {
     @Test
     public void testParseQuotedWithRemaining() throws Exception {
         type = string();
-        ParsedArgument<Object, String> result = type.parse("\"hello \\\" world\" with remaining");
+        ParsedArgument<Object, String> result = type.parse("\"hello \\\" world\" with remaining", new CommandContextBuilder<>(dispatcher, source));
 
         assertThat(result.getRaw(), is("\"hello \\\" world\""));
         assertThat(result.getResult(source), is("hello \" world"));
@@ -117,7 +122,7 @@ public class StringArgumentTypeTest {
     @Test
     public void testParseNotQuoted() throws Exception {
         type = string();
-        ParsedArgument<Object, String> result = type.parse("hello world");
+        ParsedArgument<Object, String> result = type.parse("hello world", new CommandContextBuilder<>(dispatcher, source));
 
         assertThat(result.getRaw(), is("hello"));
         assertThat(result.getResult(source), is("hello"));
@@ -127,7 +132,7 @@ public class StringArgumentTypeTest {
     public void testParseInvalidQuote_earlyUnquote() throws Exception {
         try {
             type = string();
-            type.parse("\"hello \"world");
+            type.parse("\"hello \"world", new CommandContextBuilder<>(dispatcher, source));
             fail();
         } catch (CommandException e) {
             assertThat(e.getType(), is(ERROR_UNEXPECTED_END_OF_QUOTE));
@@ -138,7 +143,7 @@ public class StringArgumentTypeTest {
     @Test
     public void testParseQuote_earlyUnquoteWithRemaining() throws Exception {
         type = string();
-        ParsedArgument<Object, String> result = type.parse("\"hello\" world");
+        ParsedArgument<Object, String> result = type.parse("\"hello\" world", new CommandContextBuilder<>(dispatcher, source));
 
         assertThat(result.getRaw(), is("\"hello\""));
         assertThat(result.getResult(source), is("hello"));
@@ -148,7 +153,7 @@ public class StringArgumentTypeTest {
     public void testParseInvalidQuote_lateQuote() throws Exception {
         try {
             type = string();
-            type.parse("hello\" world\"");
+            type.parse("hello\" world\"", new CommandContextBuilder<>(dispatcher, source));
             fail();
         } catch (CommandException e) {
             assertThat(e.getType(), is(ERROR_UNEXPECTED_START_OF_QUOTE));
@@ -159,7 +164,7 @@ public class StringArgumentTypeTest {
     @Test
     public void testParseQuote_lateQuoteWithRemaining() throws Exception {
         type = string();
-        ParsedArgument<Object, String> result = type.parse("hello \"world\"");
+        ParsedArgument<Object, String> result = type.parse("hello \"world\"", new CommandContextBuilder<>(dispatcher, source));
 
         assertThat(result.getRaw(), is("hello"));
         assertThat(result.getResult(source), is("hello"));
@@ -169,7 +174,7 @@ public class StringArgumentTypeTest {
     public void testParseInvalidQuote_middleQuote() throws Exception {
         try {
             type = string();
-            type.parse("hel\"lo");
+            type.parse("hel\"lo", new CommandContextBuilder<>(dispatcher, source));
             fail();
         } catch (CommandException e) {
             assertThat(e.getType(), is(ERROR_UNEXPECTED_START_OF_QUOTE));
@@ -181,7 +186,7 @@ public class StringArgumentTypeTest {
     public void testParseInvalidQuote_noUnquote() throws Exception {
         try {
             type = string();
-            type.parse("\"hello world");
+            type.parse("\"hello world", new CommandContextBuilder<>(dispatcher, source));
             fail();
         } catch (CommandException e) {
             assertThat(e.getType(), is(ERROR_EXPECTED_END_OF_QUOTE));
@@ -192,7 +197,7 @@ public class StringArgumentTypeTest {
     @Test
     public void testParseEmpty() throws Exception {
         type = string();
-        ParsedArgument<Object, String> result = type.parse("");
+        ParsedArgument<Object, String> result = type.parse("", new CommandContextBuilder<>(dispatcher, source));
 
         assertThat(result.getRaw(), is(""));
         assertThat(result.getResult(source), is(""));
@@ -202,7 +207,7 @@ public class StringArgumentTypeTest {
     public void testParseInvalidEscape_onlyEscape() throws Exception {
         try {
             type = string();
-            type.parse("\\");
+            type.parse("\\", new CommandContextBuilder<>(dispatcher, source));
             fail();
         } catch (CommandException e) {
             assertThat(e.getType(), is(ERROR_UNEXPECTED_ESCAPE));
@@ -214,7 +219,7 @@ public class StringArgumentTypeTest {
     public void testParseInvalidEscape_unknownSequence() throws Exception {
         try {
             type = string();
-            type.parse("\"\\n\"");
+            type.parse("\"\\n\"", new CommandContextBuilder<>(dispatcher, source));
             fail();
         } catch (CommandException e) {
             assertThat(e.getType(), is(ERROR_INVALID_ESCAPE));
@@ -226,7 +231,7 @@ public class StringArgumentTypeTest {
     public void testParseInvalidEscape_notQuoted() throws Exception {
         try {
             type = string();
-            type.parse("hel\\\\o");
+            type.parse("hel\\\\o", new CommandContextBuilder<>(dispatcher, source));
             fail();
         } catch (CommandException e) {
             assertThat(e.getType(), is(ERROR_UNEXPECTED_ESCAPE));
@@ -238,7 +243,8 @@ public class StringArgumentTypeTest {
     public void testSuggestions() throws Exception {
         type = string();
         Set<String> set = Sets.newHashSet();
-        type.listSuggestions("", set);
+        @SuppressWarnings("unchecked") final CommandContextBuilder<Object> context = Mockito.mock(CommandContextBuilder.class);
+        type.listSuggestions("", set, context);
         assertThat(set, is(empty()));
     }
 

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

@@ -2,6 +2,7 @@ package com.mojang.brigadier.context;
 
 import com.google.common.testing.EqualsTester;
 import com.mojang.brigadier.Command;
+import com.mojang.brigadier.CommandDispatcher;
 import com.mojang.brigadier.tree.CommandNode;
 import org.junit.Before;
 import org.junit.Test;
@@ -10,7 +11,6 @@ import org.mockito.Mock;
 import org.mockito.runners.MockitoJUnitRunner;
 
 import java.util.function.Function;
-import java.util.function.Supplier;
 
 import static com.mojang.brigadier.arguments.IntegerArgumentType.integer;
 import static com.mojang.brigadier.builder.LiteralArgumentBuilder.literal;
@@ -27,10 +27,12 @@ public class CommandContextTest {
     private CommandContextBuilder<Object> builder;
     @Mock
     private Object source;
+    @Mock
+    private CommandDispatcher<Object> dispatcher;
 
     @Before
     public void setUp() throws Exception {
-        builder = new CommandContextBuilder<>(source);
+        builder = new CommandContextBuilder<>(dispatcher, source);
     }
 
     @Test(expected = IllegalArgumentException.class)
@@ -40,13 +42,13 @@ public class CommandContextTest {
 
     @Test(expected = IllegalArgumentException.class)
     public void testGetArgument_wrongType() throws Exception {
-        CommandContext<Object> context = builder.withArgument("foo", integer().parse("123")).build();
+        CommandContext<Object> context = builder.withArgument("foo", integer().parse("123", new CommandContextBuilder<>(dispatcher, source))).build();
         context.getArgument("foo", String.class);
     }
 
     @Test
     public void testGetArgument() throws Exception {
-        CommandContext<Object> context = builder.withArgument("foo", integer().parse("123")).build();
+        CommandContext<Object> context = builder.withArgument("foo", integer().parse("123", new CommandContextBuilder<>(dispatcher, source))).build();
         assertThat(context.getArgument("foo", int.class), is(123));
     }
 
@@ -64,13 +66,13 @@ public class CommandContextTest {
         CommandNode<Object> node = mock(CommandNode.class);
         CommandNode<Object> otherNode = mock(CommandNode.class);
         new EqualsTester()
-            .addEqualityGroup(new CommandContextBuilder<>(source).build(), new CommandContextBuilder<>(source).build())
+            .addEqualityGroup(new CommandContextBuilder<>(dispatcher, source).build(), new CommandContextBuilder<>(dispatcher, source).build())
-            .addEqualityGroup(new CommandContextBuilder<>(otherSource).build(), new CommandContextBuilder<>(otherSource).build())
+            .addEqualityGroup(new CommandContextBuilder<>(dispatcher, otherSource).build(), new CommandContextBuilder<>(dispatcher, otherSource).build())
-            .addEqualityGroup(new CommandContextBuilder<>(source).withCommand(command).build(), new CommandContextBuilder<>(source).withCommand(command).build())
+            .addEqualityGroup(new CommandContextBuilder<>(dispatcher, source).withCommand(command).build(), new CommandContextBuilder<>(dispatcher, source).withCommand(command).build())
-            .addEqualityGroup(new CommandContextBuilder<>(source).withCommand(otherCommand).build(), new CommandContextBuilder<>(source).withCommand(otherCommand).build())
+            .addEqualityGroup(new CommandContextBuilder<>(dispatcher, source).withCommand(otherCommand).build(), new CommandContextBuilder<>(dispatcher, source).withCommand(otherCommand).build())
-            .addEqualityGroup(new CommandContextBuilder<>(source).withArgument("foo", integer().parse("123")).build(), new CommandContextBuilder<>(source).withArgument("foo", integer().parse("123")).build())
+            .addEqualityGroup(new CommandContextBuilder<>(dispatcher, source).withArgument("foo", integer().parse("123", new CommandContextBuilder<>(dispatcher, source))).build(), new CommandContextBuilder<>(dispatcher, source).withArgument("foo", integer().parse("123", new CommandContextBuilder<>(dispatcher, source))).build())
-            .addEqualityGroup(new CommandContextBuilder<>(source).withNode(node, "foo").withNode(otherNode, "bar").build(), new CommandContextBuilder<>(source).withNode(node, "foo").withNode(otherNode, "bar").build())
+            .addEqualityGroup(new CommandContextBuilder<>(dispatcher, source).withNode(node, "foo").withNode(otherNode, "bar").build(), new CommandContextBuilder<>(dispatcher, source).withNode(node, "foo").withNode(otherNode, "bar").build())
-            .addEqualityGroup(new CommandContextBuilder<>(source).withNode(otherNode, "bar").withNode(node, "foo").build(), new CommandContextBuilder<>(source).withNode(otherNode, "bar").withNode(node, "foo").build())
+            .addEqualityGroup(new CommandContextBuilder<>(dispatcher, source).withNode(otherNode, "bar").withNode(node, "foo").build(), new CommandContextBuilder<>(dispatcher, source).withNode(otherNode, "bar").withNode(node, "foo").build())
             .testEquals();
     }
 

src/test/java/com/mojang/brigadier/tree/ArgumentCommandNodeTest.java

@@ -4,12 +4,14 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Sets;
 import com.google.common.testing.EqualsTester;
 import com.mojang.brigadier.Command;
+import com.mojang.brigadier.CommandDispatcher;
 import com.mojang.brigadier.arguments.IntegerArgumentType;
 import com.mojang.brigadier.builder.RequiredArgumentBuilder;
 import com.mojang.brigadier.context.CommandContextBuilder;
 import com.mojang.brigadier.exceptions.CommandException;
 import org.junit.Before;
 import org.junit.Test;
+import org.mockito.Mockito;
 
 import java.util.Set;
 
@@ -34,7 +36,7 @@ public class ArgumentCommandNodeTest extends AbstractCommandNodeTest {
     @Before
     public void setUp() throws Exception {
         node = argument("foo", integer()).build();
-        contextBuilder = new CommandContextBuilder<>(new Object());
+        contextBuilder = new CommandContextBuilder<>(new CommandDispatcher<>(), new Object());
     }
 
     @Test
@@ -72,7 +74,8 @@ public class ArgumentCommandNodeTest extends AbstractCommandNodeTest {
     @Test
     public void testSuggestions() throws Exception {
         Set<String> set = Sets.newHashSet();
-        node.listSuggestions("", set);
+        @SuppressWarnings("unchecked") final CommandContextBuilder<Object> context = Mockito.mock(CommandContextBuilder.class);
+        node.listSuggestions("", set, context);
         assertThat(set, is(empty()));
     }
 

src/test/java/com/mojang/brigadier/tree/LiteralCommandNodeTest.java

@@ -4,12 +4,13 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Sets;
 import com.google.common.testing.EqualsTester;
 import com.mojang.brigadier.Command;
-import com.mojang.brigadier.builder.ArgumentBuilder;
+import com.mojang.brigadier.CommandDispatcher;
 import com.mojang.brigadier.builder.LiteralArgumentBuilder;
 import com.mojang.brigadier.context.CommandContextBuilder;
 import com.mojang.brigadier.exceptions.CommandException;
 import org.junit.Before;
 import org.junit.Test;
+import org.mockito.Mockito;
 
 import java.util.Set;
 
@@ -33,7 +34,7 @@ public class LiteralCommandNodeTest extends AbstractCommandNodeTest {
     @Before
     public void setUp() throws Exception {
         node = literal("foo").build();
-        contextBuilder = new CommandContextBuilder<>(new Object());
+        contextBuilder = new CommandContextBuilder<>(new CommandDispatcher<>(), new Object());
     }
 
     @Test
@@ -76,19 +77,21 @@ public class LiteralCommandNodeTest extends AbstractCommandNodeTest {
     @Test
     public void testSuggestions() throws Exception {
         Set<String> set = Sets.newHashSet();
-        node.listSuggestions("", set);
+        @SuppressWarnings("unchecked") final CommandContextBuilder<Object> context = Mockito.mock(CommandContextBuilder.class);
+
+        node.listSuggestions("", set, context);
         assertThat(set, equalTo(Sets.newHashSet("foo")));
 
         set.clear();
-        node.listSuggestions("foo", set);
+        node.listSuggestions("foo", set, context);
         assertThat(set, equalTo(Sets.newHashSet("foo")));
 
         set.clear();
-        node.listSuggestions("food", set);
+        node.listSuggestions("food", set, context);
         assertThat(set, is(empty()));
 
         set.clear();
-        node.listSuggestions("b", set);
+        node.listSuggestions("b", set, context);
         assertThat(set, is(empty()));
     }
 

src/test/java/com/mojang/brigadier/tree/RootCommandNodeTest.java

@@ -2,9 +2,11 @@ package com.mojang.brigadier.tree;
 
 import com.google.common.collect.Sets;
 import com.google.common.testing.EqualsTester;
+import com.mojang.brigadier.CommandDispatcher;
 import com.mojang.brigadier.context.CommandContextBuilder;
 import org.junit.Before;
 import org.junit.Test;
+import org.mockito.Mockito;
 
 import java.util.Set;
 
@@ -28,7 +30,7 @@ public class RootCommandNodeTest extends AbstractCommandNodeTest {
 
     @Test
     public void testParse() throws Exception {
-        assertThat(node.parse("foo bar baz", new CommandContextBuilder<>(new Object())), is("foo bar baz"));
+        assertThat(node.parse("foo bar baz", new CommandContextBuilder<>(new CommandDispatcher<>(), new Object())), is("foo bar baz"));
     }
 
     @Test(expected = UnsupportedOperationException.class)
@@ -44,7 +46,8 @@ public class RootCommandNodeTest extends AbstractCommandNodeTest {
     @Test
     public void testSuggestions() throws Exception {
         Set<String> set = Sets.newHashSet();
-        node.listSuggestions("", set);
+        @SuppressWarnings("unchecked") final CommandContextBuilder<Object> context = Mockito.mock(CommandContextBuilder.class);
+        node.listSuggestions("", set, context);
         assertThat(set, is(empty()));
     }