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 {
-        String raw = SPLITTER.split(command).iterator().next();
+    public <S> ParsedArgument<S, Integer> parse(String command, CommandContextBuilder<S> contextBuilder) throws CommandException {
+        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) {
-        type.listSuggestions(command, output);
+    public void listSuggestions(String command, Set<String> output, CommandContextBuilder<S> contextBuilder) {
+        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