Paper/patches/server/0729-Add-paper-mobcaps-and-paper-playermobcaps.patch

From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: Jason Penilla <11360596+jpenilla@users.noreply.github.com>
Date: Mon, 16 Aug 2021 01:31:54 -0500
Subject: [PATCH] Add '/paper mobcaps' and '/paper playermobcaps'

Add commands to get the mobcaps for a world, as well as the mobcaps for
each player when per-player mob spawning is enabled.

Also has a hover text on each mob category listing what entity types are
in said category

diff --git a/src/main/java/com/destroystokyo/paper/PaperCommand.java b/src/main/java/com/destroystokyo/paper/PaperCommand.java
index 27c350c123b4349b8b401a06f204de46c88da7ac..008bb8896657892fcaf64e134684cc49e62f23f6 100644
--- a/src/main/java/com/destroystokyo/paper/PaperCommand.java
+++ b/src/main/java/com/destroystokyo/paper/PaperCommand.java
@@ -3,6 +3,7 @@ package com.destroystokyo.paper;
 import com.destroystokyo.paper.io.SyncLoadFinder;
 import com.google.common.base.Functions;
 import com.google.common.base.Joiner;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
@@ -12,6 +13,12 @@ import net.minecraft.resources.ResourceLocation;
 import com.google.gson.JsonObject;
 import com.google.gson.internal.Streams;
 import com.google.gson.stream.JsonWriter;
+import net.kyori.adventure.text.Component;
+import net.kyori.adventure.text.ComponentLike;
+import net.kyori.adventure.text.JoinConfiguration;
+import net.kyori.adventure.text.TextComponent;
+import net.kyori.adventure.text.format.NamedTextColor;
+import net.kyori.adventure.text.format.TextColor;
 import net.minecraft.server.MinecraftServer;
 import net.minecraft.server.level.ChunkHolder;
 import net.minecraft.server.level.ServerChunkCache;
@@ -19,8 +26,10 @@ import net.minecraft.server.level.ServerLevel;
 import net.minecraft.server.level.ServerPlayer;
 import net.minecraft.server.level.ThreadedLevelLightEngine;
 import net.minecraft.world.entity.EntityType;
+import net.minecraft.world.entity.MobCategory;
 import net.minecraft.world.level.ChunkPos;
 import net.minecraft.server.MCUtil;
+import net.minecraft.world.level.NaturalSpawner;
 import org.apache.commons.lang3.tuple.MutablePair;
 import org.apache.commons.lang3.tuple.Pair;
 import org.bukkit.Bukkit;
@@ -52,6 +61,7 @@ import java.util.List;
 import java.util.Locale;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.ToIntFunction;
 import java.util.stream.Collectors;
 
 import static net.kyori.adventure.text.Component.text;
