From 928912de68fee8c6cdc4ce201d9575a819feb630 Mon Sep 17 00:00:00 2001
From: Nathan Adams <dinnerbone@dinnerbone.com>
Date: Wed, 17 Jan 2018 12:46:19 +0100
Subject: [PATCH] Catch redirect errors when forked

---
 .../mojang/brigadier/CommandDispatcher.java   | 19 +++++++++++++------
 .../brigadier/SingleRedirectModifier.java     | 11 +++++++++++
 .../brigadier/builder/ArgumentBuilder.java    |  5 ++---
 3 files changed, 26 insertions(+), 9 deletions(-)
 create mode 100644 src/main/java/com/mojang/brigadier/SingleRedirectModifier.java

diff --git a/src/main/java/com/mojang/brigadier/CommandDispatcher.java b/src/main/java/com/mojang/brigadier/CommandDispatcher.java
index f7f8bb5..2be3124 100644
--- a/src/main/java/com/mojang/brigadier/CommandDispatcher.java
+++ b/src/main/java/com/mojang/brigadier/CommandDispatcher.java
@@ -111,13 +111,20 @@ public class CommandDispatcher<S> {
                             }
                             next.add(child.copyFor(context.getSource()));
                         } else {
-                            final Collection<S> results = modifier.apply(context);
-                            if (!results.isEmpty()) {
-                                if (next == null) {
-                                    next = new ArrayList<>(results.size());
+                            try {
+                                final Collection<S> results = modifier.apply(context);
+                                if (!results.isEmpty()) {
+                                    if (next == null) {
+                                        next = new ArrayList<>(results.size());
+                                    }
+                                    for (final S source : results) {
+                                        next.add(child.copyFor(source));
+                                    }
                                 }
-                                for (final S source : results) {
-                                    next.add(child.copyFor(source));
+                            } catch (final CommandSyntaxException ex) {
+                                consumer.onCommandComplete(context, false, 0);
+                                if (!forked) {
+                                    throw ex;
                                 }
                             }
                         }
diff --git a/src/main/java/com/mojang/brigadier/SingleRedirectModifier.java b/src/main/java/com/mojang/brigadier/SingleRedirectModifier.java
new file mode 100644
index 0000000..2d1b55f
--- /dev/null
+++ b/src/main/java/com/mojang/brigadier/SingleRedirectModifier.java
@@ -0,0 +1,11 @@
+package com.mojang.brigadier;
+
+import com.mojang.brigadier.context.CommandContext;
+import com.mojang.brigadier.exceptions.CommandSyntaxException;
+
+import java.util.Collection;
+
+@FunctionalInterface
+public interface SingleRedirectModifier<S> {
+    S apply(CommandContext<S> context) throws CommandSyntaxException;
+}
diff --git a/src/main/java/com/mojang/brigadier/builder/ArgumentBuilder.java b/src/main/java/com/mojang/brigadier/builder/ArgumentBuilder.java
index 327b3e7..227053f 100644
--- a/src/main/java/com/mojang/brigadier/builder/ArgumentBuilder.java
+++ b/src/main/java/com/mojang/brigadier/builder/ArgumentBuilder.java
@@ -2,13 +2,12 @@ package com.mojang.brigadier.builder;
 
 import com.mojang.brigadier.Command;
 import com.mojang.brigadier.RedirectModifier;
-import com.mojang.brigadier.context.CommandContext;
+import com.mojang.brigadier.SingleRedirectModifier;
 import com.mojang.brigadier.tree.CommandNode;
 import com.mojang.brigadier.tree.RootCommandNode;
 
 import java.util.Collection;
 import java.util.Collections;
-import java.util.function.Function;
 import java.util.function.Predicate;
 
 public abstract class ArgumentBuilder<S, T extends ArgumentBuilder<S, T>> {
@@ -63,7 +62,7 @@ public abstract class ArgumentBuilder<S, T extends ArgumentBuilder<S, T>> {
         return forward(target, null, false);
     }
 
-    public T redirect(final CommandNode<S> target, final Function<CommandContext<S>, S> modifier) {
+    public T redirect(final CommandNode<S> target, final SingleRedirectModifier<S> modifier) {
         return forward(target, modifier == null ? null : o -> Collections.singleton(modifier.apply(o)), false);
     }