Added methods for getting path by node & nodes by path

SoniEx2 · Nov 13, 2017 · 1ce9894 · 1ce9894
1 parent 6b57d5f
commit 1ce9894
Show file tree

Hide file tree

Showing 7 changed files with 86 additions and 17 deletions.
diff --git a/src/main/java/com/mojang/brigadier/CommandDispatcher.java b/src/main/java/com/mojang/brigadier/CommandDispatcher.java
@@ -354,4 +354,44 @@ public CompletableFuture<CommandSuggestions> getCompletionSuggestions(final Pars
     public RootCommandNode<S> getRoot() {
         return root;
     }
+
+    public Collection<String> getPath(final CommandNode<S> target) {
+        final List<List<CommandNode<S>>> nodes = new ArrayList<>();
+        addPaths(root, nodes, new ArrayList<>());
+
+        for (final List<CommandNode<S>> list : nodes) {
+            if (list.get(list.size() - 1) == target) {
+                final List<String> result = new ArrayList<>(list.size());
+                for (final CommandNode<S> node : list) {
+                    if (node != root) {
+                        result.add(node.getName());
+                    }
+                }
+                return result;
+            }
+        }
+
+        return Collections.emptyList();
+    }
+
+    public CommandNode<S> findNode(final Collection<String> path) {
+        CommandNode<S> node = root;
+        for (final String name : path) {
+            node = node.getChild(name);
+            if (node == null) {
+                return null;
+            }
+        }
+        return node;
+    }
+
+    private void addPaths(final CommandNode<S> node, final List<List<CommandNode<S>>> result, final List<CommandNode<S>> parents) {
+        final List<CommandNode<S>> current = new ArrayList<>(parents);
+        current.add(node);
+        result.add(current);
+
+        for (final CommandNode<S> child : node.getChildren()) {
+            addPaths(child, result, current);
+        }
+    }
 }
diff --git a/src/main/java/com/mojang/brigadier/builder/ArgumentBuilder.java b/src/main/java/com/mojang/brigadier/builder/ArgumentBuilder.java
@@ -7,8 +7,6 @@
 
 import java.util.Collection;
 import java.util.Collections;