@@ -64,7 +74,7 @@ import static net.kyori.adventure.text.format.NamedTextColor.YELLOW;
 
 public class PaperCommand extends Command {
     private static final String BASE_PERM = "bukkit.command.paper.";
-    private static final ImmutableSet<String> SUBCOMMANDS = ImmutableSet.<String>builder().add("heap", "entity", "reload", "version", "debug", "chunkinfo", "fixlight", "syncloadinfo", "dumpitem").build();
+    private static final ImmutableSet<String> SUBCOMMANDS = ImmutableSet.<String>builder().add("heap", "entity", "reload", "version", "debug", "chunkinfo", "fixlight", "syncloadinfo", "dumpitem", "mobcaps", "playermobcaps").build();
 
     public PaperCommand(String name) {
         super(name);
@@ -97,6 +107,10 @@ public class PaperCommand extends Command {
                     return getListMatchingLast(sender, args, "help", "chunks");
                 }
                 break;
+            case "mobcaps":
+                return getListMatchingLast(sender, args, this.suggestMobcaps(sender, args));
+            case "playermobcaps":
+                return getListMatchingLast(sender, args, this.suggestPlayerMobcaps(sender, args));
             case "chunkinfo":
                 List<String> worldNames = new ArrayList<>();
                 worldNames.add("*");
@@ -193,6 +207,12 @@ public class PaperCommand extends Command {
             case "syncloadinfo":
                 this.doSyncLoadInfo(sender, args);
                 break;
+            case "mobcaps":
+                this.printMobcaps(sender, args);
+                break;
+            case "playermobcaps":
+                this.printPlayerMobcaps(sender, args);
+                break;
             case "ver":
                 if (!testPermission(sender, "version")) break; // "ver" needs a special check because it's an alias. All other commands are checked up before the switch statement (because they are present in the SUBCOMMANDS set)
             case "version":
@@ -251,6 +271,184 @@ public class PaperCommand extends Command {
         }
     }
 
+    public static final Map<MobCategory, TextColor> MOB_CATEGORY_COLORS = ImmutableMap.<MobCategory, TextColor>builder()
+        .put(MobCategory.MONSTER, NamedTextColor.RED)
+        .put(MobCategory.CREATURE, NamedTextColor.GREEN)
+        .put(MobCategory.AMBIENT, NamedTextColor.GRAY)
+        .put(MobCategory.AXOLOTLS, TextColor.color(0x7324FF))
+        .put(MobCategory.UNDERGROUND_WATER_CREATURE, TextColor.color(0x3541E6))
+        .put(MobCategory.WATER_CREATURE, TextColor.color(0x006EFF))
+        .put(MobCategory.WATER_AMBIENT, TextColor.color(0x00B3FF))
+        .put(MobCategory.MISC, TextColor.color(0x636363))
+        .build();
+
+    private List<String> suggestMobcaps(CommandSender sender, String[] args) {
+        if (args.length == 2) {
+            final List<String> worlds = new ArrayList<>(Bukkit.getWorlds().stream().map(World::getName).toList());
+            worlds.add("*");
+            return worlds;
+        }
+
+        return Collections.emptyList();
+    }
+
+    private List<String> suggestPlayerMobcaps(CommandSender sender, String[] args) {
+        if (args.length == 2) {
+            final List<String> list = new ArrayList<>();
+            for (final Player player : Bukkit.getOnlinePlayers()) {
+                if (!(sender instanceof Player senderPlayer) || senderPlayer.canSee(player)) {
+                    list.add(player.getName());
+                }
+            }
+            return list;
+        }
+
+        return Collections.emptyList();
+    }
+
+    private void printMobcaps(CommandSender sender, String[] args) {
+        final List<World> worlds;
+        if (args.length == 1) {
+            if (sender instanceof Player player) {
+                worlds = List.of(player.getWorld());
+            } else {
+                sender.sendMessage(Component.text("Must specify a world! ex: '/paper mobcaps world'", NamedTextColor.RED));
+                return;
+            }
+        } else if (args.length == 2) {
+            final String input = args[1];
+            if (input.equals("*")) {
+                worlds = Bukkit.getWorlds();
+            } else {
+                final World world = Bukkit.getWorld(input);
+                if (world == null) {
+                    sender.sendMessage(Component.text("'" + input + "' is not a valid world!", NamedTextColor.RED));
+                    return;
+                } else {
+                    worlds = List.of(world);
+                }
+            }
+        } else {
+            sender.sendMessage(Component.text("Too many arguments!", NamedTextColor.RED));
+            return;
+        }
+
+        for (final World world : worlds) {
+            final ServerLevel level = ((CraftWorld) world).getHandle();
+            final NaturalSpawner.SpawnState state = level.getChunkSource().getLastSpawnState();
+
+            final int chunks;
+            if (state == null) {
+                chunks = 0;
+            } else {
+                chunks = state.getSpawnableChunkCount();
+            }
+            sender.sendMessage(Component.join(JoinConfiguration.noSeparators(),
+                Component.text("Mobcaps for world: "),
+                Component.text(world.getName(), NamedTextColor.AQUA),
+                Component.text(" (" + chunks + " spawnable chunks)")
+            ));
+
+            sender.sendMessage(this.buildMobcapsComponent(
+                category -> {
+                    if (state == null) {
+                        return 0;
+                    } else {
+                        return state.getMobCategoryCounts().getOrDefault(category, 0);
+                    }
+                },
+                category -> NaturalSpawner.globalLimitForCategory(level, category, chunks)
+            ));
+        }
+    }
+
+    private void printPlayerMobcaps(CommandSender sender, String[] args) {
+        final Player player;
+        if (args.length == 1) {
+            if (sender instanceof Player pl) {
+                player = pl;
+            } else {
+                sender.sendMessage(Component.text("Must specify a player! ex: '/paper playermobcount playerName'", NamedTextColor.RED));
+                return;
+            }
+        } else if (args.length == 2) {
+            final String input = args[1];
+            player = Bukkit.getPlayerExact(input);
+            if (player == null) {
+                sender.sendMessage(Component.text("Could not find player named '" + input + "'", NamedTextColor.RED));
+                return;
+            }
+        } else {
+            sender.sendMessage(Component.text("Too many arguments!", NamedTextColor.RED));
+            return;
+        }
+
+        final ServerPlayer serverPlayer = ((CraftPlayer) player).getHandle();
+        final ServerLevel level = serverPlayer.getLevel();
+
+        if (!level.paperConfig().entities.spawning.perPlayerMobSpawns) {
+            sender.sendMessage(Component.text("Use '/paper mobcaps' for worlds where per-player mob spawning is disabled.", NamedTextColor.RED));
+            return;
+        }
+
+        sender.sendMessage(Component.join(JoinConfiguration.noSeparators(), Component.text("Mobcaps for player: "), Component.text(player.getName(), NamedTextColor.GREEN)));
+        sender.sendMessage(this.buildMobcapsComponent(
+            category -> level.chunkSource.chunkMap.getMobCountNear(serverPlayer, category),
+            category -> level.getWorld().getSpawnLimitUnsafe(org.bukkit.craftbukkit.util.CraftSpawnCategory.toBukkit(category))
+        ));
+    }
+
+    private Component buildMobcapsComponent(final ToIntFunction<MobCategory> countGetter, final ToIntFunction<MobCategory> limitGetter) {
+        return MOB_CATEGORY_COLORS.entrySet().stream()
+            .map(entry -> {
+                final MobCategory category = entry.getKey();
+                final TextColor color = entry.getValue();
+
+                final Component categoryHover = Component.join(JoinConfiguration.noSeparators(),
+                    Component.text("Entity types in category ", TextColor.color(0xE0E0E0)),
+                    Component.text(category.getName(), color),
+                    Component.text(':', NamedTextColor.GRAY),
+                    Component.newline(),
+                    Component.newline(),
+                    Registry.ENTITY_TYPE.entrySet().stream()
+                        .filter(it -> it.getValue().getCategory() == category)
+                        .map(it -> Component.translatable(it.getValue().getDescriptionId()))
+                        .collect(Component.toComponent(Component.text(", ", NamedTextColor.GRAY)))
+                );
+
+                final Component categoryComponent = Component.text()
+                    .content("  " + category.getName())
+                    .color(color)
+                    .hoverEvent(categoryHover)
+                    .build();
+
+                final TextComponent.Builder builder = Component.text()
+                    .append(
+                        categoryComponent,
+                        Component.text(": ", NamedTextColor.GRAY)
+                    );
+                final int limit = limitGetter.applyAsInt(category);
+                if (limit != -1) {
+                    builder.append(
+                        Component.text(countGetter.applyAsInt(category)),
+                        Component.text("/", NamedTextColor.GRAY),
+                        Component.text(limit)
+                    );
+                } else {
+                    builder.append(Component.text()
+                        .append(
+                            Component.text('n'),
+                            Component.text("/", NamedTextColor.GRAY),
+                            Component.text('a')
+                        )
+                        .hoverEvent(Component.text("This category does not naturally spawn.")));
+                }
+                return builder;
+            })
+            .map(ComponentLike::asComponent)
+            .collect(Component.toComponent(Component.newline()));
+    }
+
     private void doChunkInfo(CommandSender sender, String[] args) {
         List<org.bukkit.World> worlds;
         if (args.length < 2 || args[1].equals("*")) {
diff --git a/src/main/java/net/minecraft/world/level/NaturalSpawner.java b/src/main/java/net/minecraft/world/level/NaturalSpawner.java
index fa23e9c476d4edc6176d8b8a6cb13c52d2f66a87..4150e8cd7197eac53042d56f0a53a4951f8824ce 100644
--- a/src/main/java/net/minecraft/world/level/NaturalSpawner.java
+++ b/src/main/java/net/minecraft/world/level/NaturalSpawner.java
@@ -191,6 +191,16 @@ public final class NaturalSpawner {
         world.getProfiler().pop();
     }
 
+    // Paper start
+    public static int globalLimitForCategory(final ServerLevel level, final MobCategory category, final int spawnableChunkCount) {
+        final int categoryLimit = level.getWorld().getSpawnLimitUnsafe(CraftSpawnCategory.toBukkit(category));
+        if (categoryLimit < 1) {
+            return categoryLimit;
+        }
+        return categoryLimit * spawnableChunkCount / NaturalSpawner.MAGIC_NUMBER;
+    }
+    // Paper end
+
     public static void spawnCategoryForChunk(MobCategory group, ServerLevel world, LevelChunk chunk, NaturalSpawner.SpawnPredicate checker, NaturalSpawner.AfterSpawnCallback runner) {
         // Paper start - add parameters and int ret type
         spawnCategoryForChunk(group, world, chunk, checker, runner, Integer.MAX_VALUE, null);
diff --git a/src/main/java/org/bukkit/craftbukkit/CraftServer.java b/src/main/java/org/bukkit/craftbukkit/CraftServer.java
index 8417136c411c91e4c163446d7e1c2ced5f0084e1..874f317ebae64ba9f42573e3a412dedbeeb90cc7 100644
--- a/src/main/java/org/bukkit/craftbukkit/CraftServer.java
+++ b/src/main/java/org/bukkit/craftbukkit/CraftServer.java
@@ -2154,6 +2154,11 @@ public final class CraftServer implements Server {
 
     @Override
     public int getSpawnLimit(SpawnCategory spawnCategory) {
+        // Paper start
+        return this.getSpawnLimitUnsafe(spawnCategory);
+    }
+    public int getSpawnLimitUnsafe(final SpawnCategory spawnCategory) {
+        // Paper end
         return this.spawnCategoryLimit.getOrDefault(spawnCategory, -1);
     }
 
diff --git a/src/main/java/org/bukkit/craftbukkit/CraftWorld.java b/src/main/java/org/bukkit/craftbukkit/CraftWorld.java
index 9d3f7f196c4331662c4c78cac0b047bcd2ff5e77..21bf39ef5b20ecc989881b07c4e6b90c68540afd 100644
--- a/src/main/java/org/bukkit/craftbukkit/CraftWorld.java
+++ b/src/main/java/org/bukkit/craftbukkit/CraftWorld.java
@@ -1712,9 +1712,14 @@ public class CraftWorld extends CraftRegionAccessor implements World {
         Validate.notNull(spawnCategory, "SpawnCategory cannot be null");
         Validate.isTrue(CraftSpawnCategory.isValidForLimits(spawnCategory), "SpawnCategory." + spawnCategory + " are not supported.");
 
+        // Paper start
+        return this.getSpawnLimitUnsafe(spawnCategory);
+    }
+    public final int getSpawnLimitUnsafe(final SpawnCategory spawnCategory) {
         int limit = this.spawnCategoryLimit.getOrDefault(spawnCategory, -1);
         if (limit < 0) {
-            limit = this.server.getSpawnLimit(spawnCategory);
+            limit = this.server.getSpawnLimitUnsafe(spawnCategory);
+            // Paper end
         }
         return limit;
     }
diff --git a/src/test/java/io/papermc/paper/PaperCommandTest.java b/src/test/java/io/papermc/paper/PaperCommandTest.java
new file mode 100644
index 0000000000000000000000000000000000000000..4b5b368ef17bdb90f50e6ccc1f814cf93c7c0590
--- /dev/null
+++ b/src/test/java/io/papermc/paper/PaperCommandTest.java
@@ -0,0 +1,21 @@
+package io.papermc.paper;
+
+import com.destroystokyo.paper.PaperCommand;
+import java.util.HashSet;
+import java.util.Set;
+import net.minecraft.world.entity.MobCategory;
+import org.junit.Assert;
+import org.junit.Test;
+
+public class PaperCommandTest {
+    @Test
+    public void testMobCategoryColors() {
+        final Set<String> missing = new HashSet<>();
+        for (final MobCategory value : MobCategory.values()) {
+            if (!PaperCommand.MOB_CATEGORY_COLORS.containsKey(value)) {
+                missing.add(value.getName());
+            }
+        }
+        Assert.assertTrue("PaperCommand.MOB_CATEGORY_COLORS map missing TextColors for [" + String.join(", ", missing + "]"), missing.isEmpty());
+    }
+}