Folia/patches/server/0022-fixup-Rewrite-chunk-system.patch

From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: Spottedleaf <Spottedleaf@users.noreply.github.com>
Date: Sun, 18 Jun 2023 23:25:29 -0700
Subject: [PATCH] fixup! Rewrite chunk system


diff --git a/src/main/java/io/papermc/paper/chunk/system/entity/EntityLookup.java b/src/main/java/io/papermc/paper/chunk/system/entity/EntityLookup.java
index 26d1e1af418f980b61a57479cbc64b5bc59e0864..0cdf092d5c3f65a23e26f201768c0e2fea1ffe09 100644
--- a/src/main/java/io/papermc/paper/chunk/system/entity/EntityLookup.java
+++ b/src/main/java/io/papermc/paper/chunk/system/entity/EntityLookup.java
@@ -17,6 +17,7 @@ import net.minecraft.util.AbortableIterationConsumer;
 import net.minecraft.util.Mth;
 import net.minecraft.world.entity.Entity;
 import net.minecraft.world.entity.EntityType;
+import net.minecraft.world.level.ChunkPos;
 import net.minecraft.world.level.entity.EntityInLevelCallback;
 import net.minecraft.world.level.entity.EntityTypeTest;
 import net.minecraft.world.level.entity.LevelCallback;
@@ -24,6 +25,7 @@ import net.minecraft.server.level.FullChunkStatus;
 import net.minecraft.world.level.entity.LevelEntityGetter;
 import net.minecraft.world.level.entity.Visibility;
 import net.minecraft.world.phys.AABB;
+import net.minecraft.world.phys.Vec3;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.slf4j.Logger;
@@ -319,22 +321,55 @@ public final class EntityLookup implements LevelEntityGetter<Entity> {
         this.getChunk(x, z).updateStatus(newStatus, this);
     }
 
-    public void addLegacyChunkEntities(final List<Entity> entities) {
-        for (int i = 0, len = entities.size(); i < len; ++i) {
-            this.addEntity(entities.get(i), true);
-        }
+    public void addLegacyChunkEntities(final List<Entity> entities, final ChunkPos forChunk) {
+        this.addEntityChunk(entities, forChunk, true);
     }
 
-    public void addEntityChunkEntities(final List<Entity> entities) {
-        for (int i = 0, len = entities.size(); i < len; ++i) {
-            this.addEntity(entities.get(i), true);
+    public void addEntityChunkEntities(final List<Entity> entities, final ChunkPos forChunk) {
+        this.addEntityChunk(entities, forChunk, true);
+    }
+
+    public void addWorldGenChunkEntities(final List<Entity> entities, final ChunkPos forChunk) {
+        this.addEntityChunk(entities, forChunk, false);
+    }
+
+    private void addRecursivelySafe(final Entity root, final boolean fromDisk) {
+        if (!this.addEntity(root, fromDisk)) {
+            // possible we are a passenger, and so should dismount from any valid entity in the world
+            root.stopRiding(true);
+            return;
+        }
+        for (final Entity passenger : root.getPassengers()) {
+            this.addRecursivelySafe(passenger, fromDisk);
         }
     }
 
-    public void addWorldGenChunkEntities(final List<Entity> entities) {
+    private void addEntityChunk(final List<Entity> entities, final ChunkPos forChunk, final boolean fromDisk) {
         for (int i = 0, len = entities.size(); i < len; ++i) {
-            this.addEntity(entities.get(i), false);
-        }
+            final Entity entity = entities.get(i);
+            if (entity.isPassenger()) {
+                continue;
+            }
+
+            if (!entity.chunkPosition().equals(forChunk)) {
+                LOGGER.warn("Root entity " + entity + " is outside of serialized chunk " + forChunk);
+                // can't set removed here, as we may not own the chunk position
+                // skip the entity
+                continue;
+            }
+
+            final Vec3 rootPosition = entity.position();
+
+            // always adjust positions before adding passengers in case plugins access the entity, and so that
+            // they are added to the right entity chunk
+            for (final Entity passenger : entity.getIndirectPassengers()) {
+                if (!passenger.chunkPosition().equals(forChunk)) {
+                    passenger.setPosRaw(rootPosition.x, rootPosition.y, rootPosition.z, true);
+                }
+            }
+
+            this.addRecursivelySafe(entity, fromDisk);
+         }
     }
 
     public boolean addNewEntity(final Entity entity) {
diff --git a/src/main/java/io/papermc/paper/chunk/system/scheduling/ChunkFullTask.java b/src/main/java/io/papermc/paper/chunk/system/scheduling/ChunkFullTask.java
index 0b78d1eb90500e0123b7281d722805dc65d551d0..8a48238fac0949cfddcdd9ccf179d16befe9c4d4 100644
--- a/src/main/java/io/papermc/paper/chunk/system/scheduling/ChunkFullTask.java
+++ b/src/main/java/io/papermc/paper/chunk/system/scheduling/ChunkFullTask.java
@@ -83,7 +83,7 @@ public final class ChunkFullTask extends ChunkProgressionTask implements Runnabl
                 final ServerLevel world = this.world;
                 final ProtoChunk protoChunk = (ProtoChunk)this.fromChunk;
                 chunk = new LevelChunk(this.world, protoChunk, (final LevelChunk unused) -> {
-                    ChunkMap.postLoadProtoChunk(world, protoChunk.getEntities());
+                    ChunkMap.postLoadProtoChunk(world, protoChunk.getEntities(), protoChunk.getPos()); // Paper - rewrite chunk system
                 });
             }
 
diff --git a/src/main/java/io/papermc/paper/chunk/system/scheduling/NewChunkHolder.java b/src/main/java/io/papermc/paper/chunk/system/scheduling/NewChunkHolder.java
index f1c68d9850ece7532a8607db955eaa4fc3a4bf05..08075b8895f816420c2a940bf551dfada3c0cd9e 100644
--- a/src/main/java/io/papermc/paper/chunk/system/scheduling/NewChunkHolder.java
+++ b/src/main/java/io/papermc/paper/chunk/system/scheduling/NewChunkHolder.java
@@ -115,7 +115,7 @@ public final class NewChunkHolder {
             if (entityChunk != null) {
                 final List<Entity> entities = EntityStorage.readEntities(this.world, entityChunk);
 
-                this.world.getEntityLookup().addEntityChunkEntities(entities);
+                this.world.getEntityLookup().addEntityChunkEntities(entities, new ChunkPos(this.chunkX, this.chunkZ));
             }
         }
 
diff --git a/src/main/java/net/minecraft/server/level/ChunkMap.java b/src/main/java/net/minecraft/server/level/ChunkMap.java
index 779c03f17769039cbaaa1210bd0e95163c9386b1..f1f52f9d9f8eafd65b50e3c06c92bcd8fda64286 100644
--- a/src/main/java/net/minecraft/server/level/ChunkMap.java
+++ b/src/main/java/net/minecraft/server/level/ChunkMap.java
@@ -597,7 +597,7 @@ public class ChunkMap extends ChunkStorage implements ChunkHolder.PlayerProvider
         return chunkstatus1;
     }
 
-    public static void postLoadProtoChunk(ServerLevel world, List<CompoundTag> nbt) { // Paper - public
+    public static void postLoadProtoChunk(ServerLevel world, List<CompoundTag> nbt, ChunkPos position) { // Paper - public and add chunk position parameter
         if (!nbt.isEmpty()) {
             // CraftBukkit start - these are spawned serialized (DefinedStructure) and we don't call an add event below at the moment due to ordering complexities
             world.addWorldGenChunkEntities(EntityType.loadEntitiesRecursive(nbt, world).filter((entity) -> {
@@ -613,7 +613,7 @@ public class ChunkMap extends ChunkStorage implements ChunkHolder.PlayerProvider
                 }
                 checkDupeUUID(world, entity); // Paper
                 return !needsRemoval;
-            }));
+            }), position); // Paper - rewrite chunk system
             // CraftBukkit end
         }
 
