From 4f16c3445f22654d43ca5b2006f45fac0e17203b Mon Sep 17 00:00:00 2001
From: Colin Godsey <crgodsey@gmail.com>
Date: Wed, 8 Aug 2018 10:10:06 -0600
Subject: [PATCH] Add entity count cache


diff --git a/src/main/java/net/minecraft/server/EntityInsentient.java b/src/main/java/net/minecraft/server/EntityInsentient.java
index 7b64ec27c..55606756e 100644
--- a/src/main/java/net/minecraft/server/EntityInsentient.java
+++ b/src/main/java/net/minecraft/server/EntityInsentient.java
@@ -42,6 +42,7 @@ public abstract class EntityInsentient extends EntityLiving {
     public float[] dropChanceArmor;
     // public boolean canPickUpLoot; // CraftBukkit - moved up to EntityLiving
     public boolean persistent;
+    public boolean countsAgainstSpawnLimit = true; // Paper
     private final Map<PathType, Float> bH;
     private MinecraftKey bI;
     private long bJ;
diff --git a/src/main/java/net/minecraft/server/EnumCreatureType.java b/src/main/java/net/minecraft/server/EnumCreatureType.java
index 79e52f7ba..288f59ed3 100644
--- a/src/main/java/net/minecraft/server/EnumCreatureType.java
+++ b/src/main/java/net/minecraft/server/EnumCreatureType.java
@@ -16,10 +16,17 @@ public enum EnumCreatureType {
         this.h = flag1;
     }
 
+    public Class<? extends IAnimal> innerClass() { return this.a(); } // Paper - OBFHELPER
     public Class<? extends IAnimal> a() {
         return this.e;
     }
 
+    // Paper start
+    public boolean matches(Entity entity) {
+        return innerClass().isAssignableFrom(entity.getClass());
+    }
+    // Paper end
+
     public int b() {
         return this.f;
     }
diff --git a/src/main/java/net/minecraft/server/SpawnerCreature.java b/src/main/java/net/minecraft/server/SpawnerCreature.java
index f525fd1b4..494759a1c 100644
--- a/src/main/java/net/minecraft/server/SpawnerCreature.java
+++ b/src/main/java/net/minecraft/server/SpawnerCreature.java
@@ -113,7 +113,7 @@ public final class SpawnerCreature {
                 // CraftBukkit end
 
                 if ((!enumcreaturetype.c() || flag1) && (enumcreaturetype.c() || flag) && (!enumcreaturetype.d() || flag2)) {
-                    k = worldserver.a(enumcreaturetype.a());
+                    k = worldserver.getCreatureCount(enumcreaturetype); // Paper - entity count cache
                     int l1 = limit * i / b; // CraftBukkit - use per-world limits
 
                     if (k <= l1) {
diff --git a/src/main/java/net/minecraft/server/World.java b/src/main/java/net/minecraft/server/World.java
index 004c3ec47..dc301c6f4 100644
--- a/src/main/java/net/minecraft/server/World.java
+++ b/src/main/java/net/minecraft/server/World.java
@@ -11,6 +11,7 @@ import java.util.Iterator;
 import java.util.List;
 import java.util.Random;
 import java.util.UUID;
+import java.util.EnumMap; // Paper
 
 import java.util.function.Function;
 import java.util.function.Predicate;
@@ -50,10 +51,41 @@ public abstract class World implements GeneratorAccess, IIBlockAccess, AutoClose
     // Spigot start - guard entity list from removals
     public final List<Entity> entityList = new java.util.ArrayList<Entity>()
     {
+        // Paper start - entity count cache
+        @Override
+        public boolean addAll(Collection<? extends Entity> c) {
+            for (Entity e : c) {
+                updateEntityCount(e, true);
+            }
+
+            return super.addAll(c);
+        }
+
+        @Override
+        public boolean removeAll(Collection<?> c) {
+            for (Object e : c) {
+                if (e instanceof Entity) {
+                    updateEntityCount((Entity)e, false);
+                }
+            }
+
+            return super.removeAll(c);
+        }
+
+        @Override
+        public boolean add(Entity e) {
+            updateEntityCount(e, true);
+
+            return super.add(e);
+        }
+
+        // Paper end
+
         @Override
         public Entity remove(int index)
         {
             guard();
+            updateEntityCount(get(index), false); // Paper
             return super.remove( index );
         }
 
@@ -61,6 +93,7 @@ public abstract class World implements GeneratorAccess, IIBlockAccess, AutoClose
         public boolean remove(Object o)
         {
             guard();
+            if (o instanceof Entity) updateEntityCount((Entity)o, false); // Paper
             return super.remove( o );
         }
 
@@ -2464,6 +2497,53 @@ public abstract class World implements GeneratorAccess, IIBlockAccess, AutoClose
         return i;
     }
 
+    // Paper start - entity count cache
+    private Map<EnumCreatureType, Integer> countCache = new EnumMap(EnumCreatureType.class);
+
+    public int getCreatureCount(EnumCreatureType type) {
+        Integer count = countCache.get(type);
+
+        return count == null ? 0 : count;
+    }
+
+    protected void updateEntityCount(EnumCreatureType type, boolean incr) {
+        Integer countObject = countCache.get(type);
+
+        int count = countObject == null ? 0 : countObject;
+
+        if (incr) count++;
+        else count--;
+
+        if (count < 0) {
+            e.warn("Paper - Entity count cache has gone negative");
+            count = 0;
+        }
+
+        countCache.put(type, count);
+    }
+
+    protected void updateEntityCount(Entity entity, boolean incr) {
+        if (entity == null || !(entity instanceof IAnimal)) return;
+
+        if (entity instanceof EntityInsentient) {
+            EntityInsentient entityinsentient = (EntityInsentient) entity;
+            if (incr && entityinsentient.isTypeNotPersistent() && entityinsentient.isPersistent()) {
+                entityinsentient.countsAgainstSpawnLimit = false;
+
+                return;
+            } else if (!incr && !entityinsentient.countsAgainstSpawnLimit) {
+                return;
+            }
+        }
+
+        for (EnumCreatureType type : EnumCreatureType.values()) {
+            if (type.matches(entity)) {
+                updateEntityCount(type, incr);
+            }
+        }
+    }
+    // Paper end
+
     public void addChunkEntities(Collection<Entity> collection) { a(collection); } // Paper - OBFHELPER
     public void a(Collection<Entity> collection) {
         org.spigotmc.AsyncCatcher.catchOp( "entity world add"); // Spigot
-- 
2.15.2 (Apple Git-101.1)