From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: Aikar <aikar@aikar.co>
Date: Wed, 4 Jul 2018 02:10:36 -0400
Subject: [PATCH] Store reference to current Chunk for Entity and Block
 Entities

This enables us a fast reference to the entities current chunk instead
of having to look it up by hashmap lookups.

diff --git a/src/main/java/net/minecraft/server/Chunk.java b/src/main/java/net/minecraft/server/Chunk.java
index 3b3e94cff2c6def318864bb5b7baefcf9de2cae7..6bd9457ad5c9fb1c0f12ff345b531be0673e866e 100644
--- a/src/main/java/net/minecraft/server/Chunk.java
+++ b/src/main/java/net/minecraft/server/Chunk.java
@@ -56,11 +56,36 @@ public class Chunk implements IChunkAccess {
         this(world, chunkcoordintpair, biomestorage, ChunkConverter.a, TickListEmpty.b(), TickListEmpty.b(), 0L, (ChunkSection[]) null, (Consumer) null);
     }
 
+    // Paper start
+    private class TileEntityHashMap extends java.util.HashMap<BlockPosition, TileEntity> {
+        @Override
+        public TileEntity put(BlockPosition key, TileEntity value) {
+            TileEntity replaced = super.put(key, value);
+            if (replaced != null) {
+                replaced.setCurrentChunk(null);
+            }
+            if (value != null) {
+                value.setCurrentChunk(Chunk.this);
+            }
+            return replaced;
+        }
+
+        @Override
+        public TileEntity remove(Object key) {
+            TileEntity removed = super.remove(key);
+            if (removed != null) {
+                removed.setCurrentChunk(null);
+            }
+            return removed;
+        }
+    }
+    // Paper end
+
     public Chunk(World world, ChunkCoordIntPair chunkcoordintpair, BiomeStorage biomestorage, ChunkConverter chunkconverter, TickList<Block> ticklist, TickList<FluidType> ticklist1, long i, @Nullable ChunkSection[] achunksection, @Nullable Consumer<Chunk> consumer) {
         this.sections = new ChunkSection[16];
         this.e = Maps.newHashMap();
         this.heightMap = Maps.newEnumMap(HeightMap.Type.class);
-        this.tileEntities = Maps.newHashMap();
+        this.tileEntities = new TileEntityHashMap(); // Paper
         this.l = Maps.newHashMap();
         this.m = Maps.newHashMap();
         this.n = new ShortList[16];
@@ -466,6 +491,7 @@ public class Chunk implements IChunkAccess {
         }
 
         entity.inChunk = true;
+        entity.setCurrentChunk(this); // Paper
         entity.chunkX = this.loc.x;
         entity.chunkY = k;
         entity.chunkZ = this.loc.z;
@@ -478,6 +504,7 @@ public class Chunk implements IChunkAccess {
         ((HeightMap) this.heightMap.get(heightmap_type)).a(along);
     }
 
+    public final void removeEntity(Entity entity) { this.b(entity); } // Paper - OBFHELPER
     public void b(Entity entity) {
         this.a(entity, entity.chunkY);
     }
@@ -491,7 +518,12 @@ public class Chunk implements IChunkAccess {
             i = this.entitySlices.length - 1;
         }
 
-        this.entitySlices[i].remove(entity);
+        // Paper start
+        if (entity.currentChunk != null && entity.currentChunk.get() == this) entity.setCurrentChunk(null);
+        if (!this.entitySlices[i].remove(entity)) {
+            return;
+        }
+        // Paper end
         this.entities.remove(entity); // Paper
     }
 
diff --git a/src/main/java/net/minecraft/server/Entity.java b/src/main/java/net/minecraft/server/Entity.java
index d6ee588d8752a9c3561f1cc61c9dccce16f07207..aaeb862fd93110bc19e0c2e212f67a19d5462dbb 100644
--- a/src/main/java/net/minecraft/server/Entity.java
+++ b/src/main/java/net/minecraft/server/Entity.java
@@ -175,7 +175,7 @@ public abstract class Entity implements INamableTileEntity, ICommandListener, Ke
     }
 
     public boolean isChunkLoaded() {
-        return world.isChunkLoaded((int) Math.floor(this.locX()) >> 4, (int) Math.floor(this.locZ()) >> 4);
+        return getCurrentChunk() != null;
     }
     // CraftBukkit end
 
