Paper/patches/server/0365-implement-optional-per-player-mob-spawns.patch

From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: kickash32 <kickash32@gmail.com>
Date: Mon, 19 Aug 2019 01:27:58 +0500
Subject: [PATCH] implement optional per player mob spawns


diff --git a/src/main/java/com/destroystokyo/paper/PaperWorldConfig.java b/src/main/java/com/destroystokyo/paper/PaperWorldConfig.java
index 68ae6664541d1dbc611520e7e8612b4f5162c696..6b0391743cd9e249c66796e7fe7a4da8c6b81b2e 100644
--- a/src/main/java/com/destroystokyo/paper/PaperWorldConfig.java
+++ b/src/main/java/com/destroystokyo/paper/PaperWorldConfig.java
@@ -613,4 +613,12 @@ public class PaperWorldConfig {
             }
         }
     }
+
+    public boolean perPlayerMobSpawns = false;
+    private void perPlayerMobSpawns() {
+        if (PaperConfig.version < 22) {
+            set("per-player-mob-spawns", Boolean.TRUE);
+        }
+        perPlayerMobSpawns = getBoolean("per-player-mob-spawns", true);
+    }
 }
diff --git a/src/main/java/com/destroystokyo/paper/util/PooledHashSets.java b/src/main/java/com/destroystokyo/paper/util/PooledHashSets.java
new file mode 100644
index 0000000000000000000000000000000000000000..11de56afaf059b00fa5bec293516bcdce7c4b2b9
--- /dev/null
+++ b/src/main/java/com/destroystokyo/paper/util/PooledHashSets.java
@@ -0,0 +1,241 @@
+package com.destroystokyo.paper.util;
+
+import it.unimi.dsi.fastutil.objects.Object2ObjectOpenHashMap;
+import it.unimi.dsi.fastutil.objects.ObjectLinkedOpenHashSet;
+import java.lang.ref.WeakReference;
+import java.util.Iterator;
+
+/** @author Spottedleaf */
+public class PooledHashSets<E> {
+
+    // we really want to avoid that equals() check as much as possible...
+    protected final Object2ObjectOpenHashMap<PooledObjectLinkedOpenHashSet<E>, PooledObjectLinkedOpenHashSet<E>> mapPool = new Object2ObjectOpenHashMap<>(64, 0.25f);
+
+    protected void decrementReferenceCount(final PooledObjectLinkedOpenHashSet<E> current) {
+        if (current.referenceCount == 0) {
+            throw new IllegalStateException("Cannot decrement reference count for " + current);
+        }
+        if (current.referenceCount == -1 || --current.referenceCount > 0) {
+            return;
+        }
+
+        this.mapPool.remove(current);
+        return;
+    }
+
+    public PooledObjectLinkedOpenHashSet<E> findMapWith(final PooledObjectLinkedOpenHashSet<E> current, final E object) {
+        final PooledObjectLinkedOpenHashSet<E> cached = current.getAddCache(object);
+
+        if (cached != null) {
+            if (cached.referenceCount != -1) {
+                ++cached.referenceCount;
+            }
+
+            decrementReferenceCount(current);
+
+            return cached;
+        }
+
+        if (!current.add(object)) {
+            return current;
+        }
+
+        // we use get/put since we use a different key on put
+        PooledObjectLinkedOpenHashSet<E> ret = this.mapPool.get(current);
+
+        if (ret == null) {
+            ret = new PooledObjectLinkedOpenHashSet<>(current);
+            current.remove(object);
+            this.mapPool.put(ret, ret);
+            ret.referenceCount = 1;
+        } else {
+            if (ret.referenceCount != -1) {
+                ++ret.referenceCount;
+            }
+            current.remove(object);
+        }
+
+        current.updateAddCache(object, ret);
+
+        decrementReferenceCount(current);
+        return ret;
+    }
+
+    // rets null if current.size() == 1
+    public PooledObjectLinkedOpenHashSet<E> findMapWithout(final PooledObjectLinkedOpenHashSet<E> current, final E object) {
+        if (current.set.size() == 1) {
+            decrementReferenceCount(current);
+            return null;
+        }
+
+        final PooledObjectLinkedOpenHashSet<E> cached = current.getRemoveCache(object);
+
+        if (cached != null) {
+            if (cached.referenceCount != -1) {
+                ++cached.referenceCount;
+            }
+
+            decrementReferenceCount(current);
+
+            return cached;
+        }
+
+        if (!current.remove(object)) {
+            return current;
+        }
+
+        // we use get/put since we use a different key on put
+        PooledObjectLinkedOpenHashSet<E> ret = this.mapPool.get(current);
+
+        if (ret == null) {
+            ret = new PooledObjectLinkedOpenHashSet<>(current);
+            current.add(object);
+            this.mapPool.put(ret, ret);
+            ret.referenceCount = 1;
+        } else {
+            if (ret.referenceCount != -1) {
+                ++ret.referenceCount;
+            }
+            current.add(object);
+        }
+
+        current.updateRemoveCache(object, ret);
+
+        decrementReferenceCount(current);
+        return ret;
+    }
+
+    public static final class PooledObjectLinkedOpenHashSet<E> implements Iterable<E> {
+
+        private static final WeakReference NULL_REFERENCE = new WeakReference(null);
+
+        final ObjectLinkedOpenHashSet<E> set;
+        int referenceCount; // -1 if special
+        int hash; // optimize hashcode
+
+        // add cache
+        WeakReference<E> lastAddObject = NULL_REFERENCE;
+        WeakReference<PooledObjectLinkedOpenHashSet<E>> lastAddMap = NULL_REFERENCE;
+
+        // remove cache
+        WeakReference<E> lastRemoveObject = NULL_REFERENCE;
+        WeakReference<PooledObjectLinkedOpenHashSet<E>> lastRemoveMap = NULL_REFERENCE;
+
+        public PooledObjectLinkedOpenHashSet() {
+            this.set = new ObjectLinkedOpenHashSet<>(2, 0.6f);
+        }
+
+        public PooledObjectLinkedOpenHashSet(final E single) {
+            this();
+            this.referenceCount = -1;
+            this.add(single);
+        }
+
+        public PooledObjectLinkedOpenHashSet(final PooledObjectLinkedOpenHashSet<E> other) {
+            this.set = other.set.clone();
+            this.hash = other.hash;
+        }
+
+        // from https://github.com/Spottedleaf/ConcurrentUtil/blob/master/src/main/java/ca/spottedleaf/concurrentutil/util/IntegerUtil.java
+        // generated by https://github.com/skeeto/hash-prospector
+        static int hash0(int x) {
+            x *= 0x36935555;
+            x ^= x >>> 16;
+            return x;
+        }
+
+        public PooledObjectLinkedOpenHashSet<E> getAddCache(final E element) {
+            final E currentAdd = this.lastAddObject.get();
+
+            if (currentAdd == null || !(currentAdd == element || currentAdd.equals(element))) {
+                return null;
+            }
+
+            final PooledObjectLinkedOpenHashSet<E> map = this.lastAddMap.get();
+            if (map == null || map.referenceCount == 0) {
+                // we need to ret null if ref count is zero as calling code will assume the map is in use
+                return null;
+            }
+
+            return map;
+        }
+
+        public PooledObjectLinkedOpenHashSet<E> getRemoveCache(final E element) {
+            final E currentRemove = this.lastRemoveObject.get();
+
+            if (currentRemove == null || !(currentRemove == element || currentRemove.equals(element))) {
+                return null;
+            }
+
+            final PooledObjectLinkedOpenHashSet<E> map = this.lastRemoveMap.get();
+            if (map == null || map.referenceCount == 0) {
+                // we need to ret null if ref count is zero as calling code will assume the map is in use
+                return null;
+            }
+
+            return map;
+        }
+
+        public void updateAddCache(final E element, final PooledObjectLinkedOpenHashSet<E> map) {
+            this.lastAddObject = new WeakReference<>(element);
+            this.lastAddMap = new WeakReference<>(map);
+        }
+
+        public void updateRemoveCache(final E element, final PooledObjectLinkedOpenHashSet<E> map) {
+            this.lastRemoveObject = new WeakReference<>(element);
+            this.lastRemoveMap = new WeakReference<>(map);
+        }
+
+        boolean add(final E element) {
+            boolean added =  this.set.add(element);
+
+            if (added) {
+                this.hash += hash0(element.hashCode());
+            }
+
+            return added;
+        }
+
+        boolean remove(Object element) {
+            boolean removed = this.set.remove(element);
+
+            if (removed) {
+                this.hash -= hash0(element.hashCode());
+            }
+
+            return removed;
+        }
+
+        @Override
+        public Iterator<E> iterator() {
+            return this.set.iterator();
+        }
+
+        @Override
+        public int hashCode() {
+            return this.hash;
+        }
+
+        @Override
+        public boolean equals(final Object other) {
+            if (!(other instanceof PooledObjectLinkedOpenHashSet)) {
+                return false;
+            }
+            if (this.referenceCount == 0) {
+                return other == this;
+            } else {
+                if (other == this) {
+                    // Unfortunately we are never equal to our own instance while in use!
+                    return false;
+                }
+                return this.hash == ((PooledObjectLinkedOpenHashSet)other).hash && this.set.equals(((PooledObjectLinkedOpenHashSet)other).set);
+            }
+        }
+
+        @Override
+        public String toString() {
+            return "PooledHashSet: size: " + this.set.size() + ", reference count: " + this.referenceCount + ", hash: " +
+                this.hashCode() + ", identity: " + System.identityHashCode(this) + " map: " + this.set.toString();
+        }
+    }
+} 
diff --git a/src/main/java/net/minecraft/server/level/ChunkMap.java b/src/main/java/net/minecraft/server/level/ChunkMap.java
index 6dae228c0d0460d74ed49aee5a38dfd7b145d891..c7114fc8045feab770dde30669d2711313bff189 100644
--- a/src/main/java/net/minecraft/server/level/ChunkMap.java
+++ b/src/main/java/net/minecraft/server/level/ChunkMap.java
@@ -155,6 +155,7 @@ public class ChunkMap extends ChunkStorage implements ChunkHolder.PlayerProvider
     private final Long2LongMap chunkSaveCooldowns;
     private final Queue<Runnable> unloadQueue;
     int viewDistance;