diff --git a/src/main/java/net/minecraft/server/level/ServerLevel.java b/src/main/java/net/minecraft/server/level/ServerLevel.java
index bb07ad1bb895297356b88dfc4cd17e5e93795c38..60b409e873d8f5d7975cce6e86ee2ed5525d348c 100644
--- a/src/main/java/net/minecraft/server/level/ServerLevel.java
+++ b/src/main/java/net/minecraft/server/level/ServerLevel.java
@@ -2663,12 +2663,12 @@ public class ServerLevel extends Level implements WorldGenLevel {
         return this.entityLookup; // Paper - rewrite chunk system
     }
 
-    public void addLegacyChunkEntities(Stream<Entity> entities) {
-        this.entityLookup.addLegacyChunkEntities(entities.toList()); // Paper - rewrite chunk system
+    public void addLegacyChunkEntities(Stream<Entity> entities, ChunkPos forChunk) { // Paper - rewrite chunk system
+        this.entityLookup.addLegacyChunkEntities(entities.toList(), forChunk); // Paper - rewrite chunk system
     }
 
-    public void addWorldGenChunkEntities(Stream<Entity> entities) {
-        this.entityLookup.addWorldGenChunkEntities(entities.toList()); // Paper - rewrite chunk system
+    public void addWorldGenChunkEntities(Stream<Entity> entities, ChunkPos forChunk) { // Paper - rewrite chunk system
+        this.entityLookup.addWorldGenChunkEntities(entities.toList(), forChunk); // Paper - rewrite chunk system
     }
 
     public void startTickingChunk(LevelChunk chunk) {
diff --git a/src/main/java/net/minecraft/world/level/chunk/storage/ChunkSerializer.java b/src/main/java/net/minecraft/world/level/chunk/storage/ChunkSerializer.java
index 035b06d00aadc67eb247efd7b25aea92ba0532a2..97533c4a31ab6137072b79dc4377fd602fef9ea8 100644
--- a/src/main/java/net/minecraft/world/level/chunk/storage/ChunkSerializer.java
+++ b/src/main/java/net/minecraft/world/level/chunk/storage/ChunkSerializer.java
@@ -678,7 +678,7 @@ public class ChunkSerializer {
 
         return nbttaglist == null && nbttaglist1 == null ? null : (chunk) -> {
             if (nbttaglist != null) {
-                world.addLegacyChunkEntities(EntityType.loadEntitiesRecursive(nbttaglist, world));
+                world.addLegacyChunkEntities(EntityType.loadEntitiesRecursive(nbttaglist, world), chunk.getPos()); // Paper - rewrite chunk system
             }
 
             if (nbttaglist1 != null) {