From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: Jake Potrebic <jake.m.potrebic@gmail.com>
Date: Sun, 22 Aug 2021 15:21:57 -0700
Subject: [PATCH] Fix entity type tags suggestions in selectors

This would really be better as a client fix because just to fix it
all EntityArguments have been told to ask the server for completions
when if this was fixed on the client, that wouldn't be needed.

Mojira Issue: https://bugs.mojang.com/browse/MC-235045

diff --git a/src/main/java/com/destroystokyo/paper/PaperConfig.java b/src/main/java/com/destroystokyo/paper/PaperConfig.java
index 609f7ff543353ede53dc46dafc5a2fd0b0622cd8..d5aa95846e7f52108a03e3731023527281b21d73 100644
--- a/src/main/java/com/destroystokyo/paper/PaperConfig.java
+++ b/src/main/java/com/destroystokyo/paper/PaperConfig.java
@@ -525,6 +525,11 @@ public class PaperConfig {
         itemValidationBookPageLength = getInt("settings.item-validation.book.page", itemValidationBookPageLength);
     }
 
+    public static boolean fixTargetSelectorTagCompletion = true;
+    private static void fixTargetSelectorTagCompletion() {
+        fixTargetSelectorTagCompletion = getBoolean("settings.fix-target-selector-tag-completion", fixTargetSelectorTagCompletion);
+    }
+
     public static final class PacketLimit {
         public final double packetLimitInterval;
         public final double maxPacketRate;
diff --git a/src/main/java/net/minecraft/commands/CommandSourceStack.java b/src/main/java/net/minecraft/commands/CommandSourceStack.java
index cb0045fc4ddd738c45dee89d57b213a633b9a136..098182d2426a25cef0bc285356bc346db0af8172 100644
--- a/src/main/java/net/minecraft/commands/CommandSourceStack.java
+++ b/src/main/java/net/minecraft/commands/CommandSourceStack.java
@@ -322,4 +322,20 @@ public class CommandSourceStack implements SharedSuggestionProvider, com.destroy
         return this.source.getBukkitSender(this);
     }
     // CraftBukkit end
+    // Paper start - override getSelectedEntities
+    @Override
+    public Collection<String> getSelectedEntities() {
+        if (com.destroystokyo.paper.PaperConfig.fixTargetSelectorTagCompletion && this.source instanceof ServerPlayer player) {
+            double pickDistance = player.gameMode.getGameModeForPlayer().isCreative() ? 5.0F : 4.5F;
+            Vec3 min = player.getEyePosition(1.0F);
+            Vec3 viewVector = player.getViewVector(1.0F);
+            Vec3 max = min.add(viewVector.x * pickDistance, viewVector.y * pickDistance, viewVector.z * pickDistance);
+            net.minecraft.world.phys.AABB aabb = player.getBoundingBox().expandTowards(viewVector.scale(pickDistance)).inflate(1.0D, 1.0D, 1.0D);
+            pickDistance = player.gameMode.getGameModeForPlayer().isCreative() ? 6.0F : pickDistance;
+            net.minecraft.world.phys.EntityHitResult hitResult = net.minecraft.world.entity.projectile.ProjectileUtil.getEntityHitResult(player, min, max, aabb, (e) -> !e.isSpectator() && e.isPickable(), pickDistance);
+            return hitResult != null ? java.util.Collections.singletonList(hitResult.getEntity().getStringUUID()) : SharedSuggestionProvider.super.getSelectedEntities();
+        }
+        return SharedSuggestionProvider.super.getSelectedEntities();
+    }
+    // Paper end
 }
