From 99e0adb83bd59b94783876948a95b792f8cfdf6b Mon Sep 17 00:00:00 2001
From: CraftBukkit/Spigot <noreply+git-craftbukkit@papermc.io>
Date: Sun, 30 Jan 2022 10:52:42 +1100
Subject: [PATCH] SPIGOT-6891: Crash when importing 1.16 chunks with entities
 above the world, when a BlockPopulator is active

By: DerFrZocker <derrieple@gmail.com>
---
 .../generator/CraftLimitedRegion.java         | 58 ++++++++++++++-----
 1 file changed, 42 insertions(+), 16 deletions(-)

diff --git a/paper-server/src/main/java/org/bukkit/craftbukkit/generator/CraftLimitedRegion.java b/paper-server/src/main/java/org/bukkit/craftbukkit/generator/CraftLimitedRegion.java
index ed8f86bfe4..19ef0e242d 100644
--- a/paper-server/src/main/java/org/bukkit/craftbukkit/generator/CraftLimitedRegion.java
+++ b/paper-server/src/main/java/org/bukkit/craftbukkit/generator/CraftLimitedRegion.java
@@ -38,30 +38,20 @@ public class CraftLimitedRegion extends CraftRegionAccessor implements LimitedRe
     // there the order is {..., FEATURES, LIQUID_CARVERS, STRUCTURE_STARTS, ...}
     private final int buffer = 16;
     private final BoundingBox region;
+    boolean entitiesLoaded = false;
     // Minecraft saves the entities as NBTTagCompound during chunk generation. This causes that
     // changes made to the returned bukkit entity are not saved. To combat this we keep them and
     // save them when the population is finished.
     private final List<net.minecraft.world.entity.Entity> entities = new ArrayList<>();
+    // SPIGOT-6891: Save outside Entities extra, since they are not part of the region.
+    // Prevents crash for chunks which are converting from 1.17 to 1.18
+    private final List<net.minecraft.world.entity.Entity> outsideEntities = new ArrayList<>();
 
     public CraftLimitedRegion(GeneratorAccessSeed access, ChunkCoordIntPair center) {
         this.weakAccess = new WeakReference<>(access);
         centerChunkX = center.x;
         centerChunkZ = center.z;
 
-        // load entities which are already present
-        for (int x = -(buffer >> 4); x <= (buffer >> 4); x++) {
-            for (int z = -(buffer >> 4); z <= (buffer >> 4); z++) {
-                ProtoChunk chunk = (ProtoChunk) access.getChunk(centerChunkX + x, centerChunkZ + z);
-                for (NBTTagCompound compound : chunk.getEntities()) {
-                    EntityTypes.loadEntityRecursive(compound, access.getMinecraftWorld(), (entity) -> {
-                        entity.generation = true;
-                        entities.add(entity);
-                        return entity;
-                    });
-                }
-            }
-        }
-
         World world = access.getMinecraftWorld().getWorld();
         int xCenter = centerChunkX << 4;
         int zCenter = centerChunkZ << 4;
@@ -83,12 +73,42 @@ public class CraftLimitedRegion extends CraftRegionAccessor implements LimitedRe
         return handle;
     }
 
-    public void saveEntities() {
+    public void loadEntities() {
+        if (entitiesLoaded) {
+            return;
+        }
+
         GeneratorAccessSeed access = getHandle();
+        // load entities which are already present
         for (int x = -(buffer >> 4); x <= (buffer >> 4); x++) {
             for (int z = -(buffer >> 4); z <= (buffer >> 4); z++) {
                 ProtoChunk chunk = (ProtoChunk) access.getChunk(centerChunkX + x, centerChunkZ + z);
-                chunk.getEntities().clear();
+                for (NBTTagCompound compound : chunk.getEntities()) {
+                    EntityTypes.loadEntityRecursive(compound, access.getMinecraftWorld(), (entity) -> {
+                        if (region.contains(entity.getX(), entity.getY(), entity.getZ())) {
+                            entity.generation = true;
+                            entities.add(entity);
+                        } else {
+                            outsideEntities.add(entity);
+                        }
+                        return entity;
+                    });
+                }
+            }
+        }
+
+        entitiesLoaded = true;
+    }
+
+    public void saveEntities() {
+        GeneratorAccessSeed access = getHandle();
+        // We don't clear existing entities when they are not loaded and therefore not modified
+        if (entitiesLoaded) {
+            for (int x = -(buffer >> 4); x <= (buffer >> 4); x++) {
+                for (int z = -(buffer >> 4); z <= (buffer >> 4); z++) {
+                    ProtoChunk chunk = (ProtoChunk) access.getChunk(centerChunkX + x, centerChunkZ + z);
+                    chunk.getEntities().clear();
+                }
             }
         }
 
@@ -99,6 +119,10 @@ public class CraftLimitedRegion extends CraftRegionAccessor implements LimitedRe
                 access.addFreshEntity(entity);
             }
         }
+
+        for (net.minecraft.world.entity.Entity entity : outsideEntities) {
+            access.addFreshEntity(entity);
+        }
     }
 
     public void breakLink() {
@@ -187,6 +211,8 @@ public class CraftLimitedRegion extends CraftRegionAccessor implements LimitedRe
 
     @Override
     public Collection<net.minecraft.world.entity.Entity> getNMSEntities() {
+        // Only load entities if we need them
+        loadEntities();
         return new ArrayList<>(entities);
     }