+    public final com.destroystokyo.paper.util.misc.PlayerAreaMap playerMobDistanceMap; // Paper
 
     // CraftBukkit start - recursion-safe executor for Chunk loadCallback() and unloadCallback()
     public final CallbackExecutor callbackExecutor = new CallbackExecutor();
@@ -184,16 +185,31 @@ public class ChunkMap extends ChunkStorage implements ChunkHolder.PlayerProvider
         int chunkX = MCUtil.getChunkCoordinate(player.getX());
         int chunkZ = MCUtil.getChunkCoordinate(player.getZ());
         // Note: players need to be explicitly added to distance maps before they can be updated
+        // Paper start - per player mob spawning
+        if (this.playerMobDistanceMap != null) {
+            this.playerMobDistanceMap.add(player, chunkX, chunkZ, this.distanceManager.getSimulationDistance());
+        }
+        // Paper end - per player mob spawning
     }
 
     void removePlayerFromDistanceMaps(ServerPlayer player) {
 
+        // Paper start - per player mob spawning
+        if (this.playerMobDistanceMap != null) {
+            this.playerMobDistanceMap.remove(player);
+        }
+        // Paper end - per player mob spawning
     }
 
     void updateMaps(ServerPlayer player) {
         int chunkX = MCUtil.getChunkCoordinate(player.getX());
         int chunkZ = MCUtil.getChunkCoordinate(player.getZ());
         // Note: players need to be explicitly added to distance maps before they can be updated
+        // Paper start - per player mob spawning
+        if (this.playerMobDistanceMap != null) {
+            this.playerMobDistanceMap.update(player, chunkX, chunkZ, this.distanceManager.getSimulationDistance());
+        }
+        // Paper end - per player mob spawning
     }
     // Paper end
     // Paper start
