Pass S (source type, was sometimes T) through to all nodes, so Command can have source type

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

@@ -2,6 +2,6 @@ package com.mojang.brigadier;
 
 import com.mojang.brigadier.context.CommandContext;
 
-public interface Command {
+public interface Command<S> {
-    void run(CommandContext context);
+    void run(CommandContext<S> context);
 }

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

@@ -17,14 +17,7 @@ import java.util.Set;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
 
-public class CommandDispatcher<T> {
+public class CommandDispatcher<S> {
-    private static final Predicate<CommandNode> HAS_COMMAND = new Predicate<CommandNode>() {
-        @Override
-        public boolean test(CommandNode input) {
-            return input != null && (input.getCommand() != null || input.getChildren().stream().anyMatch(HAS_COMMAND));
-        }
-    };
-
     public static final SimpleCommandExceptionType ERROR_UNKNOWN_COMMAND = new SimpleCommandExceptionType("command.unknown", "Unknown command");
     public static final String ARGUMENT_SEPARATOR = " ";
     private static final String USAGE_OPTIONAL_OPEN = "[";
@@ -33,23 +26,29 @@ public class CommandDispatcher<T> {
     private static final String USAGE_REQUIRED_CLOSE = ")";
     private static final String USAGE_OR = "|";
 
-    private final RootCommandNode root = new RootCommandNode();
+    private final RootCommandNode<S> root = new RootCommandNode<>();
+    private final Predicate<CommandNode<S>> hasCommand = new Predicate<CommandNode<S>>() {
+        @Override
+        public boolean test(CommandNode<S> input) {
+            return input != null && (input.getCommand() != null || input.getChildren().stream().anyMatch(hasCommand));
+        }
+    };
 
-    public void register(LiteralArgumentBuilder command) {
+    public void register(LiteralArgumentBuilder<S> command) {
         root.addChild(command.build());
     }
 
-    public void execute(String command, T source) throws CommandException {
+    public void execute(String command, S source) throws CommandException {
-        CommandContext<T> context = parseNodes(root, command, new CommandContextBuilder<>(source));
+        CommandContext<S> context = parseNodes(root, command, new CommandContextBuilder<>(source));
         context.getCommand().run(context);
     }
 
-    private CommandContext<T> parseNodes(CommandNode node, String command, CommandContextBuilder<T> contextBuilder) throws CommandException {
+    private CommandContext<S> parseNodes(CommandNode<S> node, String command, CommandContextBuilder<S> contextBuilder) throws CommandException {
         CommandException exception = null;
 
-        for (CommandNode child : node.getChildren()) {
+        for (CommandNode<S> child : node.getChildren()) {
             try {
-                CommandContextBuilder<T> context = contextBuilder.copy();
+                CommandContextBuilder<S> context = contextBuilder.copy();
                 String remaining = child.parse(command, context);
                 if (child.getCommand() != null) {
                     context.withCommand(child.getCommand());
@@ -74,10 +73,10 @@ public class CommandDispatcher<T> {
         return contextBuilder.build();
     }
 
-    public String getUsage(String command, T source) throws CommandException {
+    public String getUsage(String command, S source) throws CommandException {
-        CommandContext<T> context = parseNodes(root, command, new CommandContextBuilder<>(source));
+        CommandContext<S> context = parseNodes(root, command, new CommandContextBuilder<>(source));
-        CommandNode base = Iterables.getLast(context.getNodes().keySet());
+        CommandNode<S> base = Iterables.getLast(context.getNodes().keySet());
-        List<CommandNode> children = base.getChildren().stream().filter(HAS_COMMAND).collect(Collectors.toList());
+        List<CommandNode<S>> children = base.getChildren().stream().filter(hasCommand).collect(Collectors.toList());
         boolean optional = base.getCommand() != null;
 
         if (children.isEmpty()) {
@@ -113,10 +112,10 @@ public class CommandDispatcher<T> {
         return result.toString();
     }
 
-    private Set<String> findSuggestions(CommandNode node, String command, CommandContextBuilder<T> contextBuilder, Set<String> result) {
+    private Set<String> findSuggestions(CommandNode<S> node, String command, CommandContextBuilder<S> contextBuilder, Set<String> result) {
-        for (CommandNode child : node.getChildren()) {
+        for (CommandNode<S> child : node.getChildren()) {
             try {
-                CommandContextBuilder<T> context = contextBuilder.copy();
+                CommandContextBuilder<S> context = contextBuilder.copy();
                 String remaining = child.parse(command, context);
                 if (remaining.isEmpty()) {
                     child.listSuggestions(command, result);
@@ -131,7 +130,7 @@ public class CommandDispatcher<T> {
         return result;
     }
 
-    public String[] getCompletionSuggestions(String command, T source) {
+    public String[] getCompletionSuggestions(String command, S source) {
         final Set<String> nodes = findSuggestions(root, command, new CommandContextBuilder<>(source), Sets.newLinkedHashSet());
 
         return nodes.toArray(new String[nodes.size()]);

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

@@ -6,29 +6,29 @@ import com.mojang.brigadier.tree.RootCommandNode;
 
 import java.util.Collection;
 
-public abstract class ArgumentBuilder<T extends ArgumentBuilder<?>> {
+public abstract class ArgumentBuilder<S, T extends ArgumentBuilder<S, ?>> {
-    private final RootCommandNode arguments = new RootCommandNode();
+    private final RootCommandNode<S> arguments = new RootCommandNode<>();
-    private Command command;
+    private Command<S> command;
 
     protected abstract T getThis();
 
-    public T then(ArgumentBuilder argument) {
+    public T then(ArgumentBuilder<S, ?> argument) {
         arguments.addChild(argument.build());
         return getThis();
     }
 
-    public Collection<CommandNode> getArguments() {
+    public Collection<CommandNode<S>> getArguments() {
         return arguments.getChildren();
     }
 
-    public T executes(Command command) {
+    public T executes(Command<S> command) {
         this.command = command;
         return getThis();
     }
 
-    protected Command getCommand() {
+    protected Command<S> getCommand() {
         return command;
     }
 
-    public abstract CommandNode build();
+    public abstract CommandNode<S> build();
 }

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

@@ -3,19 +3,19 @@ package com.mojang.brigadier.builder;
 import com.mojang.brigadier.tree.CommandNode;
 import com.mojang.brigadier.tree.LiteralCommandNode;
 
-public class LiteralArgumentBuilder extends ArgumentBuilder<LiteralArgumentBuilder> {
+public class LiteralArgumentBuilder<S> extends ArgumentBuilder<S, LiteralArgumentBuilder<S>> {
     private final String literal;
 
     protected LiteralArgumentBuilder(String literal) {
         this.literal = literal;
     }
 
-    public static LiteralArgumentBuilder literal(String name) {
+    public static <S> LiteralArgumentBuilder<S> literal(String name) {
-        return new LiteralArgumentBuilder(name);
+        return new LiteralArgumentBuilder<>(name);
     }
 
     @Override
-    protected LiteralArgumentBuilder getThis() {
+    protected LiteralArgumentBuilder<S> getThis() {
         return this;
     }
 
@@ -24,10 +24,10 @@ public class LiteralArgumentBuilder extends ArgumentBuilder<LiteralArgumentBuild
     }
 
     @Override
-    public LiteralCommandNode build() {
+    public LiteralCommandNode<S> build() {
-        LiteralCommandNode result = new LiteralCommandNode(getLiteral(), getCommand());
+        LiteralCommandNode<S> result = new LiteralCommandNode<>(getLiteral(), getCommand());
 
-        for (CommandNode argument : getArguments()) {
+        for (CommandNode<S> argument : getArguments()) {
             result.addChild(argument);
         }
 

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

@@ -4,7 +4,7 @@ import com.mojang.brigadier.arguments.CommandArgumentType;
 import com.mojang.brigadier.tree.ArgumentCommandNode;
 import com.mojang.brigadier.tree.CommandNode;
 
-public class RequiredArgumentBuilder<T> extends ArgumentBuilder<RequiredArgumentBuilder<T>> {
+public class RequiredArgumentBuilder<S, T> extends ArgumentBuilder<S, RequiredArgumentBuilder<S, T>> {
     private final String name;
     private final CommandArgumentType<T> type;
 
@@ -13,12 +13,12 @@ public class RequiredArgumentBuilder<T> extends ArgumentBuilder<RequiredArgument
         this.type = type;
     }
 
-    public static <T> RequiredArgumentBuilder<T> argument(String name, CommandArgumentType<T> type) {
+    public static <S, T> RequiredArgumentBuilder<S, T> argument(String name, CommandArgumentType<T> type) {
         return new RequiredArgumentBuilder<>(name, type);
     }
 
     @Override
-    protected RequiredArgumentBuilder<T> getThis() {
+    protected RequiredArgumentBuilder<S, T> getThis() {
         return this;
     }
 
@@ -30,10 +30,10 @@ public class RequiredArgumentBuilder<T> extends ArgumentBuilder<RequiredArgument
         return name;
     }
 
-    public ArgumentCommandNode<T> build() {
+    public ArgumentCommandNode<S, T> build() {
-        ArgumentCommandNode<T> result = new ArgumentCommandNode<>(getName(), getType(), getCommand());
+        ArgumentCommandNode<S, T> result = new ArgumentCommandNode<>(getName(), getType(), getCommand());
 
-        for (CommandNode argument : getArguments()) {
+        for (CommandNode<S> argument : getArguments()) {
             result.addChild(argument);
         }
 

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

@@ -9,26 +9,26 @@ import com.mojang.brigadier.tree.CommandNode;
 
 import java.util.Map;
 
-public class CommandContext<T> {
+public class CommandContext<S> {
     private final Joiner JOINER = Joiner.on(CommandDispatcher.ARGUMENT_SEPARATOR);
 
-    private final T source;
+    private final S source;
     private final Map<String, ParsedArgument<?>> arguments;
-    private final Command command;
+    private final Command<S> command;
-    private final Map<CommandNode, String> nodes;
+    private final Map<CommandNode<S>, String> nodes;
 
-    public CommandContext(T source, Map<String, ParsedArgument<?>> arguments, Command command, Map<CommandNode, String> nodes) {
+    public CommandContext(S source, Map<String, ParsedArgument<?>> arguments, Command<S> command, Map<CommandNode<S>, String> nodes) {
         this.source = source;
         this.arguments = arguments;
         this.command = command;
         this.nodes = nodes;
     }
 
-    public Command getCommand() {
+    public Command<S> getCommand() {
         return command;
     }
 
-    public T getSource() {
+    public S getSource() {
         return source;
     }
 
@@ -75,7 +75,7 @@ public class CommandContext<T> {
         return JOINER.join(nodes.values());
     }
 
-    public Map<CommandNode, String> getNodes() {
+    public Map<CommandNode<S>, String> getNodes() {
         return nodes;
     }
 }

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

@@ -6,17 +6,17 @@ import com.mojang.brigadier.tree.CommandNode;
 
 import java.util.Map;
 
-public class CommandContextBuilder<T> {
+public class CommandContextBuilder<S> {
     private final Map<String, ParsedArgument<?>> arguments = Maps.newHashMap();
-    private final Map<CommandNode, String> nodes = Maps.newLinkedHashMap();
+    private final Map<CommandNode<S>, String> nodes = Maps.newLinkedHashMap();
-    private final T source;
+    private final S source;
-    private Command command;
+    private Command<S> command;
 
-    public CommandContextBuilder(T source) {
+    public CommandContextBuilder(S source) {
         this.source = source;
     }
 
-    public CommandContextBuilder<T> withArgument(String name, ParsedArgument<?> argument) {
+    public CommandContextBuilder<S> withArgument(String name, ParsedArgument<?> argument) {
         this.arguments.put(name, argument);
         return this;
     }
@@ -25,25 +25,25 @@ public class CommandContextBuilder<T> {
         return arguments;
     }
 
-    public CommandContextBuilder<T> withCommand(Command command) {
+    public CommandContextBuilder<S> withCommand(Command<S> command) {
         this.command = command;
         return this;
     }
 
-    public CommandContextBuilder<T> withNode(CommandNode node, String raw) {
+    public CommandContextBuilder<S> withNode(CommandNode<S> node, String raw) {
         this.nodes.put(node, raw);
         return this;
     }
 
-    public CommandContextBuilder<T> copy() {
+    public CommandContextBuilder<S> copy() {
-        CommandContextBuilder<T> copy = new CommandContextBuilder<>(source);
+        CommandContextBuilder<S> copy = new CommandContextBuilder<>(source);
         copy.command = this.command;
         copy.arguments.putAll(this.arguments);
         copy.nodes.putAll(this.nodes);
         return copy;
     }
 
-    public CommandContext<T> build() {
+    public CommandContext<S> build() {
         return new CommandContext<>(source, arguments, command, nodes);
     }
 }

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

@@ -8,14 +8,14 @@ import com.mojang.brigadier.exceptions.CommandException;
 
 import java.util.Set;
 
-public class ArgumentCommandNode<T> extends CommandNode {
+public class ArgumentCommandNode<S, T> extends CommandNode<S> {
     private static final String USAGE_ARGUMENT_OPEN = "<";
     private static final String USAGE_ARGUMENT_CLOSE = ">";
 
     private final String name;
     private final CommandArgumentType<T> type;
 
-    public ArgumentCommandNode(String name, CommandArgumentType<T> type, Command command) {
+    public ArgumentCommandNode(String name, CommandArgumentType<T> type, Command<S> command) {
         super(command);
         this.name = name;
         this.type = type;
@@ -40,7 +40,7 @@ public class ArgumentCommandNode<T> extends CommandNode {
     }
 
     @Override
-    public String parse(String command, CommandContextBuilder<?> contextBuilder) throws CommandException {
+    public String parse(String command, CommandContextBuilder<S> contextBuilder) throws CommandException {
         ParsedArgument<T> parsed = type.parse(command);
         int start = parsed.getRaw().length();
 

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

@@ -9,30 +9,30 @@ import java.util.Collection;
 import java.util.Map;
 import java.util.Set;
 
-public abstract class CommandNode {
+public abstract class CommandNode<S> {
-    private final Map<Object, CommandNode> children = Maps.newLinkedHashMap();
+    private final Map<Object, CommandNode<S>> children = Maps.newLinkedHashMap();
-    private Command command;
+    private Command<S> command;
 
-    protected CommandNode(Command command) {
+    protected CommandNode(Command<S> command) {
         this.command = command;
     }
 
-    public Command getCommand() {
+    public Command<S> getCommand() {
         return command;
     }
 
-    public Collection<CommandNode> getChildren() {
+    public Collection<CommandNode<S>> getChildren() {
         return children.values();
     }
 
-    public void addChild(CommandNode node) {
+    public void addChild(CommandNode<S> node) {
-        CommandNode child = children.get(node.getMergeKey());
+        CommandNode<S> child = children.get(node.getMergeKey());
         if (child != null) {
             // We've found something to merge onto
             if (node.getCommand() != null) {
                 child.command = node.getCommand();
             }
-            for (CommandNode grandchild : node.getChildren()) {
+            for (CommandNode<S> grandchild : node.getChildren()) {
                 child.addChild(grandchild);
             }
         } else {
@@ -45,7 +45,7 @@ public abstract class CommandNode {
         if (this == o) return true;
         if (!(o instanceof CommandNode)) return false;
 
-        CommandNode that = (CommandNode) o;
+        CommandNode<S> that = (CommandNode<S>) o;
 
         if (!children.equals(that.children)) return false;
         if (command != null ? !command.equals(that.command) : that.command != null) return false;
@@ -62,7 +62,7 @@ public abstract class CommandNode {
 
     public abstract String getUsageText();
 
-    public abstract String parse(String command, CommandContextBuilder<?> contextBuilder) throws CommandException;
+    public abstract String parse(String command, CommandContextBuilder<S> contextBuilder) throws CommandException;
 
     public abstract void listSuggestions(String command, Set<String> output);
 }

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

@@ -8,12 +8,12 @@ import com.mojang.brigadier.exceptions.ParameterizedCommandExceptionType;
 
 import java.util.Set;
 
-public class LiteralCommandNode extends CommandNode {
+public class LiteralCommandNode<S> extends CommandNode<S> {
     public static final ParameterizedCommandExceptionType ERROR_INCORRECT_LITERAL = new ParameterizedCommandExceptionType("argument.literal.incorrect", "Expected literal ${expected}", "expected");
 
     private final String literal;
 
-    public LiteralCommandNode(String literal, Command command) {
+    public LiteralCommandNode(String literal, Command<S> command) {
         super(command);
         this.literal = literal;
     }
@@ -28,7 +28,7 @@ public class LiteralCommandNode extends CommandNode {
     }
 
     @Override
-    public String parse(String command, CommandContextBuilder<?> contextBuilder) throws CommandException {
+    public String parse(String command, CommandContextBuilder<S> contextBuilder) throws CommandException {
         String expected = literal + (command.length() > literal.length() ? CommandDispatcher.ARGUMENT_SEPARATOR : "");
 
         if (!command.startsWith(expected)) {

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

@@ -5,7 +5,7 @@ import com.mojang.brigadier.exceptions.CommandException;
 
 import java.util.Set;
 
-public class RootCommandNode extends CommandNode {
+public class RootCommandNode<S> extends CommandNode<S> {
     public RootCommandNode() {
         super(null);
     }
@@ -21,7 +21,7 @@ public class RootCommandNode extends CommandNode {
     }
 
     @Override
-    public String parse(String command, CommandContextBuilder<?> contextBuilder) throws CommandException {
+    public String parse(String command, CommandContextBuilder<S> contextBuilder) throws CommandException {
         return command;
     }
 

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

@@ -26,7 +26,7 @@ import static org.mockito.Mockito.*;
 public class CommandDispatcherTest {
     private CommandDispatcher<Object> subject;
     @Mock
-    private Command command;
+    private Command<Object> command;
     @Mock
     private Object source;
 
@@ -35,6 +35,7 @@ public class CommandDispatcherTest {
         subject = new CommandDispatcher<>();
     }
 
+    @SuppressWarnings("unchecked")
     @Test
     public void testCreateAndExecuteCommand() throws Exception {
         subject.register(literal("foo").executes(command));
@@ -43,6 +44,7 @@ public class CommandDispatcherTest {
         verify(command).run(any(CommandContext.class));
     }
 
+    @SuppressWarnings("unchecked")
     @Test
     public void testCreateAndMergeCommands() throws Exception {
         subject.register(literal("base").then(literal("foo")).executes(command));
@@ -53,11 +55,12 @@ public class CommandDispatcherTest {
         verify(command, times(2)).run(any(CommandContext.class));
     }
 
+    @SuppressWarnings("unchecked")
     @Test
     public void testCreateAndExecuteOverlappingCommands() throws Exception {
-        Command one = mock(Command.class);
+        Command<Object> one = mock(Command.class);
-        Command two = mock(Command.class);
+        Command<Object> two = mock(Command.class);
-        Command three = mock(Command.class);
+        Command<Object> three = mock(Command.class);
 
         subject.register(
             literal("foo").then(
@@ -109,9 +112,10 @@ public class CommandDispatcherTest {
         }
     }
 
+    @SuppressWarnings("unchecked")
     @Test
     public void testExecuteSubcommand() throws Exception {
-        Command subCommand = mock(Command.class);
+        Command<Object> subCommand = mock(Command.class);
 
         subject.register(literal("foo").then(
             literal("a")

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

@@ -24,7 +24,7 @@ public class CommandDispatcherUsagesTest {
     @Mock
     private Object source;
     @Mock
-    private Command command;
+    private Command<Object>  command;
 
     @Before
     public void setUp() throws Exception {

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

@@ -8,34 +8,35 @@ import static com.mojang.brigadier.arguments.IntegerArgumentType.integer;
 import static com.mojang.brigadier.builder.RequiredArgumentBuilder.argument;
 import static org.hamcrest.Matchers.hasItems;
 import static org.hamcrest.Matchers.hasSize;
+import static org.hamcrest.core.IsCollectionContaining.hasItem;
 import static org.junit.Assert.assertThat;
 
 public class ArgumentBuilderTest {
-    private TestableArgumentBuilder builder;
+    private TestableArgumentBuilder<Object> builder;
 
     @Before
     public void setUp() throws Exception {
-        builder = new TestableArgumentBuilder();
+        builder = new TestableArgumentBuilder<>();
     }
 
     @Test
     public void testArguments() throws Exception {
-        RequiredArgumentBuilder argument = argument("bar", integer());
+        RequiredArgumentBuilder<Object, ?> argument = argument("bar", integer());
 
         builder.then(argument);
 
         assertThat(builder.getArguments(), hasSize(1));
-        assertThat(builder.getArguments(), hasItems((CommandNode) argument.build()));
+        assertThat(builder.getArguments(), hasItem((CommandNode<Object>) argument.build()));
     }
 
-    private static class TestableArgumentBuilder extends ArgumentBuilder<TestableArgumentBuilder> {
+    private static class TestableArgumentBuilder<S> extends ArgumentBuilder<S, TestableArgumentBuilder<S>> {
         @Override
-        protected TestableArgumentBuilder getThis() {
+        protected TestableArgumentBuilder<S> getThis() {
             return this;
         }
 
         @Override
-        public CommandNode build() {
+        public CommandNode<S> build() {
             return null;
         }
     }

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

@@ -13,26 +13,26 @@ import static org.hamcrest.Matchers.is;
 import static org.junit.Assert.assertThat;
 
 public class LiteralArgumentBuilderTest {
-    private LiteralArgumentBuilder builder;
+    private LiteralArgumentBuilder<Object> builder;
     @Mock
     private
-    Command command;
+    Command<Object> command;
 
     @Before
     public void setUp() throws Exception {
-        builder = new LiteralArgumentBuilder("foo");
+        builder = new LiteralArgumentBuilder<>("foo");
     }
 
     @Test
     public void testBuild() throws Exception {
-        LiteralCommandNode node = builder.build();
+        LiteralCommandNode<Object> node = builder.build();
 
         assertThat(node.getLiteral(), is("foo"));
     }
 
     @Test
     public void testBuildWithExecutor() throws Exception {
-        LiteralCommandNode node = builder.executes(command).build();
+        LiteralCommandNode<Object> node = builder.executes(command).build();
 
         assertThat(node.getLiteral(), is("foo"));
         assertThat(node.getCommand(), is(command));
@@ -42,7 +42,7 @@ public class LiteralArgumentBuilderTest {
     public void testBuildWithChildren() throws Exception {
         builder.then(argument("bar", integer()));
         builder.then(argument("baz", integer()));
-        LiteralCommandNode node = builder.build();
+        LiteralCommandNode<Object> node = builder.build();
 
         assertThat(node.getChildren(), hasSize(2));
     }

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

@@ -16,10 +16,10 @@ import static org.junit.Assert.assertThat;
 public class RequiredArgumentBuilderTest {
     @Mock
     private CommandArgumentType<Integer> type;
-    private RequiredArgumentBuilder<Integer> builder;
+    private RequiredArgumentBuilder<Object, Integer> builder;
     @Mock
     private
-    Command command;
+    Command<Object> command;
 
     @Before
     public void setUp() throws Exception {
@@ -28,7 +28,7 @@ public class RequiredArgumentBuilderTest {
 
     @Test
     public void testBuild() throws Exception {
-        ArgumentCommandNode<Integer> node = builder.build();
+        ArgumentCommandNode<Object, Integer> node = builder.build();
 
         assertThat(node.getName(), is("foo"));
         assertThat(node.getType(), is(type));
@@ -36,7 +36,7 @@ public class RequiredArgumentBuilderTest {
 
     @Test
     public void testBuildWithExecutor() throws Exception {
-        ArgumentCommandNode<Integer> node = builder.executes(command).build();
+        ArgumentCommandNode<Object, Integer> node = builder.executes(command).build();
 
         assertThat(node.getName(), is("foo"));
         assertThat(node.getType(), is(type));
@@ -47,7 +47,7 @@ public class RequiredArgumentBuilderTest {
     public void testBuildWithChildren() throws Exception {
         builder.then(argument("bar", integer()));
         builder.then(argument("baz", integer()));
-        ArgumentCommandNode node = builder.build();
+        ArgumentCommandNode<Object, Integer> node = builder.build();
 
         assertThat(node.getChildren(), hasSize(2));
     }

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

@@ -49,13 +49,14 @@ public class CommandContextTest {
         assertThat(builder.build().getSource(), is(source));
     }
 
+    @SuppressWarnings("unchecked")
     @Test
     public void testEquals() throws Exception {
         Object otherSource = new Object();
-        Command command = mock(Command.class);
+        Command<Object> command = mock(Command.class);
-        Command otherCommand = mock(Command.class);
+        Command<Object> otherCommand = mock(Command.class);
-        CommandNode node = mock(CommandNode.class);
+        CommandNode<Object> node = mock(CommandNode.class);
-        CommandNode otherNode = mock(CommandNode.class);
+        CommandNode<Object> otherNode = mock(CommandNode.class);
         new EqualsTester()
             .addEqualityGroup(new CommandContextBuilder<>(source).build(), new CommandContextBuilder<>(source).build())
             .addEqualityGroup(new CommandContextBuilder<>(otherSource).build(), new CommandContextBuilder<>(otherSource).build())

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

@@ -15,11 +15,11 @@ import static org.junit.Assert.assertThat;
 public abstract class AbstractCommandNodeTest {
     @Mock private Command command;
 
-    protected abstract CommandNode getCommandNode();
+    protected abstract CommandNode<Object> getCommandNode();
 
     @Test
     public void testAddChild() throws Exception {
-        CommandNode node = getCommandNode();
+        CommandNode<Object> node = getCommandNode();
 
         node.addChild(literal("child1").build());
         node.addChild(literal("child2").build());
@@ -30,7 +30,7 @@ public abstract class AbstractCommandNodeTest {
 
     @Test
     public void testAddChildMergesGrandchildren() throws Exception {
-        CommandNode node = getCommandNode();
+        CommandNode<Object> node = getCommandNode();
 
         node.addChild(literal("child").then(
             literal("grandchild1")
@@ -46,7 +46,7 @@ public abstract class AbstractCommandNodeTest {
 
     @Test
     public void testAddChildPreservesCommand() throws Exception {
-        CommandNode node = getCommandNode();
+        CommandNode<Object> node = getCommandNode();
 
         node.addChild(literal("child").executes(command).build());
         node.addChild(literal("child").build());
@@ -56,7 +56,7 @@ public abstract class AbstractCommandNodeTest {
 
     @Test
     public void testAddChildOverwritesCommand() throws Exception {
-        CommandNode node = getCommandNode();
+        CommandNode<Object> node = getCommandNode();
 
         node.addChild(literal("child").build());
         node.addChild(literal("child").executes(command).build());

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

@@ -7,28 +7,25 @@ import com.mojang.brigadier.Command;
 import com.mojang.brigadier.arguments.IntegerArgumentType;
 import com.mojang.brigadier.context.CommandContextBuilder;
 import com.mojang.brigadier.exceptions.CommandException;
-import com.mojang.brigadier.exceptions.CommandExceptionType;
 import org.junit.Before;
 import org.junit.Test;
 
-import java.util.Map;
 import java.util.Set;
 
 import static com.mojang.brigadier.arguments.IntegerArgumentType.integer;
 import static com.mojang.brigadier.builder.RequiredArgumentBuilder.argument;
 import static org.hamcrest.Matchers.empty;
-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;
 
 public class ArgumentCommandNodeTest extends AbstractCommandNodeTest {
-    private ArgumentCommandNode<Integer> node;
+    private ArgumentCommandNode<Object, Integer> node;
     private CommandContextBuilder<Object> contextBuilder;
 
     @Override
-    protected CommandNode getCommandNode() {
+    protected CommandNode<Object> getCommandNode() {
         return node;
     }
 
@@ -79,7 +76,7 @@ public class ArgumentCommandNodeTest extends AbstractCommandNodeTest {
 
     @Test
     public void testEquals() throws Exception {
-        Command command = mock(Command.class);
+        @SuppressWarnings("unchecked") Command<Object> command = (Command<Object>) mock(Command.class);
 
         new EqualsTester()
             .addEqualityGroup(

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

@@ -6,11 +6,9 @@ import com.google.common.testing.EqualsTester;
 import com.mojang.brigadier.Command;
 import com.mojang.brigadier.context.CommandContextBuilder;
 import com.mojang.brigadier.exceptions.CommandException;
-import com.mojang.brigadier.exceptions.CommandExceptionType;
 import org.junit.Before;
 import org.junit.Test;
 
-import java.util.Map;
 import java.util.Set;
 
 import static com.mojang.brigadier.builder.LiteralArgumentBuilder.literal;
@@ -22,11 +20,11 @@ import static org.junit.Assert.fail;
 import static org.mockito.Mockito.mock;
 
 public class LiteralCommandNodeTest extends AbstractCommandNodeTest {
-    private LiteralCommandNode node;
+    private LiteralCommandNode<Object> node;
     private CommandContextBuilder<Object> contextBuilder;
 
     @Override
-    protected CommandNode getCommandNode() {
+    protected CommandNode<Object> getCommandNode() {
         return node;
     }
 
@@ -94,7 +92,7 @@ public class LiteralCommandNodeTest extends AbstractCommandNodeTest {
 
     @Test
     public void testEquals() throws Exception {
-        Command command = mock(Command.class);
+        @SuppressWarnings("unchecked") Command<Object> command = mock(Command.class);
 
         new EqualsTester()
             .addEqualityGroup(

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

@@ -14,16 +14,16 @@ import static org.hamcrest.Matchers.is;
 import static org.junit.Assert.assertThat;
 
 public class RootCommandNodeTest extends AbstractCommandNodeTest {
-    private RootCommandNode node;
+    private RootCommandNode<Object> node;
 
     @Override
-    protected CommandNode getCommandNode() {
+    protected CommandNode<Object> getCommandNode() {
         return node;
     }
 
     @Before
     public void setUp() throws Exception {
-        node = new RootCommandNode();
+        node = new RootCommandNode<>();
     }
 
     @Test
@@ -33,7 +33,7 @@ public class RootCommandNodeTest extends AbstractCommandNodeTest {
 
     @Test(expected = UnsupportedOperationException.class)
     public void testAddChildNoRoot() throws Exception {
-        node.addChild(new RootCommandNode());
+        node.addChild(new RootCommandNode<>());
     }
 
     @Test
@@ -52,14 +52,14 @@ public class RootCommandNodeTest extends AbstractCommandNodeTest {
     public void testEquals() throws Exception {
         new EqualsTester()
             .addEqualityGroup(
-                new RootCommandNode(),
+                new RootCommandNode<>(),
-                new RootCommandNode()
+                new RootCommandNode<>()
             )
             .addEqualityGroup(
-                new RootCommandNode() {{
+                new RootCommandNode<Object>() {{
                     addChild(literal("foo").build());
                 }},
-                new RootCommandNode() {{
+                new RootCommandNode<Object>() {{
                     addChild(literal("foo").build());
                 }}
             )