From ee8815d2f8b593add99906d41760a1cb70556840 Mon Sep 17 00:00:00 2001
From: KrystilizeNevaDies
 <57762380+KrystilizeNevaDies@users.noreply.github.com>
Date: Fri, 12 Feb 2021 10:37:53 +1000
Subject: [PATCH] Add shouldLoadChunks in BlockBatch#flush

---
 .../server/instance/batch/BlockBatch.java     | 99 ++++++++++++-------
 1 file changed, 64 insertions(+), 35 deletions(-)

diff --git a/src/main/java/net/minestom/server/instance/batch/BlockBatch.java b/src/main/java/net/minestom/server/instance/batch/BlockBatch.java
index 2d49418b5..cc62ebdfd 100644
--- a/src/main/java/net/minestom/server/instance/batch/BlockBatch.java
+++ b/src/main/java/net/minestom/server/instance/batch/BlockBatch.java
@@ -6,6 +6,8 @@ import net.minestom.server.instance.Instance;
 import net.minestom.server.instance.InstanceContainer;
 import net.minestom.server.instance.block.CustomBlock;
 import net.minestom.server.utils.block.CustomBlockUtils;
+import net.minestom.server.utils.chunk.ChunkUtils;
+
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -13,6 +15,7 @@ import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
 
 /**
@@ -76,44 +79,70 @@ public class BlockBatch implements InstanceBatch {
     }
 
     public void flush(@Nullable Runnable callback) {
+    	this.flush(callback, false);
+    }
+    
+    public void flush(@Nullable Runnable callback, boolean shouldLoadChunks) {
         synchronized (data) {
-            AtomicInteger counter = new AtomicInteger();
-            for (Map.Entry<Chunk, List<BlockData>> entry : data.entrySet()) {
-                final Chunk chunk = entry.getKey();
-                final List<BlockData> dataList = entry.getValue();
-                BLOCK_BATCH_POOL.execute(() -> {
-                    synchronized (chunk) {
-                        if (!chunk.isLoaded())
-                            return;
-
-                        if (batchOption.isFullChunk()) {
-                            chunk.reset();
-                        }
-
-                        for (BlockData data : dataList) {
-                            data.apply(chunk);
-                        }
-
-                        // Refresh chunk for viewers
-                        if (batchOption.isFullChunk()) {
-                            chunk.sendChunk();
-                        } else {
-                            chunk.sendChunkUpdate();
-                        }
-
-                        final boolean isLast = counter.incrementAndGet() == data.size();
-
-                        // Execute the callback if this was the last chunk to process
-                        if (isLast) {
-                            this.instance.refreshLastBlockChangeTime();
-                            if (callback != null) {
-                                this.instance.scheduleNextTick(inst -> callback.run());
-                            }
-                        }
+        	// Load chunks if applicable
+        	if (shouldLoadChunks) {
+        		
+        		// Get chunks
+            	Chunk[] chunks = data.keySet().toArray(new Chunk[data.size()]);
+            	
+            	// Get chunk indexs
+            	Long[] indexs = new Long[chunks.length];
+            	for (int i = 0; i < chunks.length; i++) {
+            		indexs[i] = ChunkUtils.getChunkIndex(chunks[i].getChunkX(), chunks[i].getChunkZ());
+            	}
+            	
+            	// Load all chunks + flush block data
+        		ChunkUtils.optionalLoadAll(instance, null, null, (chunk) -> {
+        			flushBlockData(callback);
+        		});
+        	} else {
+        		flushBlockData(callback);
+        	}
+        }
+    }
+    
+    private void flushBlockData(@Nullable Runnable callback) {
+    	AtomicInteger counter = new AtomicInteger();
+        for (Map.Entry<Chunk, List<BlockData>> entry : data.entrySet()) {
+            final Chunk chunk = entry.getKey();
+            final List<BlockData> dataList = entry.getValue();
+            BLOCK_BATCH_POOL.execute(() -> {
+                synchronized (chunk) {
+                    if (!chunk.isLoaded())
+                        return;
 
+                    if (batchOption.isFullChunk()) {
+                        chunk.reset();
                     }
-                });
-            }
+
+                    for (BlockData data : dataList) {
+                        data.apply(chunk);
+                    }
+
+                    // Refresh chunk for viewers
+                    if (batchOption.isFullChunk()) {
+                        chunk.sendChunk();
+                    } else {
+                        chunk.sendChunkUpdate();
+                    }
+
+                    final boolean isLast = counter.incrementAndGet() == data.size();
+
+                    // Execute the callback if this was the last chunk to process
+                    if (isLast) {
+                        this.instance.refreshLastBlockChangeTime();
+                        if (callback != null) {
+                            this.instance.scheduleNextTick(inst -> callback.run());
+                        }
+                    }
+
+                }
+            });
         }
     }