From b3d6644972dcaf708a4ced8106b57459afe00d6d Mon Sep 17 00:00:00 2001
From: Kenzie Togami <ket1999@gmail.com>
Date: Sat, 25 Jun 2016 11:37:07 -0700
Subject: [PATCH] [Forge] Fix //regen + improve performance.

---
 .../com/sk89q/worldedit/forge/ForgeWorld.java | 58 ++++++++++++-------
 .../main/resources/META-INF/worldedit_at.cfg  |  1 +
 2 files changed, 39 insertions(+), 20 deletions(-)

diff --git a/worldedit-forge/src/main/java/com/sk89q/worldedit/forge/ForgeWorld.java b/worldedit-forge/src/main/java/com/sk89q/worldedit/forge/ForgeWorld.java
index 6fe10f8ff..49a004a6d 100644
--- a/worldedit-forge/src/main/java/com/sk89q/worldedit/forge/ForgeWorld.java
+++ b/worldedit-forge/src/main/java/com/sk89q/worldedit/forge/ForgeWorld.java
@@ -42,7 +42,6 @@ import java.lang.ref.WeakReference;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Random;
-import java.util.Set;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import javax.annotation.Nullable;
@@ -54,15 +53,21 @@ import net.minecraft.block.BlockPlanks;
 import net.minecraft.block.state.IBlockState;
 import net.minecraft.entity.EntityList;
 import net.minecraft.entity.item.EntityItem;
+import net.minecraft.entity.player.EntityPlayerMP;
 import net.minecraft.init.Blocks;
 import net.minecraft.inventory.IInventory;
 import net.minecraft.nbt.NBTTagCompound;
+import net.minecraft.network.play.server.SPacketChunkData;
+import net.minecraft.server.management.PlayerChunkMap;
+import net.minecraft.server.management.PlayerChunkMapEntry;
 import net.minecraft.tileentity.TileEntity;
 import net.minecraft.util.math.BlockPos;
 import net.minecraft.util.math.ChunkPos;
 import net.minecraft.world.World;
+import net.minecraft.world.WorldServer;
 import net.minecraft.world.biome.Biome;
 import net.minecraft.world.chunk.Chunk;
+import net.minecraft.world.chunk.IChunkGenerator;
 import net.minecraft.world.chunk.IChunkProvider;
 import net.minecraft.world.gen.ChunkProviderServer;
 import net.minecraft.world.gen.feature.WorldGenBigMushroom;
@@ -238,6 +243,11 @@ public class ForgeWorld extends AbstractWorld {
 
     @Override
     public boolean regenerate(Region region, EditSession editSession) {
+        // Don't even try to regen if it's going to fail.
+        IChunkProvider provider = getWorld().getChunkProvider();
+        if (!(provider instanceof ChunkProviderServer)) {
+            return false;
+        }
         BaseBlock[] history = new BaseBlock[256 * (getMaxY() + 1)];
 
         for (Vector2D chunk : region.getChunks()) {
@@ -252,28 +262,29 @@ public class ForgeWorld extends AbstractWorld {
                     }
                 }
             }
+            PlayerChunkMap playerManager = ((WorldServer) getWorld()).getPlayerChunkMap();
+            List<EntityPlayerMP> oldWatchers = null;
+            Chunk mcChunk = null;
             try {
-                Set<Vector2D> chunks = region.getChunks();
-                IChunkProvider provider = getWorld().getChunkProvider();
-                if (!(provider instanceof ChunkProviderServer)) {
-                    return false;
-                }
                 ChunkProviderServer chunkServer = (ChunkProviderServer) provider;
-                for (Vector2D coord : chunks) {
-                    long pos = ChunkPos.chunkXZ2Int(coord.getBlockX(), coord.getBlockZ());
-                    Chunk mcChunk;
-                    if (chunkServer.chunkExists(coord.getBlockX(), coord.getBlockZ())) {
-                        mcChunk = chunkServer.loadChunk(coord.getBlockX(), coord.getBlockZ());
-                        mcChunk.onChunkUnload();
-                    }
-                    chunkServer.droppedChunksSet.remove(pos);
-                    chunkServer.id2ChunkMap.remove(pos);
-                    mcChunk = chunkServer.provideChunk(coord.getBlockX(), coord.getBlockZ());
-                    chunkServer.id2ChunkMap.put(pos, mcChunk);
-                    if (mcChunk != null) {
-                        mcChunk.onChunkLoad();
-                        mcChunk.populateChunk(chunkServer, chunkServer.chunkGenerator);
+                IChunkGenerator gen = chunkServer.chunkGenerator;
+                long pos = ChunkPos.chunkXZ2Int(chunk.getBlockX(), chunk.getBlockZ());
+                if (chunkServer.chunkExists(chunk.getBlockX(), chunk.getBlockZ())) {
+                    mcChunk = chunkServer.loadChunk(chunk.getBlockX(), chunk.getBlockZ());
+                    PlayerChunkMapEntry entry = playerManager.getEntry(chunk.getBlockX(), chunk.getBlockZ());
+                    if (entry != null) {
+                        oldWatchers = entry.players;
+                        playerManager.removeEntry(entry);
                     }
+                    mcChunk.onChunkUnload();
+                }
+                chunkServer.droppedChunksSet.remove(pos);
+                chunkServer.id2ChunkMap.remove(pos);
+                mcChunk = gen.provideChunk(chunk.getBlockX(), chunk.getBlockZ());
+                chunkServer.id2ChunkMap.put(pos, mcChunk);
+                if (mcChunk != null) {
+                    mcChunk.onChunkLoad();
+                    mcChunk.populateChunk(chunkServer, chunkServer.chunkGenerator);
                 }
             } catch (Throwable t) {
                 logger.log(Level.WARNING, "Failed to generate chunk", t);
@@ -294,6 +305,13 @@ public class ForgeWorld extends AbstractWorld {
                     }
                 }
             }
+            // We don't need to recreate the ChunkMapEntry unless there are players
+            // but addPlayer handles that for us
+            if (oldWatchers != null) {
+                for (EntityPlayerMP player : oldWatchers) {
+                    playerManager.addPlayer(player);
+                }
+            }
         }
 
         return false;
diff --git a/worldedit-forge/src/main/resources/META-INF/worldedit_at.cfg b/worldedit-forge/src/main/resources/META-INF/worldedit_at.cfg
index 837da7c36..1d2c1b157 100644
--- a/worldedit-forge/src/main/resources/META-INF/worldedit_at.cfg
+++ b/worldedit-forge/src/main/resources/META-INF/worldedit_at.cfg
@@ -1 +1,2 @@
 public net.minecraft.world.gen.ChunkProviderServer field_73248_b # droppedChunksSet
+public net.minecraft.server.management.PlayerChunkMapEntry field_187283_c # players