Overhauled usage generation

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

@@ -2,6 +2,7 @@ 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;
 import com.mojang.brigadier.builder.LiteralArgumentBuilder;
@@ -14,6 +15,9 @@ import com.mojang.brigadier.tree.CommandNode;
 import com.mojang.brigadier.tree.LiteralCommandNode;
 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;
@@ -41,7 +45,8 @@ public class CommandDispatcher<S> {
     };
 
     public void register(LiteralArgumentBuilder<S> command) {
-        root.addChild(command.build());
+        final LiteralCommandNode<S> build = command.build();
+        root.addChild(build);
     }
 
     public int execute(String input, S source) throws CommandException {
@@ -97,48 +102,88 @@ public class CommandDispatcher<S> {
         return new ParseResults<>(contextBuilder, command, errors);
     }
 
-    public String getUsage(String command, S source) throws CommandException {
+    public String[] getAllUsage(CommandNode<S> node, S source) {
-        final ParseResults<S> parse = parseNodes(root, command, new CommandContextBuilder<>(source));
+        final ArrayList<String> result = Lists.newArrayList();
-        if (parse.getContext().getNodes().isEmpty()) {
+        getAllUsage(node, source, result,  "");
-            throw ERROR_UNKNOWN_COMMAND.create();
+        return result.toArray(new String[result.size()]);
+    }
+
+    private void getAllUsage(CommandNode<S> node, S source, ArrayList<String> result, String prefix) {
+        if (!node.canUse(source)) {
+            return;
         }
 
-        CommandContext<S> context = parse.getContext().build();
+        if (node.getCommand() != null) {
-        CommandNode<S> base = Iterables.getLast(context.getNodes().keySet());
+            result.add(prefix);
-        List<CommandNode<S>> children = base.getChildren().stream().filter(hasCommand).collect(Collectors.toList());
-        boolean optional = base.getCommand() != null;
-
-        if (children.isEmpty()) {
-            return context.getInput();
         }
 
-        children.sort((o1, o2) -> ComparisonChain.start()
+        if (!node.getChildren().isEmpty()) {
-            .compareTrueFirst(o1 instanceof LiteralCommandNode, o2 instanceof LiteralCommandNode)
+            for (final CommandNode<S> child : node.getChildren()) {
-            .result());
+                getAllUsage(child, source, result, prefix.isEmpty() ? child.getUsageText() : prefix + ARGUMENT_SEPARATOR + child.getUsageText());
+            }
+        }
+    }
 
-        StringBuilder result = new StringBuilder(context.getInput());
+    public Map<CommandNode<S>, String> getSmartUsage(CommandNode<S> node, S source) {
-        result.append(ARGUMENT_SEPARATOR);
+        Map<CommandNode<S>, String> result = Maps.newLinkedHashMap();
-        if (optional) {
+
-            result.append(USAGE_OPTIONAL_OPEN);
+        final boolean optional = node.getCommand() != null;
-        } else if (children.size() > 1) {
+        for (CommandNode<S> child : node.getChildren()) {
-            result.append(USAGE_REQUIRED_OPEN);
+            String usage = getSmartUsage(child, source, optional, false);
+            if (usage != null) {
+                result.put(child, usage);
+            }
+        }
+        return result;
+    }
+
+    private String getSmartUsage(CommandNode<S> node, S source, boolean optional, boolean deep) {
+        if (!node.canUse(source)) {
+            return null;
         }
 
-        for (int i = 0; i < children.size(); i++) {
+        String self = optional ? USAGE_OPTIONAL_OPEN + node.getUsageText() + USAGE_OPTIONAL_CLOSE : node.getUsageText();
-            result.append(children.get(i).getUsageText());
+        boolean childOptional = node.getCommand() != null;
+        String open = childOptional ? USAGE_OPTIONAL_OPEN : USAGE_REQUIRED_OPEN;
+        String close = childOptional ? USAGE_OPTIONAL_CLOSE : USAGE_REQUIRED_CLOSE;
 
-            if (i < children.size() - 1) {
+        if (!deep) {
-                result.append(USAGE_OR);
+            final Collection<CommandNode<S>> children = node.getChildren().stream().filter(c -> c.canUse(source)).collect(Collectors.toList());
+            if (children.size() == 1) {
+                final String usage = getSmartUsage(children.iterator().next(), source, childOptional, childOptional);
+                if (usage != null) {
+                    return self + ARGUMENT_SEPARATOR + usage;
+                }
+            } else if (children.size() > 1) {
+                Set<String> childUsage = Sets.newLinkedHashSet();
+                for (final CommandNode<S> child : children) {
+                    final String usage = getSmartUsage(child, source, childOptional, true);
+                    if (usage != null) {
+                        childUsage.add(usage);
+                    }
+                }
+                if (childUsage.size() == 1) {
+                    final String usage = childUsage.iterator().next();
+                    return self + ARGUMENT_SEPARATOR + (childOptional ? USAGE_OPTIONAL_OPEN + usage + USAGE_OPTIONAL_CLOSE : usage);
+                } else if (childUsage.size() > 1) {
+                    StringBuilder builder = new StringBuilder(open);
+                    int count = 0;
+                    for (final CommandNode<S> child : children) {
+                        if (count > 0) {
+                            builder.append(USAGE_OR);
+                        }
+                        builder.append(child.getUsageText());
+                        count++;
+                    }
+                    if (count > 0) {
+                        builder.append(close);
+                        return self + ARGUMENT_SEPARATOR + builder.toString();
+                    }
+                }
             }
         }
 
-        if (optional) {
+        return self;
-            result.append(USAGE_OPTIONAL_CLOSE);
-        } else if (children.size() > 1) {
-            result.append(USAGE_REQUIRED_CLOSE);
-        }
-
-        return result.toString();
     }
 
     private Set<String> findSuggestions(CommandNode<S> node, String command, CommandContextBuilder<S> contextBuilder, Set<String> result) {

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

@@ -1,18 +1,23 @@
 package com.mojang.brigadier;
 
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Iterators;
+import com.google.common.collect.Maps;
 import com.mojang.brigadier.exceptions.CommandException;
-import com.mojang.brigadier.exceptions.CommandExceptionType;
+import com.mojang.brigadier.tree.CommandNode;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.runners.MockitoJUnitRunner;
 
-import java.util.Collections;
+import java.util.Map;
 
 import static com.mojang.brigadier.arguments.IntegerArgumentType.integer;
 import static com.mojang.brigadier.builder.LiteralArgumentBuilder.literal;
-import static com.mojang.brigadier.builder.RequiredArgumentBuilder.argument;
+import static org.hamcrest.Matchers.empty;
+import static org.hamcrest.Matchers.emptyArray;
+import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasToString;
 import static org.hamcrest.Matchers.is;
 import static org.junit.Assert.assertThat;
@@ -24,226 +29,144 @@ public class CommandDispatcherUsagesTest {
     @Mock
     private Object source;
     @Mock
-    private Command<Object>  command;
+    private Command<Object> command;
 
     @Before
     public void setUp() throws Exception {
         subject = new CommandDispatcher<>();
+        subject.register(
+            literal("a")
+                .then(
+                    literal("1")
+                        .then(literal("i").executes(command))
+                        .then(literal("ii").executes(command))
+                )
+                .then(
+                    literal("2")
+                        .then(literal("i").executes(command))
+                        .then(literal("ii").executes(command))
+                )
+        );
+        subject.register(literal("b").then(literal("1").executes(command)));
+        subject.register(literal("c").executes(command));
+        subject.register(literal("d").requires(s -> false).executes(command));
+        subject.register(
+            literal("e")
+                .executes(command)
+                .then(
+                    literal("1")
+                        .executes(command)
+                        .then(literal("i").executes(command))
+                        .then(literal("ii").executes(command))
+                )
+        );
+        subject.register(
+            literal("f")
+                .then(
+                    literal("1")
+                        .then(literal("i").executes(command))
+                        .then(literal("ii").executes(command).requires(s -> false))
+                )
+                .then(
+                    literal("2")
+                        .then(literal("i").executes(command).requires(s -> false))
+                        .then(literal("ii").executes(command))
+                )
+        );
+        subject.register(
+            literal("g")
+                .executes(command)
+                .then(literal("1").then(literal("i").executes(command)))
+        );
+        subject.register(
+            literal("h")
+                .executes(command)
+                .then(literal("1").then(literal("i").executes(command)))
+                .then(literal("2").then(literal("i").then(literal("ii").executes(command))))
+                .then(literal("3").executes(command))
+        );
+        subject.register(
+            literal("i")
+                .executes(command)
+                .then(literal("1").executes(command))
+                .then(literal("2").executes(command))
+        );
     }
 
-    @Test
+    private CommandNode<Object> get(String command) {
-    public void testUnknownCommand() throws Exception {
         try {
-            subject.getUsage("foo", source);
+            return Iterators.getLast(subject.parse(command, source).getContext().getNodes().keySet().iterator());
-            fail();
+        } catch (CommandException e) {
-        } catch (CommandException ex) {
+            throw new AssertionError("get() failed unexpectedly", e);
-            assertThat(ex.getType(), is(CommandDispatcher.ERROR_UNKNOWN_COMMAND));
-            assertThat(ex.getData(), is(Collections.<String, Object>emptyMap()));
         }
     }
 
     @Test
-    public void testNoCommand() throws Exception {
+    public void testAllUsage_noCommands() throws Exception {
-        try {
+        subject = new CommandDispatcher<>();
-            subject.getUsage("", source);
+        final String[] results = subject.getAllUsage(subject.getRoot(), source);
-            fail();
+        assertThat(results, is(emptyArray()));
-        } catch (CommandException ex) {
-            assertThat(ex.getType(), is(CommandDispatcher.ERROR_UNKNOWN_COMMAND));
-            assertThat(ex.getData(), is(Collections.<String, Object>emptyMap()));
-        }
     }
 
     @Test
-    public void testUnknownSubcommand() throws Exception {
+    public void testSmartUsage_noCommands() throws Exception {
-        subject.register(
+        subject = new CommandDispatcher<>();
-            literal("base").then(
+        final Map<CommandNode<Object>, String> results = subject.getSmartUsage(subject.getRoot(), source);
-                literal("foo").executes(command)
+        assertThat(results.entrySet(), is(empty()));
-            ).then(
-                literal("bar").then(
-                    literal("baz").executes(command)
-                ).then(
-                    literal("qux").then(
-                        literal("not_runnable")
-                    )
-                ).then(
-                    literal("quux").then(
-                        literal("corge").executes(command)
-                    )
-                ).executes(command)
-            ).executes(command)
-        );
-
-        assertThat(subject.getUsage("base unknown", source), hasToString("base [foo|bar]"));
     }
 
     @Test
-    public void testInaccessibleCommand() throws Exception {
+    public void testAllUsage_root() throws Exception {
-        subject.register(literal("foo").requires(s -> false));
+        final String[] results = subject.getAllUsage(subject.getRoot(), source);
-
+        assertThat(results, equalTo(new String[] {
-        try {
+            "a 1 i",
-            subject.getUsage("foo", source);
+            "a 1 ii",
-            fail();
+            "a 2 i",
-        } catch (CommandException ex) {
+            "a 2 ii",
-            assertThat(ex.getType(), is(CommandDispatcher.ERROR_UNKNOWN_COMMAND));
+            "b 1",
-            assertThat(ex.getData(), is(Collections.<String, Object>emptyMap()));
+            "c",
-        }
+            "e",
+            "e 1",
+            "e 1 i",
+            "e 1 ii",
+            "f 1 i",
+            "f 2 ii",
+            "g",
+            "g 1 i",
+            "h",
+            "h 1 i",
+            "h 2 i ii",
+            "h 3",
+            "i",
+            "i 1",
+            "i 2",
+        }));
     }
 
     @Test
-    public void testSubcommandUsage() throws Exception {
+    public void testSmartUsage_root() throws Exception {
-        subject.register(
+        final Map<CommandNode<Object>, String> results = subject.getSmartUsage(subject.getRoot(), source);
-            literal("base").then(
+        assertThat(results, equalTo(ImmutableMap.builder()
-                literal("foo").executes(command)
+                .put(get("a"), "a (1|2)")
-            ).then(
+                .put(get("b"), "b 1")
-                literal("bar").then(
+                .put(get("c"), "c")
-                    literal("baz").executes(command)
+                .put(get("e"), "e [1]")
-                ).then(
+                .put(get("f"), "f (1|2)")
-                    literal("qux").then(
+                .put(get("g"), "g [1]")
-                        literal("not_runnable")
+                .put(get("h"), "h [1|2|3]")
-                    )
+                .put(get("i"), "i [1|2]")
-                ).then(
+                .build()
-                    literal("quux").then(
+        ));
-                        literal("corge").executes(command)
-                    )
-                ).executes(command)
-            ).executes(command)
-        );
-
-        assertThat(subject.getUsage("base bar", source), hasToString("base bar [baz|quux]"));
     }
 
     @Test
-    public void testOptionalSingleLiteral() throws Exception {
+    public void testSmartUsage_h() throws Exception {
-        subject.register(
+        final Map<CommandNode<Object>, String> results = subject.getSmartUsage(get("h"), source);
-            literal("base").then(
+        assertThat(results, equalTo(ImmutableMap.builder()
-                literal("foo").executes(command)
+            .put(get("h 1"), "[1] i")
-            ).executes(command)
+            .put(get("h 2"), "[2] i ii")
-        );
+            .put(get("h 3"), "[3]")
-
+            .build()
-        assertThat(subject.getUsage("base", source), hasToString("base [foo]"));
+        ));
-    }
-
-    @Test
-    public void testNoArguments() throws Exception {
-        subject.register(
-            literal("base").executes(command)
-        );
-
-        assertThat(subject.getUsage("base", source), hasToString("base"));
-    }
-
-    @Test
-    public void testRequiredSingleLiteral() throws Exception {
-        subject.register(
-            literal("base").then(
-                literal("foo").executes(command)
-            )
-        );
-
-        assertThat(subject.getUsage("base", source), hasToString("base foo"));
-    }
-
-    @Test
-    public void testOptionalTwoLiterals() throws Exception {
-        subject.register(
-            literal("base").then(
-                literal("foo").executes(command)
-            ).then(
-                literal("bar").executes(command)
-            ).executes(command)
-        );
-
-        assertThat(subject.getUsage("base", source), hasToString("base [foo|bar]"));
-    }
-
-    @Test
-    public void testRequiredTwoLiterals() throws Exception {
-        subject.register(
-            literal("base").then(
-                literal("foo").executes(command)
-            ).then(
-                literal("bar").executes(command)
-            )
-        );
-
-        assertThat(subject.getUsage("base", source), hasToString("base (foo|bar)"));
-    }
-
-    @Test
-    public void testOptionalOneArgument() throws Exception {
-        subject.register(
-            literal("base").then(
-                argument("foo", integer()).executes(command)
-            ).executes(command)
-        );
-
-        assertThat(subject.getUsage("base", source), hasToString("base [<foo>]"));
-    }
-
-    @Test
-    public void testRequiredOneArgument() throws Exception {
-        subject.register(
-            literal("base").then(
-                argument("foo", integer()).executes(command)
-            )
-        );
-
-        assertThat(subject.getUsage("base", source), hasToString("base <foo>"));
-    }
-
-    @Test
-    public void testOptionalTwoArguments() throws Exception {
-        subject.register(
-            literal("base").then(
-                argument("foo", integer()).executes(command)
-            ).then(
-                argument("bar", integer()).executes(command)
-            ).executes(command)
-        );
-
-        assertThat(subject.getUsage("base", source), hasToString("base [<foo>|<bar>]"));
-    }
-
-    @Test
-    public void testRequiredTwoArguments() throws Exception {
-        subject.register(
-            literal("base").then(
-                argument("foo", integer()).executes(command)
-            ).then(
-                argument("bar", integer()).executes(command)
-            )
-        );
-
-        assertThat(subject.getUsage("base", source), hasToString("base (<foo>|<bar>)"));
-    }
-
-    @Test
-    public void testOptionalLiteralOrArgument() throws Exception {
-        subject.register(
-            literal("base").then(
-                literal("foo").executes(command)
-            ).then(
-                argument("bar", integer()).executes(command)
-            ).then(
-                literal("baz").executes(command)
-            ).executes(command)
-        );
-
-        assertThat(subject.getUsage("base", source), hasToString("base [foo|baz|<bar>]"));
-    }
-
-    @Test
-    public void testRequiredLiteralOrArgument() throws Exception {
-        subject.register(
-            literal("base").then(
-                literal("foo").executes(command)
-            ).then(
-                argument("bar", integer()).executes(command)
-            ).then(
-                literal("baz").executes(command)
-            )
-        );
-
-        assertThat(subject.getUsage("base", source), hasToString("base (foo|baz|<bar>)"));
     }
 }