From 895a4fd4bd0ce4ddce500784a51635620f6d22d4 Mon Sep 17 00:00:00 2001
From: Felix Cravic <themode@outlook.fr>
Date: Mon, 1 Jun 2020 23:55:11 +0200
Subject: [PATCH] Fixed instance switch while unloading chunk

---
 .../minestom/server/instance/Instance.java    | 22 ++++++++++++-------
 .../server/instance/InstanceContainer.java    |  2 ++
 2 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/src/main/java/net/minestom/server/instance/Instance.java b/src/main/java/net/minestom/server/instance/Instance.java
index 310d926da..b9e778458 100644
--- a/src/main/java/net/minestom/server/instance/Instance.java
+++ b/src/main/java/net/minestom/server/instance/Instance.java
@@ -295,9 +295,13 @@ public abstract class Instance implements BlockModifier, EventHandler, DataConta
 
     /**
      * @param chunk the chunk to get the entities from
-     * @return an unmodifiable set containing all the entities in a chunk
+     * @return an unmodifiable set containing all the entities in a chunk,
+     * if {@code chunk} is null, return an empty {@link HashSet}
      */
     public Set<Entity> getChunkEntities(Chunk chunk) {
+        if (chunk == null)
+            return new HashSet<>();
+
         long index = ChunkUtils.getChunkIndex(chunk.getChunkX(), chunk.getChunkZ());
         return Collections.unmodifiableSet(getEntitiesInChunk(index));
     }
@@ -545,14 +549,16 @@ public abstract class Instance implements BlockModifier, EventHandler, DataConta
      * @param chunk  the chunk where the entity will be removed
      */
     public void removeEntityFromChunk(Entity entity, Chunk chunk) {
-        long chunkIndex = ChunkUtils.getChunkIndex(chunk.getChunkX(), chunk.getChunkZ());
         synchronized (chunkEntities) {
-            Set<Entity> entities = getEntitiesInChunk(chunkIndex);
-            entities.remove(entity);
-            if (entities.isEmpty()) {
-                this.chunkEntities.remove(chunkIndex);
-            } else {
-                this.chunkEntities.put(chunkIndex, entities);
+            if (chunk != null) {
+                long chunkIndex = ChunkUtils.getChunkIndex(chunk.getChunkX(), chunk.getChunkZ());
+                Set<Entity> entities = getEntitiesInChunk(chunkIndex);
+                entities.remove(entity);
+                if (entities.isEmpty()) {
+                    this.chunkEntities.remove(chunkIndex);
+                } else {
+                    this.chunkEntities.put(chunkIndex, entities);
+                }
             }
 
             if (entity instanceof Player) {
diff --git a/src/main/java/net/minestom/server/instance/InstanceContainer.java b/src/main/java/net/minestom/server/instance/InstanceContainer.java
index 2e3975609..ec6930a78 100644
--- a/src/main/java/net/minestom/server/instance/InstanceContainer.java
+++ b/src/main/java/net/minestom/server/instance/InstanceContainer.java
@@ -314,6 +314,8 @@ public class InstanceContainer extends Instance {
         }
 
         this.chunks.remove(index);
+        this.chunkEntities.remove(index);
+
         chunk.unload();
     }