-import java.util.concurrent.CompletableFuture;
-import java.util.function.Function;
 import java.util.function.Predicate;
 
 public abstract class ArgumentBuilder<S, T extends ArgumentBuilder<S, T>> {
@@ -28,6 +26,14 @@ public T then(final ArgumentBuilder<S, ?> argument) {
         return getThis();
     }
 
+    public T then(final CommandNode<S> argument) {
+        if (target != null) {
+            throw new IllegalStateException("Cannot add children to a redirected node");
+        }
+        arguments.addChild(argument);
+        return getThis();
+    }
+
     public Collection<CommandNode<S>> getArguments() {
         return arguments.getChildren();
     }

diff --git a/src/main/java/com/mojang/brigadier/tree/ArgumentCommandNode.java b/src/main/java/com/mojang/brigadier/tree/ArgumentCommandNode.java
@@ -13,7 +13,6 @@
 
 import java.util.Collection;
 import java.util.concurrent.CompletableFuture;
-import java.util.function.Function;
 import java.util.function.Predicate;
 
 public class ArgumentCommandNode<S, T> extends CommandNode<S> {
@@ -31,16 +30,12 @@ public ArgumentCommandNode(final String name, final ArgumentType<T> type, final
         this.suggestionsProvider = suggestionsProvider;
     }
 
-    public String getName() {
-        return name;
-    }
-
     public ArgumentType<T> getType() {
         return type;
     }
 
     @Override
-    protected Object getMergeKey() {
+    public String getName() {
         return name;
     }
 

diff --git a/src/main/java/com/mojang/brigadier/tree/CommandNode.java b/src/main/java/com/mojang/brigadier/tree/CommandNode.java
@@ -18,7 +18,7 @@
 import java.util.stream.Collectors;
 
 public abstract class CommandNode<S> implements Comparable<CommandNode<S>> {
-    private Map<Object, CommandNode<S>> children = Maps.newLinkedHashMap();
+    private Map<String, CommandNode<S>> children = Maps.newLinkedHashMap();
     private final Predicate<S> requirement;
     private final CommandNode<S> redirect;
     private final RedirectModifier<S> modifier;
@@ -39,6 +39,10 @@ public Collection<CommandNode<S>> getChildren() {
         return children.values();
     }
 
+    public CommandNode<S> getChild(final String name) {
+        return children.get(name);
+    }
+
     public CommandNode<S> getRedirect() {
         return redirect;
     }
@@ -52,7 +56,11 @@ public boolean canUse(final S source) {
     }
 
     public void addChild(final CommandNode<S> node) {
-        final CommandNode<S> child = children.get(node.getMergeKey());
+        if (node instanceof RootCommandNode) {
+            throw new UnsupportedOperationException("Cannot add a RootCommandNode as a child to any other CommandNode");
+        }
+
+        final CommandNode<S> child = children.get(node.getName());
         if (child != null) {
             // We've found something to merge onto
             if (node.getCommand() != null) {
@@ -62,7 +70,7 @@ public void addChild(final CommandNode<S> node) {
                 child.addChild(grandchild);
             }
         } else {
-            children.put(node.getMergeKey(), node);
+            children.put(node.getName(), node);
         }
 
         children = children.entrySet().stream().sorted(Map.Entry.comparingByValue()).collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue, (e1, e2) -> e1, LinkedHashMap::new));
@@ -90,7 +98,7 @@ public Predicate<S> getRequirement() {
         return requirement;
     }
 
-    protected abstract Object getMergeKey();
+    public abstract String getName();
 
     public abstract String getUsageText();
 

diff --git a/src/main/java/com/mojang/brigadier/tree/LiteralCommandNode.java b/src/main/java/com/mojang/brigadier/tree/LiteralCommandNode.java
@@ -30,7 +30,7 @@ public String getLiteral() {
     }
 
     @Override
-    protected Object getMergeKey() {
+    public String getName() {
         return literal;
     }
 

diff --git a/src/main/java/com/mojang/brigadier/tree/RootCommandNode.java b/src/main/java/com/mojang/brigadier/tree/RootCommandNode.java
@@ -16,8 +16,8 @@ public RootCommandNode() {
     }
 
     @Override
-    protected Object getMergeKey() {
-        throw new UnsupportedOperationException("Cannot add a RootCommandNode as a child to any other CommandNode");
+    public String getName() {
+        return "";
     }
 
     @Override

diff --git a/src/test/java/com/mojang/brigadier/CommandDispatcherTest.java b/src/test/java/com/mojang/brigadier/CommandDispatcherTest.java
@@ -11,9 +11,7 @@
 import org.mockito.Mock;
 import org.mockito.runners.MockitoJUnitRunner;
 
-import java.util.Collection;
 import java.util.Collections;
-import java.util.function.Function;
 
 import static com.mojang.brigadier.arguments.IntegerArgumentType.integer;
 import static com.mojang.brigadier.builder.LiteralArgumentBuilder.literal;
@@ -22,6 +20,7 @@
 import static org.hamcrest.Matchers.hasProperty;
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.notNullValue;
+import static org.hamcrest.Matchers.nullValue;
 import static org.junit.Assert.assertThat;
 import static org.junit.Assert.fail;
 import static org.mockito.Matchers.argThat;
@@ -357,4 +356,25 @@ public void testExecuteInvalidSubcommand() throws Exception {
             assertThat(ex.getCursor(), is(4));
         }
     }
+
+    @Test
+    public void testGetPath() {
+        final LiteralCommandNode<Object> bar = literal("bar").build();
+        subject.register(literal("foo").then(bar));
+
+        assertThat(subject.getPath(bar), equalTo(Lists.newArrayList("foo", "bar")));
+    }
+
+    @Test
+    public void testFindNodeExists() {
+        final LiteralCommandNode<Object> bar = literal("bar").build();
+        subject.register(literal("foo").then(bar));
+
+        assertThat(subject.findNode(Lists.newArrayList("foo", "bar")), is(bar));
+    }
+
+    @Test
+    public void testFindNodeDoesntExist() {
+        assertThat(subject.findNode(Lists.newArrayList("foo", "bar")), is(nullValue()));
+    }
 }