From 94a96d7df1efb12127b350a0321611476b205ef1 Mon Sep 17 00:00:00 2001
From: TheMode <themode@outlook.fr>
Date: Wed, 21 Jul 2021 08:25:44 +0200
Subject: [PATCH] Improve entity chunk switch

---
 .../minestom/server/instance/Instance.java    | 49 ++++++-------------
 .../server/instance/InstanceContainer.java    |  2 +-
 .../server/utils/chunk/ChunkUtils.java        |  4 ++
 3 files changed, 20 insertions(+), 35 deletions(-)

diff --git a/src/main/java/net/minestom/server/instance/Instance.java b/src/main/java/net/minestom/server/instance/Instance.java
index b90c9441c..7287dfae6 100644
--- a/src/main/java/net/minestom/server/instance/Instance.java
+++ b/src/main/java/net/minestom/server/instance/Instance.java
@@ -629,8 +629,7 @@ public abstract class Instance implements BlockGetter, BlockSetter, Tickable, Ta
     }
 
     /**
-     * Synchronized method to execute {@link #UNSAFE_removeEntityFromChunk(Entity, Chunk)}
-     * and {@link #UNSAFE_addEntityToChunk(Entity, Chunk)} simultaneously.
+     * Changes an entity chunk.
      *
      * @param entity    the entity to change its chunk
      * @param lastChunk the last entity chunk
@@ -638,27 +637,20 @@ public abstract class Instance implements BlockGetter, BlockSetter, Tickable, Ta
      */
     @ApiStatus.Internal
     public synchronized void UNSAFE_switchEntityChunk(@NotNull Entity entity, @NotNull Chunk lastChunk, @NotNull Chunk newChunk) {
-        UNSAFE_removeEntityFromChunk(entity, lastChunk);
-        UNSAFE_addEntityToChunk(entity, newChunk);
+        Check.notNull(newChunk, "The chunk {0} is not loaded, you can make it automatic by using Instance#enableAutoChunkLoad(true)", newChunk);
+        Check.argCondition(!newChunk.isLoaded(), "Chunk {0} has been unloaded previously", newChunk);
+        final long oldIndex = ChunkUtils.getChunkIndex(lastChunk);
+        final long newIndex = ChunkUtils.getChunkIndex(newChunk);
+        synchronized (entitiesLock) {
+            getEntitiesInChunk(oldIndex).remove(entity);
+            getEntitiesInChunk(newIndex).add(entity);
+        }
     }
 
-    /**
-     * Adds the specified {@link Entity} to the instance entities cache.
-     * <p>
-     * Warning: this is done automatically when the entity move out of his chunk.
-     *
-     * @param entity the entity to add
-     * @param chunk  the chunk where the entity will be added
-     */
-    @ApiStatus.Internal
-    public void UNSAFE_addEntityToChunk(@NotNull Entity entity, @NotNull Chunk chunk) {
-        Check.notNull(chunk, "The chunk {0} is not loaded, you can make it automatic by using Instance#enableAutoChunkLoad(true)", chunk);
-        Check.argCondition(!chunk.isLoaded(), "Chunk {0} has been unloaded previously", chunk);
-        final long chunkIndex = ChunkUtils.getChunkIndex(chunk.getChunkX(), chunk.getChunkZ());
+    private void UNSAFE_addEntityToChunk(@NotNull Entity entity, @NotNull Chunk chunk) {
+        final long chunkIndex = ChunkUtils.getChunkIndex(chunk);
         synchronized (entitiesLock) {
-            Set<Entity> entities = getEntitiesInChunk(chunkIndex);
-            entities.add(entity);
-
+            getEntitiesInChunk(chunkIndex).add(entity);
             this.entities.add(entity);
             if (entity instanceof Player) {
                 this.players.add((Player) entity);
@@ -670,21 +662,10 @@ public abstract class Instance implements BlockGetter, BlockSetter, Tickable, Ta
         }
     }
 
-    /**
-     * Removes the specified {@link Entity} to the instance entities cache.
-     * <p>
-     * Warning: this is done automatically when the entity move out of his chunk.
-     *
-     * @param entity the entity to remove
-     * @param chunk  the chunk where the entity will be removed
-     */
-    @ApiStatus.Internal
-    public void UNSAFE_removeEntityFromChunk(@NotNull Entity entity, @NotNull Chunk chunk) {
+    private void UNSAFE_removeEntityFromChunk(@NotNull Entity entity, @NotNull Chunk chunk) {
+        final long chunkIndex = ChunkUtils.getChunkIndex(chunk);
         synchronized (entitiesLock) {
-            final long chunkIndex = ChunkUtils.getChunkIndex(chunk.getChunkX(), chunk.getChunkZ());
-            Set<Entity> entities = getEntitiesInChunk(chunkIndex);
-            entities.remove(entity);
-
+            getEntitiesInChunk(chunkIndex).remove(entity);
             this.entities.remove(entity);
             if (entity instanceof Player) {
                 this.players.remove(entity);
diff --git a/src/main/java/net/minestom/server/instance/InstanceContainer.java b/src/main/java/net/minestom/server/instance/InstanceContainer.java
index 54fd7f16f..6dcb2155e 100644
--- a/src/main/java/net/minestom/server/instance/InstanceContainer.java
+++ b/src/main/java/net/minestom/server/instance/InstanceContainer.java
@@ -461,7 +461,7 @@ public class InstanceContainer extends Instance {
      * @param chunk the chunk to cache
      */
     public void cacheChunk(@NotNull Chunk chunk) {
-        final long index = ChunkUtils.getChunkIndex(chunk.getChunkX(), chunk.getChunkZ());
+        final long index = ChunkUtils.getChunkIndex(chunk);
         this.chunks.put(index, chunk);
     }
 
diff --git a/src/main/java/net/minestom/server/utils/chunk/ChunkUtils.java b/src/main/java/net/minestom/server/utils/chunk/ChunkUtils.java
index 82d566e70..c30429798 100644
--- a/src/main/java/net/minestom/server/utils/chunk/ChunkUtils.java
+++ b/src/main/java/net/minestom/server/utils/chunk/ChunkUtils.java
@@ -116,6 +116,10 @@ public final class ChunkUtils {
         return (((long) chunkX) << 32) | (chunkZ & 0xffffffffL);
     }
 
+    public static long getChunkIndex(@NotNull Chunk chunk) {
+        return getChunkIndex(chunk.getChunkX(), chunk.getChunkZ());
+    }
+
     public static long getChunkIndexWithSection(int chunkX, int chunkZ, int section) {
         long l = 0L;
         l |= ((long) chunkX & 4194303L) << 42;