Initial attempt at list/repeating arguments

Author: SoniEx2

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

@@ -11,15 +11,21 @@
 public class RequiredArgumentBuilder<S, T> extends ArgumentBuilder<S, RequiredArgumentBuilder<S, T>> {
     private final String name;
     private final ArgumentType<T> type;
+    private final boolean list;
     private SuggestionProvider<S> suggestionsProvider = null;
 
-    private RequiredArgumentBuilder(final String name, final ArgumentType<T> type) {
+    private RequiredArgumentBuilder(final String name, final ArgumentType<T> type, boolean list) {
         this.name = name;
         this.type = type;
+        this.list = list;
     }
 
     public static <S, T> RequiredArgumentBuilder<S, T> argument(final String name, final ArgumentType<T> type) {
-        return new RequiredArgumentBuilder<>(name, type);
+        return new RequiredArgumentBuilder<>(name, type, false);
+    }
+
+    public static <S, T> RequiredArgumentBuilder<S, T> list(final String name, final ArgumentType<T> type) {
+        return new RequiredArgumentBuilder<>(name, type, true);
     }
 
     public RequiredArgumentBuilder<S, T> suggests(final SuggestionProvider<S> provider) {
@@ -44,8 +50,10 @@ public String getName() {
         return name;
     }
 
+    public boolean isList() { return list; }
+
     public ArgumentCommandNode<S, T> build() {
-        final ArgumentCommandNode<S, T> result = new ArgumentCommandNode<>(getName(), getType(), getCommand(), getRequirement(), getRedirect(), getRedirectModifier(), isFork(), getSuggestionsProvider());
+        final ArgumentCommandNode<S, T> result = new ArgumentCommandNode<>(getName(), getType(), getCommand(), getRequirement(), getRedirect(), getRedirectModifier(), isFork(), getSuggestionsProvider(), isList());
 
         for (final CommandNode<S> argument : getArguments()) {
             result.addChild(argument);

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

@@ -11,12 +11,15 @@
 import com.mojang.brigadier.context.CommandContext;
 import com.mojang.brigadier.context.CommandContextBuilder;
 import com.mojang.brigadier.context.ParsedArgument;
+import com.mojang.brigadier.context.StringRange;
 import com.mojang.brigadier.exceptions.CommandSyntaxException;
 import com.mojang.brigadier.suggestion.SuggestionProvider;
 import com.mojang.brigadier.suggestion.Suggestions;
 import com.mojang.brigadier.suggestion.SuggestionsBuilder;
 
+import java.util.ArrayList;
 import java.util.Collection;
+import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.function.Predicate;
 
@@ -27,12 +30,21 @@ public class ArgumentCommandNode<S, T> extends CommandNode<S> {
     private final String name;
     private final ArgumentType<T> type;
     private final SuggestionProvider<S> customSuggestions;
+    private final boolean list;
+    private Collection<? extends CommandNode<S>> cached;
+    private Collection<? extends CommandNode<S>> cachedOut;
 
+    @Deprecated
     public ArgumentCommandNode(final String name, final ArgumentType<T> type, final Command<S> command, final Predicate<S> requirement, final CommandNode<S> redirect, final RedirectModifier<S> modifier, final boolean forks, final SuggestionProvider<S> customSuggestions) {
+        this(name, type, command, requirement, redirect, modifier, forks, customSuggestions, false);
+    }
+
+    public ArgumentCommandNode(final String name, final ArgumentType<T> type, final Command<S> command, final Predicate<S> requirement, final CommandNode<S> redirect, final RedirectModifier<S> modifier, final boolean forks, final SuggestionProvider<S> customSuggestions, boolean list) {
         super(command, requirement, redirect, modifier, forks);
         this.name = name;
         this.type = type;
         this.customSuggestions = customSuggestions;
+        this.list = list;
     }
 
     public ArgumentType<T> getType() {
@@ -57,10 +69,24 @@ public SuggestionProvider<S> getCustomSuggestions() {
     public void parse(final StringReader reader, final CommandContextBuilder<S> contextBuilder) throws CommandSyntaxException {
         final int start = reader.getCursor();
         final T result = type.parse(reader);
-        final ParsedArgument<S, T> parsed = new ParsedArgument<>(start, reader.getCursor(), result);
 
-        contextBuilder.withArgument(name, parsed);
-        contextBuilder.withNode(this, parsed.getRange());
+        if (!list) {
+            final ParsedArgument<S, T> parsed = new ParsedArgument<>(start, reader.getCursor(), result);
+
+            contextBuilder.withArgument(name, parsed);
+        } else {
+            ParsedArgument<S, List<T>> parsed = (ParsedArgument<S, List<T>>) contextBuilder.getArguments().get(name);
+            if (parsed == null) {
+                parsed = new ParsedArgument<>(start, reader.getCursor(), new ArrayList<>());
+            } else {
+                parsed = new ParsedArgument<>(parsed.getRange().getStart(), reader.getCursor(), parsed.getResult());
+            }
+            parsed.getResult().add(result);
+
+            contextBuilder.withArgument(name, parsed);
+        }
+
+        contextBuilder.withNode(this, new StringRange(start, reader.getCursor()));
     }
 
     @Override
@@ -74,7 +100,7 @@ public CompletableFuture<Suggestions> listSuggestions(final CommandContext<S> co
 
     @Override
     public RequiredArgumentBuilder<S, T> createBuilder() {
-        final RequiredArgumentBuilder<S, T> builder = RequiredArgumentBuilder.argument(name, type);
+        final RequiredArgumentBuilder<S, T> builder = list ? RequiredArgumentBuilder.list(name, type) : RequiredArgumentBuilder.argument(name, type);
         builder.requires(getRequirement());
         builder.forward(getRedirect(), getRedirectModifier(), isFork());
         builder.suggests(customSuggestions);
@@ -128,4 +154,24 @@ public Collection<String> getExamples() {
     public String toString() {
         return "<argument " + name + ":" + type +">";
     }
+
+    @Override
+    public Collection<? extends CommandNode<S>> getRelevantNodes(StringReader input) {
+        if (!list) {
+            return super.getRelevantNodes(input);
+        } else {
+            Collection<? extends CommandNode<S>> relevantNodes = super.getRelevantNodes(input);
+            if (relevantNodes.size() == 1 && relevantNodes.stream().allMatch(node -> node instanceof LiteralCommandNode)) {
+                return relevantNodes;
+            } else {
+                if (cached != relevantNodes) {
+                    cached = relevantNodes;
+                    ArrayList<CommandNode<S>> commandNodes = new ArrayList<>(cached);
+                    commandNodes.add(this);
+                    cachedOut = commandNodes;
+                }
+                return cachedOut;
+            }
+        }
+    }
 }