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 2ef4b4c2ff81d0fa33d4630593266066d8e6a6f3..34bc24403a83ae578d2fc3956b4883894c618747 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;
@@ -10,6 +11,12 @@ import com.google.common.collect.Maps;
 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.TextComponent;
+import net.kyori.adventure.text.format.NamedTextColor;
+import net.kyori.adventure.text.format.TextColor;
+import net.minecraft.core.Registry;
 import net.minecraft.resources.ResourceLocation;
 import net.minecraft.server.MCUtil;
 import net.minecraft.server.MinecraftServer;
@@ -19,10 +26,12 @@ 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.network.protocol.game.ClientboundLightUpdatePacket;
 import net.minecraft.resources.ResourceLocation;
 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;
@@ -55,11 +64,12 @@ 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;
 
 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);
@@ -92,6 +102,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("*");
@@ -188,6 +202,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":
@@ -246,6 +266,183 @@ 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.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(TextComponent.ofChildren(
+                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.perPlayerMobSpawns) {
+            sender.sendMessage(Component.text("Use '/paper mobcaps' for worlds where per-player mob spawning is disabled.", NamedTextColor.RED));
+            return;
+        }
+
+        sender.sendMessage(TextComponent.ofChildren(Component.text("Mobcaps for player: "), Component.text(player.getName(), NamedTextColor.GREEN)));
+        sender.sendMessage(this.buildMobcapsComponent(
+            category -> level.chunkSource.chunkMap.getMobCountNear(serverPlayer, category),
+            category -> NaturalSpawner.limitForCategory(level, 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 = TextComponent.ofChildren(
+                    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 403f50ef908f65c62d4aaa7e5328faa0a7654480..b59f1554b4ad26d362c41c00b7a864a72efbd602 100644
--- a/src/main/java/net/minecraft/world/level/NaturalSpawner.java
+++ b/src/main/java/net/minecraft/world/level/NaturalSpawner.java
@@ -144,28 +144,15 @@ public final class NaturalSpawner {
             MobCategory enumcreaturetype = aenumcreaturetype[j];
             // CraftBukkit start - Use per-world spawn limits
             boolean spawnThisTick = true;
-            int limit = enumcreaturetype.getMaxInstancesPerChunk();
+            final int limit = limitForCategory(world, enumcreaturetype); // Paper
             switch (enumcreaturetype) {
-                case MONSTER:
-                    spawnThisTick = spawnMonsterThisTick;
-                    limit = world.getWorld().getMonsterSpawnLimit();
-                    break;
-                case CREATURE:
-                    spawnThisTick = spawnAnimalThisTick;
-                    limit = world.getWorld().getAnimalSpawnLimit();
-                    break;
-                case WATER_CREATURE:
-                    spawnThisTick = spawnWaterThisTick;
-                    limit = world.getWorld().getWaterAnimalSpawnLimit();
-                    break;
-                case AMBIENT:
-                    spawnThisTick = spawnAmbientThisTick;
-                    limit = world.getWorld().getAmbientSpawnLimit();
-                    break;
-                case WATER_AMBIENT:
-                    spawnThisTick = spawnWaterAmbientThisTick;
-                    limit = world.getWorld().getWaterAmbientSpawnLimit();
-                    break;
+                // Paper start - not mindiff so we get conflict on change
+                case MONSTER -> spawnThisTick = spawnMonsterThisTick;
+                case CREATURE -> spawnThisTick = spawnAnimalThisTick;
+                case WATER_CREATURE -> spawnThisTick = spawnWaterThisTick;
+                case AMBIENT -> spawnThisTick = spawnAmbientThisTick;
+                case WATER_AMBIENT -> spawnThisTick = spawnWaterAmbientThisTick;
+                // Paper end
             }
 
             if (!spawnThisTick || limit == 0) {
@@ -204,6 +191,23 @@ public final class NaturalSpawner {
         world.getProfiler().pop();
     }
 
+    // Paper start
+    public static int limitForCategory(final ServerLevel world, final MobCategory enumcreaturetype) {
+        return switch (enumcreaturetype) {
+            case MONSTER -> world.getWorld().getMonsterSpawnLimit();
+            case CREATURE -> world.getWorld().getAnimalSpawnLimit();
+            case WATER_CREATURE -> world.getWorld().getWaterAnimalSpawnLimit();
+            case AMBIENT -> world.getWorld().getAmbientSpawnLimit();
+            case WATER_AMBIENT -> world.getWorld().getWaterAmbientSpawnLimit();
+            default -> enumcreaturetype.getMaxInstancesPerChunk();
+        };
+    }
+
+    public static int globalLimitForCategory(final ServerLevel level, final MobCategory category, final int spawnableChunkCount) {
+        return limitForCategory(level, category) * spawnableChunkCount / NaturalSpawner.MAGIC_NUMBER;
+    }
+    // Paper end
+
     // Paper start - add parameters and int ret type
     public static void spawnCategoryForChunk(MobCategory group, ServerLevel world, LevelChunk chunk, NaturalSpawner.SpawnPredicate checker, NaturalSpawner.AfterSpawnCallback runner) {
         spawnCategoryForChunk(group, world, chunk, checker, runner);
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());
+    }
+}