diff --git a/src/main/java/net/minecraft/commands/Commands.java b/src/main/java/net/minecraft/commands/Commands.java
index f6b73f8c6638ddf79e45042f5c8902ea1f271f5c..1617437515590a32c42687d290dd11bc8fa8edf5 100644
--- a/src/main/java/net/minecraft/commands/Commands.java
+++ b/src/main/java/net/minecraft/commands/Commands.java
@@ -407,6 +407,7 @@ public class Commands {
     private void fillUsableCommands(CommandNode<CommandSourceStack> tree, CommandNode<SharedSuggestionProvider> result, CommandSourceStack source, Map<CommandNode<CommandSourceStack>, CommandNode<SharedSuggestionProvider>> resultNodes) {
         Iterator iterator = tree.getChildren().iterator();
 
+        boolean registeredAskServerSuggestionsForTree = false; // Paper - tell clients to ask server for suggestions for EntityArguments
         while (iterator.hasNext()) {
             CommandNode<CommandSourceStack> commandnode2 = (CommandNode) iterator.next();
             if ( !org.spigotmc.SpigotConfig.sendNamespaced && commandnode2.getName().contains( ":" ) ) continue; // Spigot
@@ -428,6 +429,12 @@ public class Commands {
 
                     if (requiredargumentbuilder.getSuggestionsProvider() != null) {
                         requiredargumentbuilder.suggests(SuggestionProviders.safelySwap(requiredargumentbuilder.getSuggestionsProvider()));
+                        // Paper start - tell clients to ask server for suggestions for EntityArguments
+                        registeredAskServerSuggestionsForTree = requiredargumentbuilder.getSuggestionsProvider() == net.minecraft.commands.synchronization.SuggestionProviders.ASK_SERVER;
+                    } else if (com.destroystokyo.paper.PaperConfig.fixTargetSelectorTagCompletion && !registeredAskServerSuggestionsForTree && requiredargumentbuilder.getType() instanceof net.minecraft.commands.arguments.EntityArgument) {
+                        requiredargumentbuilder.suggests(requiredargumentbuilder.getType()::listSuggestions);
+                        registeredAskServerSuggestionsForTree = true; // You can only
+                        // Paper end - tell clients to ask server for suggestions for EntityArguments
                     }
                 }
 
diff --git a/src/main/java/net/minecraft/commands/arguments/EntityArgument.java b/src/main/java/net/minecraft/commands/arguments/EntityArgument.java
index 1f3076e59bac23d428c747ae12619e4b4e5fdd5a..1d23d05d7028c5f820f172cc54153f56848e1d05 100644
--- a/src/main/java/net/minecraft/commands/arguments/EntityArgument.java
+++ b/src/main/java/net/minecraft/commands/arguments/EntityArgument.java
@@ -127,7 +127,7 @@ public class EntityArgument implements ArgumentType<EntitySelector> {
 
             stringreader.setCursor(suggestionsbuilder.getStart());
             SharedSuggestionProvider icompletionprovider = (SharedSuggestionProvider) commandcontext.getSource();
-            EntitySelectorParser argumentparserselector = new EntitySelectorParser(stringreader, icompletionprovider.hasPermission(2));
+            EntitySelectorParser argumentparserselector = new EntitySelectorParser(stringreader, icompletionprovider.hasPermission(2), true); // Paper
 
             try {
                 argumentparserselector.parse();
diff --git a/src/main/java/net/minecraft/commands/arguments/selector/EntitySelectorParser.java b/src/main/java/net/minecraft/commands/arguments/selector/EntitySelectorParser.java
index 466d1d8d80df028ff842ce21be198be6a1c77b42..015d01242a9e8e7c6ef5b6bbf1b6d6ad0c8f36ca 100644
--- a/src/main/java/net/minecraft/commands/arguments/selector/EntitySelectorParser.java
+++ b/src/main/java/net/minecraft/commands/arguments/selector/EntitySelectorParser.java
@@ -115,12 +115,19 @@ public class EntitySelectorParser {
     private boolean hasScores;
     private boolean hasAdvancements;
     private boolean usesSelectors;
+    public boolean parsingEntityArgumentSuggestions; // Paper - track when parsing EntityArgument suggestions
 
     public EntitySelectorParser(StringReader reader) {
         this(reader, true);
     }
 
     public EntitySelectorParser(StringReader reader, boolean atAllowed) {
+        // Paper start
+        this(reader, atAllowed, false);
+    }
+    public EntitySelectorParser(StringReader reader, boolean atAllowed, boolean parsingEntityArgumentSuggestions) {
+        this.parsingEntityArgumentSuggestions = parsingEntityArgumentSuggestions;
+        // Paper end
         this.distance = MinMaxBounds.Doubles.ANY;
         this.level = MinMaxBounds.Ints.ANY;
         this.rotX = WrappedMinMaxBounds.ANY;
diff --git a/src/main/java/net/minecraft/commands/arguments/selector/options/EntitySelectorOptions.java b/src/main/java/net/minecraft/commands/arguments/selector/options/EntitySelectorOptions.java
index 061181381e4eabad5fa0122f049c4ce05996ffd2..90e023be5c38a038f1c03141ef4325abb25fd615 100644
--- a/src/main/java/net/minecraft/commands/arguments/selector/options/EntitySelectorOptions.java
+++ b/src/main/java/net/minecraft/commands/arguments/selector/options/EntitySelectorOptions.java
@@ -69,6 +69,19 @@ public class EntitySelectorOptions {
     public static final DynamicCommandExceptionType ERROR_ENTITY_TYPE_INVALID = new DynamicCommandExceptionType((entity) -> {
         return new TranslatableComponent("argument.entity.options.type.invalid", entity);
     });
+    // Paper start
+    public static final DynamicCommandExceptionType ERROR_ENTITY_TAG_INVALID = new DynamicCommandExceptionType((object) -> {
+        return io.papermc.paper.adventure.PaperAdventure
+            .asVanilla(net.kyori.adventure.text.Component
+                .text("Invalid or unknown entity type tag '" + object + "'")
+                .hoverEvent(net.kyori.adventure.text.event.HoverEvent
+                    .showText(net.kyori.adventure.text.Component
+                        .text("You can disable this error in 'paper.yml'")
+                    )
+                )
+            );
+    });
+    // Paper end
 
     private static void register(String id, EntitySelectorOptions.Modifier handler, Predicate<EntitySelectorParser> condition, Component description) {
         OPTIONS.put(id, new EntitySelectorOptions.Option(handler, condition, description));
@@ -316,8 +329,20 @@ public class EntitySelectorOptions {
 
                     if (reader.isTag()) {
                         ResourceLocation resourceLocation = ResourceLocation.read(reader.getReader());
+                        // Paper start - throw error if invalid entity tag (only on suggestions to keep cmd success behavior)
+                        final net.minecraft.tags.Tag<EntityType<?>> tag;
+                        if (com.destroystokyo.paper.PaperConfig.fixTargetSelectorTagCompletion && reader.parsingEntityArgumentSuggestions) {
+                            tag = EntityTypeTags.getAllTags().getTag(resourceLocation);
+                        } else {
+                            tag = EntityTypeTags.getAllTags().getTagOrEmpty(resourceLocation);
+                        }
+                        if (tag == null) {
+                            reader.getReader().setCursor(i);
+                            throw ERROR_ENTITY_TAG_INVALID.createWithContext(reader.getReader(), resourceLocation.toString());
+                        }
+                        // Paper end
                         reader.addPredicate((entity) -> {
-                            return entity.getType().is(entity.getServer().getTags().getOrEmpty(Registry.ENTITY_TYPE_REGISTRY).getTagOrEmpty(resourceLocation)) != bl;
+                            return entity.getType().is(tag) != bl; // Paper
                         });
                     } else {
                         ResourceLocation resourceLocation2 = ResourceLocation.read(reader.getReader());