@@ -1669,6 +1669,23 @@ public abstract class Entity implements INamableTileEntity, ICommandListener, Ke
     }
 
     // Paper start
+    java.lang.ref.WeakReference<Chunk> currentChunk = null;
+
+    public void setCurrentChunk(Chunk chunk) {
+        this.currentChunk = chunk != null ? new java.lang.ref.WeakReference<>(chunk) : null;
+    }
+    /**
+     * Returns the entities current registered chunk. If the entity is not added to a chunk yet, it will return null
+     */
+    public Chunk getCurrentChunk() {
+        final Chunk chunk = currentChunk != null ? currentChunk.get() : null;
+        if (chunk != null && chunk.loaded) {
+            return chunk;
+        }
+
+        return !inChunk ? null : ((WorldServer)world).getChunkProvider().getChunkAtIfLoadedMainThreadNoCache(chunkX, chunkZ);
+    }
+
     private MinecraftKey entityKey;
     private String entityKeyString;
 
diff --git a/src/main/java/net/minecraft/server/TileEntity.java b/src/main/java/net/minecraft/server/TileEntity.java
index cfbca7cfde3160e888e14c5bb3bc48267f1cb599..67fdd560f8136d83a13f75e265d5f5dd11871375 100644
--- a/src/main/java/net/minecraft/server/TileEntity.java
+++ b/src/main/java/net/minecraft/server/TileEntity.java
@@ -51,6 +51,15 @@ public abstract class TileEntity implements KeyedObject { // Paper
         getMinecraftKey(); // Try to load if it doesn't exists.
         return tileEntityKeyString;
     }
+
+    private java.lang.ref.WeakReference<Chunk> currentChunk = null;
+    public Chunk getCurrentChunk() {
+        final Chunk chunk = currentChunk != null ? currentChunk.get() : null;
+        return chunk != null && chunk.loaded ? chunk : null;
+    }
+    public void setCurrentChunk(Chunk chunk) {
+        this.currentChunk = chunk != null ? new java.lang.ref.WeakReference<>(chunk) : null;
+    }
     // Paper end
 
     @Nullable
diff --git a/src/main/java/org/bukkit/craftbukkit/entity/CraftEntity.java b/src/main/java/org/bukkit/craftbukkit/entity/CraftEntity.java
index fc23f2c4dc3ece5fd944de72b62b1a1ede3e7e45..66878c53bfb71392bf07d2c41a2edf1bc8f7af79 100644
--- a/src/main/java/org/bukkit/craftbukkit/entity/CraftEntity.java
+++ b/src/main/java/org/bukkit/craftbukkit/entity/CraftEntity.java
@@ -143,6 +143,7 @@ import net.minecraft.server.EntityZombieHusk;
 import net.minecraft.server.EntityZombieVillager;
 import net.minecraft.server.IChatBaseComponent;
 import net.minecraft.server.NBTTagCompound;
+import org.bukkit.Chunk; // Paper
 import org.bukkit.EntityEffect;
 import org.bukkit.Location;
 import org.bukkit.Server;
@@ -184,6 +185,12 @@ public abstract class CraftEntity implements org.bukkit.entity.Entity {
         this.entity = entity;
     }
 
+    @Override
+    public Chunk getChunk() {
+        net.minecraft.server.Chunk currentChunk = entity.getCurrentChunk();
+        return currentChunk != null ? currentChunk.bukkitChunk : getLocation().getChunk();
+    }
+
     public static CraftEntity getEntity(CraftServer server, Entity entity) {
         /*
          * Order is *EXTREMELY* important -- keep it right! =D