Change BlockBatch to use ChunkBatch instead of its own internal representation

src/main/java/net/minestom/server/instance/batch/BlockBatch.java

@@ -1,35 +1,23 @@
 package net.minestom.server.instance.batch;
 
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.concurrent.atomic.AtomicInteger;
-
-import org.jetbrains.annotations.NotNull;
-import org.jetbrains.annotations.Nullable;
-
 import net.minestom.server.data.Data;
 import net.minestom.server.instance.Chunk;
-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;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
 
-/**
- * Used when the blocks you want to place need to be divided in multiple chunks,
- * use a {@link ChunkBatch} instead otherwise.
- * Can be created using {@link Instance#createBlockBatch()}, and executed with {@link #flush(Runnable)}.
- *
- * @see InstanceBatch
- */
 public class BlockBatch implements InstanceBatch {
-
     private final InstanceContainer instance;
     private final BatchOption batchOption;
 
-    private final Map<Chunk, List<BlockData>> data = new HashMap<>();
+    // In the form of <Chunk Index, Batch>
+    private final Map<Long, ChunkBatch> data = new HashMap<>();
 
     public BlockBatch(@NotNull InstanceContainer instance, @NotNull BatchOption batchOption) {
         this.instance = instance;
@@ -60,76 +48,24 @@ public class BlockBatch implements InstanceBatch {
     }
 
     private void addBlockData(@NotNull Chunk chunk, int x, int y, int z, short blockStateId, short customBlockId, @Nullable Data data) {
-        List<BlockData> blocksData = this.data.get(chunk);
+        long chunkIndex = ChunkUtils.getChunkIndex(chunk.getChunkX(), chunk.getChunkZ());
-        if (blocksData == null)
-            blocksData = new ArrayList<>();
 
-        BlockData blockData = new BlockData();
+        ChunkBatch chunkBatch = this.data.get(chunkIndex);
-        blockData.x = x;
+        if (chunkBatch == null)
-        blockData.y = y;
+            chunkBatch = new ChunkBatch(this.instance, chunk, this.batchOption, false);
-        blockData.z = z;
-        blockData.blockStateId = blockStateId;
-        blockData.customBlockId = customBlockId;
-        blockData.data = data;
 
-        blocksData.add(blockData);
+        int relativeX = x - (chunk.getChunkX() * Chunk.CHUNK_SIZE_X);
+        int relativeZ = z - (chunk.getChunkZ() * Chunk.CHUNK_SIZE_Z);
+        chunkBatch.setSeparateBlocks(relativeX, y, relativeZ, blockStateId, customBlockId, data);
 
-        this.data.put(chunk, blocksData);
+        this.data.put(chunkIndex, chunkBatch);
     }
 
     public void flush(@Nullable Runnable callback) {
-    	this.flush(callback, false);
-    }
-    
-    public void flush(@Nullable Runnable callback, boolean shouldLoadChunks) {
         synchronized (data) {
-        	// Load chunks if applicable
+            AtomicInteger counter = new AtomicInteger();
-        	if (shouldLoadChunks) {
+            for (ChunkBatch chunkBatch : data.values()) {
-        		
+                chunkBatch.flush(c -> {
-        		// 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, indexs, 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
@@ -139,23 +75,8 @@ public class BlockBatch implements InstanceBatch {
                             this.instance.scheduleNextTick(inst -> callback.run());
                         }
                     }
-
+                });
-                }
+            }
-            });
         }
     }
-
-    private static class BlockData {
-
-        private int x, y, z;
-        private short blockStateId;
-        private short customBlockId;
-        private Data data;
-
-        public void apply(Chunk chunk) {
-            chunk.UNSAFE_setBlock(x, y, z, blockStateId, customBlockId, data, CustomBlockUtils.hasUpdate(customBlockId));
-        }
-
-    }
-
 }

src/test/java/demo/Main.java

@@ -41,6 +41,7 @@ public class Main {
         commandManager.register(new PlayersCommand());
         commandManager.register(new PotionCommand());
         commandManager.register(new TitleCommand());
+        commandManager.register(new CubeBatchCommand());
         commandManager.register(new BookCommand());
         commandManager.register(new ShootCommand());
         commandManager.register(new HorseCommand());

src/test/java/demo/commands/CubeBatchCommand.java

@@ -0,0 +1,42 @@
+package demo.commands;
+
+import net.minestom.server.chat.ChatColor;
+import net.minestom.server.chat.ColoredText;
+import net.minestom.server.command.CommandSender;
+import net.minestom.server.command.builder.Arguments;
+import net.minestom.server.command.builder.Command;
+import net.minestom.server.entity.Player;
+import net.minestom.server.instance.InstanceContainer;
+import net.minestom.server.instance.batch.BlockBatch;
+
+import java.util.concurrent.ThreadLocalRandom;
+
+public class CubeBatchCommand extends Command {
+
+    public CubeBatchCommand() {
+        super("cube");
+
+        setDefaultExecutor(this::execute);
+    }
+
+    private void execute(CommandSender sender, Arguments args) {
+        if (!sender.isPlayer()) {
+            sender.sendMessage("This command may only be run by players.");
+            return;
+        }
+        Player player = sender.asPlayer();
+
+        BlockBatch batch = new BlockBatch((InstanceContainer) player.getInstance());
+
+        int offset = 50;
+        for (int x = 0; x < 50; x += 2) {
+            for (int y = 0; y < 50; y += 2) {
+                for (int z = 0; z < 50; z += 2) {
+                    batch.setBlockStateId(x + offset, y + offset, z + offset, (short) ThreadLocalRandom.current().nextInt(500));
+                }
+            }
+        }
+
+        batch.flush(() -> sender.sendMessage(ColoredText.of(ChatColor.BRIGHT_GREEN, "Created cube.")));
+    }
+}