World: Track entities separately from chunks

this allows entities to exist outside of generated chunks, with one caveat: they won't be saved in such cases.
Obviously, for player entities, this doesn't matter.

fixes #3947

src/player/Player.php

@@ -630,14 +630,11 @@ class Player extends Human implements CommandSender, ChunkListener, IPlayer{
 		$world = $world ?? $this->getWorld();
 		$index = World::chunkHash($x, $z);
 		if(isset($this->usedChunks[$index])){
-			$chunk = $world->getChunk($x, $z);
+			foreach($world->getChunkEntities($x, $z) as $entity){
-			if($chunk !== null){ //this might be a chunk that hasn't been generated yet
-				foreach($chunk->getEntities() as $entity){
 				if($entity !== $this){
 					$entity->despawnFrom($this);
 				}
 			}
-			}
 			$this->getNetworkSession()->stopUsingChunk($x, $z);
 			unset($this->usedChunks[$index]);
 		}
@@ -656,7 +653,7 @@ class Player extends Human implements CommandSender, ChunkListener, IPlayer{
 	}
 
 	protected function spawnEntitiesOnChunk(int $chunkX, int $chunkZ) : void{
-		foreach($this->getWorld()->getChunk($chunkX, $chunkZ)->getEntities() as $entity){
+		foreach($this->getWorld()->getChunkEntities($chunkX, $chunkZ) as $entity){
 			if($entity !== $this and !$entity->isFlaggedForDespawn()){
 				$entity->spawnTo($this);
 			}

src/world/World.php

@@ -93,6 +93,7 @@ use pocketmine\world\sound\Sound;
 use pocketmine\world\utils\SubChunkExplorer;
 use function abs;
 use function array_fill_keys;
+use function array_filter;
 use function array_key_exists;
 use function array_map;
 use function array_merge;
@@ -158,6 +159,12 @@ class World implements ChunkManager{
 	 */
 	private $entityLastKnownPositions = [];
 
+	/**
+	 * @var Entity[][]
+	 * @phpstan-var array<int, array<int, Entity>>
+	 */
+	private array $entitiesByChunk = [];
+
 	/** @var Entity[] */
 	public $updateEntities = [];
 	/** @var Block[][] */
@@ -516,6 +523,15 @@ class World implements ChunkManager{
 			self::getXZ($chunkHash, $chunkX, $chunkZ);
 			$this->unloadChunk($chunkX, $chunkZ, false);
 		}
+		foreach($this->entitiesByChunk as $chunkHash => $entities){
+			self::getXZ($chunkHash, $chunkX, $chunkZ);
+			if(count($entities) !== 0){
+				$this->logger->warning(count($entities) . " entities found in ungenerated chunk $chunkX $chunkZ, they won't be saved!");
+			}
+			foreach($entities as $entity){
+				$entity->close();
+			}
+		}
 
 		$this->save();
 
@@ -1070,7 +1086,7 @@ class World implements ChunkManager{
 		if($chunk === null){
 			throw new \InvalidArgumentException("Chunk is not loaded");
 		}
-		foreach($chunk->getEntities() as $entity){
+		foreach($this->getChunkEntities($chunkX, $chunkZ) as $entity){
 			$entity->onRandomUpdate();
 		}
 
@@ -1129,7 +1145,7 @@ class World implements ChunkManager{
 				self::getXZ($chunkHash, $chunkX, $chunkZ);
 				$this->provider->saveChunk($chunkX, $chunkZ, new ChunkData(
 					$chunk,
-					array_map(fn(Entity $e) => $e->saveNBT(), $chunk->getSavableEntities()),
+					array_map(fn(Entity $e) => $e->saveNBT(), array_filter($this->getChunkEntities($chunkX, $chunkZ), fn(Entity $e) => $e->canSaveWithChunk())),
 					array_map(fn(Tile $t) => $t->saveNBT(), $chunk->getTiles()),
 				));
 				$chunk->clearTerrainDirtyFlags();
@@ -1921,7 +1937,7 @@ class World implements ChunkManager{
 				if(!$this->isChunkLoaded($x, $z)){
 					continue;
 				}
-				foreach($this->getChunk($x, $z)->getEntities() as $ent){
+				foreach($this->getChunkEntities($x, $z) as $ent){
 					if($ent !== $entity and $ent->boundingBox->intersectsWith($bb)){
 						$nearby[] = $ent;
 					}
@@ -1964,7 +1980,7 @@ class World implements ChunkManager{
 				if(!$this->isChunkLoaded($x, $z)){
 					continue;
 				}
-				foreach($this->getChunk($x, $z)->getEntities() as $entity){
+				foreach($this->getChunkEntities($x, $z) as $entity){
 					if(!($entity instanceof $entityType) or $entity->isFlaggedForDespawn() or (!$includeDead and !$entity->isAlive())){
 						continue;
 					}
@@ -2043,6 +2059,13 @@ class World implements ChunkManager{
 		return $this->chunks[World::chunkHash($chunkX, $chunkZ)] ?? null;
 	}
 
+	/**
+	 * @return Entity[]
+	 */
+	public function getChunkEntities(int $chunkX, int $chunkZ) : array{
+		return $this->entitiesByChunk[World::chunkHash($chunkX, $chunkZ)] ?? [];
+	}
+
 	/**
 	 * Returns the chunk containing the given Vector3 position.
 	 */
@@ -2154,11 +2177,8 @@ class World implements ChunkManager{
 		$oldChunk = $this->loadChunk($chunkX, $chunkZ);
 		if($oldChunk !== null and $oldChunk !== $chunk){
 			if($deleteEntitiesAndTiles){
-				foreach($oldChunk->getEntities() as $entity){
+				foreach($this->getChunkEntities($chunkX, $chunkZ) as $entity){
-					if($entity instanceof Player){
+					if(!($entity instanceof Player)){
-						$chunk->addEntity($entity);
-						$oldChunk->removeEntity($entity);
-					}else{
 						$entity->close();
 					}
 				}
@@ -2166,11 +2186,6 @@ class World implements ChunkManager{
 					$tile->close();
 				}
 			}else{
-				foreach($oldChunk->getEntities() as $entity){
-					$chunk->addEntity($entity);
-					$oldChunk->removeEntity($entity);
-				}
-
 				foreach($oldChunk->getTiles() as $tile){
 					$chunk->addTile($tile);
 					$oldChunk->removeTile($tile);
@@ -2272,7 +2287,7 @@ class World implements ChunkManager{
 		if($chunk === null){
 			throw new \InvalidArgumentException("Cannot add an Entity in an ungenerated chunk");
 		}
-		$chunk->addEntity($entity);
+		$this->entitiesByChunk[World::chunkHash($pos->getFloorX() >> 4, $pos->getFloorZ() >> 4)][$entity->getId()] = $entity;
 		$this->entityLastKnownPositions[$entity->getId()] = $pos;
 
 		if($entity instanceof Player){
@@ -2294,9 +2309,12 @@ class World implements ChunkManager{
 			throw new \InvalidArgumentException("Entity is not tracked by this world (possibly already removed?)");
 		}
 		$pos = $this->entityLastKnownPositions[$entity->getId()];
-		$chunk = $this->getChunk($pos->getFloorX() >> 4, $pos->getFloorZ() >> 4);
+		$chunkHash = World::chunkHash($pos->getFloorX() >> 4, $pos->getFloorZ() >> 4);
-		if($chunk !== null){ //we don't care if the chunk already went out of scope
+		if(isset($this->entitiesByChunk[$chunkHash][$entity->getId()])){
-			$chunk->removeEntity($entity);
+			unset($this->entitiesByChunk[$chunkHash][$entity->getId()]);
+			if(count($this->entitiesByChunk[$chunkHash]) === 0){
+				unset($this->entitiesByChunk[$chunkHash]);
+			}
 		}
 		unset($this->entityLastKnownPositions[$entity->getId()]);
 
@@ -2326,21 +2344,14 @@ class World implements ChunkManager{
 		$newChunkZ = $newPosition->getFloorZ() >> 4;
 
 		if($oldChunkX !== $newChunkX || $oldChunkZ !== $newChunkZ){
-			$oldChunk = $this->getChunk($oldChunkX, $oldChunkZ);
+			$oldChunkHash = World::chunkHash($oldChunkX, $oldChunkZ);
-			if($oldChunk !== null){
+			if(isset($this->entitiesByChunk[$oldChunkHash][$entity->getId()])){
-				$oldChunk->removeEntity($entity);
+				unset($this->entitiesByChunk[$oldChunkHash][$entity->getId()]);
+				if(count($this->entitiesByChunk[$oldChunkHash]) === 0){
+					unset($this->entitiesByChunk[$oldChunkHash]);
 				}
-			$newChunk = $this->loadChunk($newChunkX, $newChunkZ);
+			}
-			if($newChunk === null){
+
-				//TODO: this is a non-ideal solution for a hard problem
-				//when this happens the entity won't be tracked by any chunk, so we can't have it hanging around in memory
-				//we also can't allow this to cause chunk generation, nor can we just create an empty ungenerated chunk
-				//for it, because an empty chunk won't get saved, so the entity will vanish anyway. Therefore, this is
-				//the cleanest way to make sure this doesn't result in leaks.
-				$this->logger->debug("Entity " . $entity->getId() . " is in ungenerated terrain, flagging for despawn");
-				$entity->flagForDespawn();
-				$entity->despawnFromAll();
-			}else{
 			$newViewers = $this->getViewersForPosition($newPosition);
 			foreach($entity->getViewers() as $player){
 				if(!isset($newViewers[spl_object_id($player)])){
@@ -2353,8 +2364,8 @@ class World implements ChunkManager{
 				$entity->spawnTo($player);
 			}
 
-				$newChunk->addEntity($entity);
+			$newChunkHash = World::chunkHash($newChunkX, $newChunkZ);
-			}
+			$this->entitiesByChunk[$newChunkHash][$entity->getId()] = $entity;
 		}
 		$this->entityLastKnownPositions[$entity->getId()] = $newPosition->asVector3();
 	}
@@ -2564,7 +2575,7 @@ class World implements ChunkManager{
 				try{
 					$this->provider->saveChunk($x, $z, new ChunkData(
 						$chunk,
-						array_map(fn(Entity $e) => $e->saveNBT(), $chunk->getSavableEntities()),
+						array_map(fn(Entity $e) => $e->saveNBT(), array_filter($this->getChunkEntities($x, $z), fn(Entity $e) => $e->canSaveWithChunk())),
 						array_map(fn(Tile $t) => $t->saveNBT(), $chunk->getTiles()),
 					));
 				}finally{
@@ -2576,6 +2587,13 @@ class World implements ChunkManager{
 				$listener->onChunkUnloaded($x, $z, $chunk);
 			}
 
+			foreach($this->getChunkEntities($x, $z) as $entity){
+				if($entity instanceof Player){
+					continue;
+				}
+				$entity->close();
+			}
+
 			$chunk->onUnload();
 		}
 

src/world/format/Chunk.php

@@ -30,10 +30,7 @@ use pocketmine\block\Block;
 use pocketmine\block\BlockLegacyIds;
 use pocketmine\block\tile\Tile;
 use pocketmine\data\bedrock\BiomeIds;
-use pocketmine\entity\Entity;
-use pocketmine\player\Player;
 use function array_fill;
-use function array_filter;
 use function array_map;
 
 class Chunk{
@@ -59,9 +56,6 @@ class Chunk{
 	/** @var Tile[] */
 	protected $tiles = [];
 
-	/** @var Entity[] */
-	protected $entities = [];
-
 	/** @var HeightArray */
 	protected $heightMap;
 
@@ -191,17 +185,6 @@ class Chunk{
 		$this->terrainDirtyFlags |= self::DIRTY_FLAG_TERRAIN;
 	}
 
-	public function addEntity(Entity $entity) : void{
-		if($entity->isClosed()){
-			throw new \InvalidArgumentException("Attempted to add a garbage closed Entity to a chunk");
-		}
-		$this->entities[$entity->getId()] = $entity;
-	}
-
-	public function removeEntity(Entity $entity) : void{
-		unset($this->entities[$entity->getId()]);
-	}
-
 	public function addTile(Tile $tile) : void{
 		if($tile->isClosed()){
 			throw new \InvalidArgumentException("Attempted to add a garbage closed Tile to a chunk");
@@ -219,22 +202,6 @@ class Chunk{
 		unset($this->tiles[Chunk::blockHash($pos->x, $pos->y, $pos->z)]);
 	}
 
-	/**
-	 * Returns an array of entities currently using this chunk.
-	 *
-	 * @return Entity[]
-	 */
-	public function getEntities() : array{
-		return $this->entities;
-	}
-
-	/**
-	 * @return Entity[]
-	 */
-	public function getSavableEntities() : array{
-		return array_filter($this->entities, function(Entity $entity) : bool{ return $entity->canSaveWithChunk(); });
-	}
-
 	/**
 	 * @return Tile[]
 	 */
@@ -257,13 +224,6 @@ class Chunk{
 	 * Called when the chunk is unloaded, closing entities and tiles.
 	 */
 	public function onUnload() : void{
-		foreach($this->getEntities() as $entity){
-			if($entity instanceof Player){
-				continue;
-			}
-			$entity->close();
-		}
-
 		foreach($this->getTiles() as $tile){
 			$tile->close();
 		}

tests/phpstan/configs/l8-baseline.neon

@@ -195,11 +195,6 @@ parameters:
 			count: 1
 			path: ../../../src/permission/DefaultPermissions.php
 
-		-
-			message: "#^Cannot call method getEntities\\(\\) on pocketmine\\\\world\\\\format\\\\Chunk\\|null\\.$#"
-			count: 1
-			path: ../../../src/player/Player.php
-
 		-
 			message: "#^Cannot call method getSpawnLocation\\(\\) on pocketmine\\\\world\\\\World\\|null\\.$#"
 			count: 1
@@ -265,11 +260,6 @@ parameters:
 			count: 1
 			path: ../../../src/world/Explosion.php
 
-		-
-			message: "#^Cannot call method getEntities\\(\\) on pocketmine\\\\world\\\\format\\\\Chunk\\|null\\.$#"
-			count: 3
-			path: ../../../src/world/World.php
-
 		-
 			message: "#^Parameter \\#3 \\$chunk of method pocketmine\\\\player\\\\Player\\:\\:onChunkChanged\\(\\) expects pocketmine\\\\world\\\\format\\\\Chunk, pocketmine\\\\world\\\\format\\\\Chunk\\|null given\\.$#"
 			count: 1