@@ -268,6 +284,7 @@ public class ChunkMap extends ChunkStorage implements ChunkHolder.PlayerProvider
         this.dataRegionManager = new io.papermc.paper.chunk.SingleThreadChunkRegionManager(this.level, 2, (1.0 / 3.0), 1, 6, "Data", DataRegionData::new, DataRegionSectionData::new);
         this.regionManagers.add(this.dataRegionManager);
         // Paper end
+        this.playerMobDistanceMap = this.level.paperConfig.perPlayerMobSpawns ? new com.destroystokyo.paper.util.misc.PlayerAreaMap(this.pooledLinkedPlayerHashSets) : null; // Paper
     }
 
     protected ChunkGenerator generator() {
@@ -285,6 +302,31 @@ public class ChunkMap extends ChunkStorage implements ChunkHolder.PlayerProvider
         });
     }
 
+    // Paper start
+    public void updatePlayerMobTypeMap(Entity entity) {
+        if (!this.level.paperConfig.perPlayerMobSpawns) {
+            return;
+        }
+        int index = entity.getType().getCategory().ordinal();
+
+        final com.destroystokyo.paper.util.misc.PooledLinkedHashSets.PooledObjectLinkedOpenHashSet<ServerPlayer> inRange = this.playerMobDistanceMap.getObjectsInRange(entity.chunkPosition());
+        if (inRange == null) {
+            return;
+        }
+        final Object[] backingSet = inRange.getBackingSet();
+        for (int i = 0; i < backingSet.length; i++) {
+            if (!(backingSet[i] instanceof final ServerPlayer player)) {
+                continue;
+            }
+            ++player.mobCounts[index];
+        }
+    }
+
+    public int getMobCountNear(ServerPlayer entityPlayer, net.minecraft.world.entity.MobCategory mobCategory) {
+        return entityPlayer.mobCounts[mobCategory.ordinal()];
+    }
+    // Paper end
+
     private static double euclideanDistanceSquared(ChunkPos pos, Entity entity) {
         double d0 = (double) SectionPos.sectionToBlockCoord(pos.x, 8);
         double d1 = (double) SectionPos.sectionToBlockCoord(pos.z, 8);
diff --git a/src/main/java/net/minecraft/server/level/DistanceManager.java b/src/main/java/net/minecraft/server/level/DistanceManager.java
index 1e066a35b53b1f71a0e6376a22d51fc4c0a412dc..6228f2f67541da62b0ae093de987662db9643740 100644
--- a/src/main/java/net/minecraft/server/level/DistanceManager.java
+++ b/src/main/java/net/minecraft/server/level/DistanceManager.java
@@ -326,6 +326,12 @@ public abstract class DistanceManager {
 
     }
 
+    // Paper start
+    public int getSimulationDistance() {
+        return this.simulationDistance;
+    }
+    // Paper end
+
     public int getNaturalSpawnChunkCount() {
         this.naturalSpawnChunkCounter.runAllUpdates();
         return this.naturalSpawnChunkCounter.chunks.size();
diff --git a/src/main/java/net/minecraft/server/level/ServerChunkCache.java b/src/main/java/net/minecraft/server/level/ServerChunkCache.java
index b6737c0b74e821c948919ca4184dfe0281a19894..cbb8adea5e867a92e1dc4c94a0e7e116b7ac2274 100644
--- a/src/main/java/net/minecraft/server/level/ServerChunkCache.java
+++ b/src/main/java/net/minecraft/server/level/ServerChunkCache.java
@@ -915,7 +915,18 @@ public class ServerChunkCache extends ChunkSource {
             gameprofilerfiller.push("naturalSpawnCount");
             this.level.timings.countNaturalMobs.startTiming(); // Paper - timings
             int l = this.distanceManager.getNaturalSpawnChunkCount();
-            NaturalSpawner.SpawnState spawnercreature_d = NaturalSpawner.createState(l, this.level.getAllEntities(), this::getFullChunk, new LocalMobCapCalculator(this.chunkMap));
+            // Paper start - per player mob spawning
+            NaturalSpawner.SpawnState spawnercreature_d; // moved down
+            if ((this.spawnFriendlies || this.spawnEnemies) && this.chunkMap.playerMobDistanceMap != null) { // don't count mobs when animals and monsters are disabled
+                // re-set mob counts
+                for (ServerPlayer player : this.level.players) {
+                    Arrays.fill(player.mobCounts, 0);
+                }
+                spawnercreature_d = NaturalSpawner.createState(l, this.level.getAllEntities(), this::getFullChunk, null, true);
+            } else {
+                spawnercreature_d = NaturalSpawner.createState(l, this.level.getAllEntities(), this::getFullChunk, this.chunkMap.playerMobDistanceMap == null ? new LocalMobCapCalculator(this.chunkMap) : null, false);
+            }
+            // Paper end
             this.level.timings.countNaturalMobs.stopTiming(); // Paper - timings
 
             this.lastSpawnState = spawnercreature_d;
diff --git a/src/main/java/net/minecraft/server/level/ServerPlayer.java b/src/main/java/net/minecraft/server/level/ServerPlayer.java
index 1d0dc77947d1b7fc89d41e5edb660aea86282562..49f272d37fb59fd4c13466217389e626b6ca1e2d 100644
--- a/src/main/java/net/minecraft/server/level/ServerPlayer.java
+++ b/src/main/java/net/minecraft/server/level/ServerPlayer.java
@@ -225,6 +225,11 @@ public class ServerPlayer extends Player {
     public boolean queueHealthUpdatePacket = false;
     public net.minecraft.network.protocol.game.ClientboundSetHealthPacket queuedHealthUpdatePacket;
     // Paper end
+    // Paper start - mob spawning rework
+    public static final int MOBCATEGORY_TOTAL_ENUMS = net.minecraft.world.entity.MobCategory.values().length;
+    public final int[] mobCounts = new int[MOBCATEGORY_TOTAL_ENUMS]; // Paper
+    public final com.destroystokyo.paper.util.PooledHashSets.PooledObjectLinkedOpenHashSet<ServerPlayer> cachedSingleMobDistanceMap;
+    // Paper end
 
     // CraftBukkit start
     public String displayName;
@@ -314,6 +319,7 @@ public class ServerPlayer extends Player {
         this.adventure$displayName = net.kyori.adventure.text.Component.text(this.getScoreboardName()); // Paper
         this.bukkitPickUpLoot = true;
         this.maxHealthCache = this.getMaxHealth();
+        this.cachedSingleMobDistanceMap = new com.destroystokyo.paper.util.PooledHashSets.PooledObjectLinkedOpenHashSet<>(this); // Paper
     }
 
     // Yes, this doesn't match Vanilla, but it's the best we can do for now.
diff --git a/src/main/java/net/minecraft/world/level/NaturalSpawner.java b/src/main/java/net/minecraft/world/level/NaturalSpawner.java
index 474dc6302d2b0f8c629b2e1f6de9d4ba6faf398b..1622450b53e0f0f48c3ca107e4d705b4ad29dadf 100644
--- a/src/main/java/net/minecraft/world/level/NaturalSpawner.java
+++ b/src/main/java/net/minecraft/world/level/NaturalSpawner.java
@@ -67,7 +67,13 @@ public final class NaturalSpawner {
 
     private NaturalSpawner() {}
 
+    // Paper start - add countMobs parameter
     public static NaturalSpawner.SpawnState createState(int spawningChunkCount, Iterable<Entity> entities, NaturalSpawner.ChunkGetter chunkSource, LocalMobCapCalculator localmobcapcalculator) {
+        return createState(spawningChunkCount, entities, chunkSource, localmobcapcalculator, false);
+    }
+
+    public static NaturalSpawner.SpawnState createState(int spawningChunkCount, Iterable<Entity> entities, NaturalSpawner.ChunkGetter chunkSource, LocalMobCapCalculator localmobcapcalculator, boolean countMobs) {
+        // Paper end
         PotentialCalculator spawnercreatureprobabilities = new PotentialCalculator();
         Object2IntOpenHashMap<MobCategory> object2intopenhashmap = new Object2IntOpenHashMap();
         Iterator iterator = entities.iterator();
@@ -103,11 +109,16 @@ public final class NaturalSpawner {
                         spawnercreatureprobabilities.addCharge(entity.blockPosition(), biomesettingsmobs_b.getCharge());
                     }
 
-                    if (entity instanceof Mob) {
+                    if (localmobcapcalculator != null && entity instanceof Mob) { // Paper
                         localmobcapcalculator.addMob(chunk.getPos(), enumcreaturetype);
                     }
 
                     object2intopenhashmap.addTo(enumcreaturetype, 1);
+                    // Paper start
+                    if (countMobs) {
+                        chunk.level.getChunkSource().chunkMap.updatePlayerMobTypeMap(entity);
+                    }
+                    // Paper end
                 });
             }
         }
@@ -142,13 +153,37 @@ public final class NaturalSpawner {
                 continue;
             }
 
-            if ((spawnAnimals || !enumcreaturetype.isFriendly()) && (spawnMonsters || enumcreaturetype.isFriendly()) && (rareSpawn || !enumcreaturetype.isPersistent()) && info.canSpawnForCategory(enumcreaturetype, chunk.getPos(), limit)) {
+            // Paper start - only allow spawns upto the limit per chunk and update count afterwards
+            int currEntityCount = info.mobCategoryCounts.getInt(enumcreaturetype);
+            int k1 = limit * info.getSpawnableChunkCount() / NaturalSpawner.MAGIC_NUMBER;
+            int difference = k1 - currEntityCount;
+
+            if (world.paperConfig.perPlayerMobSpawns) {
+                int minDiff = Integer.MAX_VALUE;
+                final com.destroystokyo.paper.util.misc.PooledLinkedHashSets.PooledObjectLinkedOpenHashSet<net.minecraft.server.level.ServerPlayer> inRange = world.getChunkSource().chunkMap.playerMobDistanceMap.getObjectsInRange(chunk.getPos());
+                if (inRange != null) {
+                    final Object[] backingSet = inRange.getBackingSet();
+                    for (int k = 0; k < backingSet.length; k++) {
+                        if (!(backingSet[k] instanceof final net.minecraft.server.level.ServerPlayer player)) {
+                            continue;
+                        }
+                        minDiff = Math.min(limit - world.getChunkSource().chunkMap.getMobCountNear(player, enumcreaturetype), minDiff);
+                    }
+                }
+                difference = (minDiff == Integer.MAX_VALUE) ? 0 : minDiff;
+            }
+            if ((spawnAnimals || !enumcreaturetype.isFriendly()) && (spawnMonsters || enumcreaturetype.isFriendly()) && (rareSpawn || !enumcreaturetype.isPersistent()) && difference > 0) {
+                // Paper end
                 // CraftBukkit end
                 Objects.requireNonNull(info);
                 NaturalSpawner.SpawnPredicate spawnercreature_c = info::canSpawn;
 
                 Objects.requireNonNull(info);
-                NaturalSpawner.spawnCategoryForChunk(enumcreaturetype, world, chunk, spawnercreature_c, info::afterSpawn);
+                // Paper start
+                int spawnCount = NaturalSpawner.spawnCategoryForChunk(enumcreaturetype, world, chunk, spawnercreature_c, info::afterSpawn,
+                    difference, world.paperConfig.perPlayerMobSpawns ? world.getChunkSource().chunkMap::updatePlayerMobTypeMap : null);
+                info.mobCategoryCounts.mergeInt(enumcreaturetype, spawnCount, Integer::sum);
+                // Paper end
             }
         }
 
@@ -156,12 +191,18 @@ public final class NaturalSpawner {
         world.getProfiler().pop();
     }
 
+    // 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, Integer.MAX_VALUE, null);
+    }
+    public static int spawnCategoryForChunk(MobCategory group, ServerLevel world, LevelChunk chunk, NaturalSpawner.SpawnPredicate checker, NaturalSpawner.AfterSpawnCallback runner, int maxSpawns, Consumer<Entity> trackEntity) {
+        // Paper end - add parameters and int ret type
         BlockPos blockposition = NaturalSpawner.getRandomPosWithin(world, chunk);
 
         if (blockposition.getY() >= world.getMinBuildHeight() + 1) {
-            NaturalSpawner.spawnCategoryForPosition(group, world, chunk, blockposition, checker, runner);
+            return NaturalSpawner.spawnCategoryForPosition(group, world, chunk, blockposition, checker, runner, maxSpawns, trackEntity); // Paper
         }
+        return 0; // Paper
     }
 
     @VisibleForDebug
@@ -172,15 +213,21 @@ public final class NaturalSpawner {
         });
     }
 
