Improve entity chunk switch

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)}
+     * Changes an entity chunk.
-     * and {@link #UNSAFE_addEntityToChunk(Entity, Chunk)} simultaneously.
      *
      * @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);
+        Check.notNull(newChunk, "The chunk {0} is not loaded, you can make it automatic by using Instance#enableAutoChunkLoad(true)", newChunk);
-        UNSAFE_addEntityToChunk(entity, 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);
+        }
     }
 
-    /**
+    private void UNSAFE_addEntityToChunk(@NotNull Entity entity, @NotNull Chunk chunk) {
-     * Adds the specified {@link Entity} to the instance entities cache.
+        final long chunkIndex = ChunkUtils.getChunkIndex(chunk);
-     * <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());
         synchronized (entitiesLock) {
-            Set<Entity> entities = getEntitiesInChunk(chunkIndex);
+            getEntitiesInChunk(chunkIndex).add(entity);
-            entities.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
         }
     }
 
-    /**
+    private void UNSAFE_removeEntityFromChunk(@NotNull Entity entity, @NotNull Chunk chunk) {
-     * Removes the specified {@link Entity} to the instance entities cache.
+        final long chunkIndex = ChunkUtils.getChunkIndex(chunk);
-     * <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) {
         synchronized (entitiesLock) {
-            final long chunkIndex = ChunkUtils.getChunkIndex(chunk.getChunkX(), chunk.getChunkZ());
+            getEntitiesInChunk(chunkIndex).remove(entity);
-            Set<Entity> entities = getEntitiesInChunk(chunkIndex);
-            entities.remove(entity);
-
             this.entities.remove(entity);
             if (entity instanceof Player) {
                 this.players.remove(entity);

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);
     }
 

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;