+    // Paper start - add maxSpawns parameter and return spawned mobs
     public static void spawnCategoryForPosition(MobCategory group, ServerLevel world, ChunkAccess chunk, BlockPos pos, NaturalSpawner.SpawnPredicate checker, NaturalSpawner.AfterSpawnCallback runner) {
+        spawnCategoryForPosition(group, world,chunk, pos, checker, runner, Integer.MAX_VALUE, null);
+    }
+    public static int spawnCategoryForPosition(MobCategory group, ServerLevel world, ChunkAccess chunk, BlockPos pos, NaturalSpawner.SpawnPredicate checker, NaturalSpawner.AfterSpawnCallback runner, int maxSpawns, Consumer<Entity> trackEntity) {
+    // Paper end - add maxSpawns parameter and return spawned mobs
         StructureFeatureManager structuremanager = world.structureFeatureManager();
         ChunkGenerator chunkgenerator = world.getChunkSource().getGenerator();
         int i = pos.getY();
         BlockState iblockdata = world.getBlockStateIfLoadedAndInBounds(pos); // Paper - don't load chunks for mob spawn
+        int j = 0; // Paper - moved up
 
         if (iblockdata != null && !iblockdata.isRedstoneConductor(chunk, pos)) { // Paper - don't load chunks for mob spawn
             BlockPos.MutableBlockPos blockposition_mutableblockposition = new BlockPos.MutableBlockPos();
-            int j = 0;
+            //int j = 0; // Paper - moved up
             int k = 0;
 
             while (k < 3) {
@@ -222,14 +269,14 @@ public final class NaturalSpawner {
                                     // Paper start
                                     Boolean doSpawning = isValidSpawnPostitionForType(world, group, structuremanager, chunkgenerator, biomesettingsmobs_c, blockposition_mutableblockposition, d2);
                                     if (doSpawning == null) {
-                                        return;
+                                        return j; // Paper
                                     }
                                     if (doSpawning && checker.test(biomesettingsmobs_c.type, blockposition_mutableblockposition, chunk)) {
                                         // Paper end
                                         Mob entityinsentient = NaturalSpawner.getMobForSpawn(world, biomesettingsmobs_c.type);
 
                                         if (entityinsentient == null) {
-                                            return;
+                                            return j; // Paper
                                         }
 
                                         entityinsentient.moveTo(d0, (double) i, d1, world.random.nextFloat() * 360.0F, 0.0F);
@@ -241,10 +288,15 @@ public final class NaturalSpawner {
                                                 ++j;
                                                 ++k1;
                                                 runner.run(entityinsentient, chunk);
+                                                // Paper start
+                                                if (trackEntity != null) {
+                                                    trackEntity.accept(entityinsentient);
+                                                }
+                                                // Paper end
                                             }
                                             // CraftBukkit end
-                                            if (j >= entityinsentient.getMaxSpawnClusterSize()) {
-                                                return;
+                                            if (j >= entityinsentient.getMaxSpawnClusterSize() || j >= maxSpawns) { // Paper
+                                                return j; // Paper
                                             }
 
                                             if (entityinsentient.isMaxGroupSizeReached(k1)) {
@@ -266,6 +318,7 @@ public final class NaturalSpawner {
             }
 
         }
+        return j; // Paper
     }
 
     private static boolean isRightDistanceToPlayerAndSpawnPoint(ServerLevel world, ChunkAccess chunk, BlockPos.MutableBlockPos pos, double squaredDistance) {
@@ -551,7 +604,7 @@ public final class NaturalSpawner {
             MobCategory enumcreaturetype = entitytypes.getCategory();
 
             this.mobCategoryCounts.addTo(enumcreaturetype, 1);
-            this.localMobCapCalculator.addMob(new ChunkPos(blockposition), enumcreaturetype);
+            if (this.localMobCapCalculator != null) this.localMobCapCalculator.addMob(new ChunkPos(blockposition), enumcreaturetype); // Paper
         }
 
         public int getSpawnableChunkCount() {
@@ -567,6 +620,7 @@ public final class NaturalSpawner {
             int i = limit * this.spawnableChunkCount / NaturalSpawner.MAGIC_NUMBER;
             // CraftBukkit end
 
+            if (this.localMobCapCalculator == null) return this.mobCategoryCounts.getInt(enumcreaturetype) < i; // Paper
             return this.mobCategoryCounts.getInt(enumcreaturetype) >= i ? false : this.localMobCapCalculator.canSpawn(enumcreaturetype, chunkcoordintpair);
         }
     }