From 36baa77ab7eb50eff8cbe9291eabc482576c1b8c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=80=E5=93=A5?= <3523206925@qq.com>
Date: Sat, 2 Sep 2023 00:29:19 +0800
Subject: [PATCH] refactor: Tidy AllayChunkService::loadChunk and add
 AllayChunkService::loadChunkImmediately method

---
 .../allay/api/world/chunk/ChunkService.java   |  4 +-
 .../server/world/chunk/AllayChunkService.java | 45 ++++++++++++-------
 2 files changed, 33 insertions(+), 16 deletions(-)

diff --git a/Allay-API/src/main/java/cn/allay/api/world/chunk/ChunkService.java b/Allay-API/src/main/java/cn/allay/api/world/chunk/ChunkService.java
index 26209e799..75b6cfec1 100644
--- a/Allay-API/src/main/java/cn/allay/api/world/chunk/ChunkService.java
+++ b/Allay-API/src/main/java/cn/allay/api/world/chunk/ChunkService.java
@@ -18,9 +18,11 @@ public interface ChunkService extends ChunkAccessible {
 
     CompletableFuture<Chunk> getOrLoadChunk(int x, int z);
 
-    @SlowOperation
     CompletableFuture<Chunk> loadChunk(int x, int z);
 
+    @SlowOperation
+    Chunk loadChunkImmediately(int x, int z);
+
     void unloadChunk(int x, int z);
 
     void unloadChunk(long chunkHash);
diff --git a/Allay-Server/src/main/java/cn/allay/server/world/chunk/AllayChunkService.java b/Allay-Server/src/main/java/cn/allay/server/world/chunk/AllayChunkService.java
index 3b956e467..307449f7c 100644
--- a/Allay-Server/src/main/java/cn/allay/server/world/chunk/AllayChunkService.java
+++ b/Allay-Server/src/main/java/cn/allay/server/world/chunk/AllayChunkService.java
@@ -19,11 +19,13 @@
 import io.netty.buffer.Unpooled;
 import it.unimi.dsi.fastutil.longs.*;
 import lombok.Getter;
+import lombok.SneakyThrows;
 import lombok.extern.slf4j.Slf4j;
 import org.cloudburstmc.protocol.bedrock.data.HeightMapDataType;
 import org.cloudburstmc.protocol.bedrock.data.SubChunkData;
 import org.cloudburstmc.protocol.bedrock.data.SubChunkRequestResult;
 import org.cloudburstmc.protocol.bedrock.packet.SubChunkPacket;
+import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.annotations.UnmodifiableView;
 import org.joml.Vector3i;
@@ -201,7 +203,6 @@ public CompletableFuture<Chunk> getOrLoadChunk(int x, int z) {
         return loadChunk(x, z);
     }
 
-    @SlowOperation
     @Override
     public CompletableFuture<Chunk> loadChunk(int x, int z) {
         var hashXZ = HashUtils.hashXZ(x, z);
@@ -214,23 +215,37 @@ public CompletableFuture<Chunk> loadChunk(int x, int z) {
         }
         var future = worldStorage.readChunk(x, z);
         loadingChunks.put(hashXZ, future);
-        future.thenApplyAsync(loadedChunk -> {
-            if (loadedChunk != null) {
-                setChunk(x, z, loadedChunk);
-                loadingChunks.remove(hashXZ);
-                return loadedChunk;
-            }
-            var unsafeChunk = AllayUnsafeChunk.builder().emptyChunk(x, z, getWorldStorage().getWorldDataCache().getDimensionInfo());
-            var chunkGenerateContext = new ChunkGenerateContext(unsafeChunk, world);
-            world.getWorldGenerator().generate(chunkGenerateContext);
-            var safeChunk = new AllayChunk(unsafeChunk);
-            setChunk(x, z, safeChunk);
-            loadingChunks.remove(hashXZ);
-            return safeChunk;
-        }, Server.getInstance().getComputeThreadPool());
+        future.thenApplyAsync(loadedChunk -> generateChunkIfNull(x, z, loadedChunk), Server.getInstance().getComputeThreadPool());
         return future;
     }
 
+    @SneakyThrows
+    @SlowOperation
+    @Override
+    public Chunk loadChunkImmediately(int x, int z) {
+        return generateChunkIfNull(
+                x, z,
+                worldStorage.readChunk(x, z).get()
+        );
+    }
+
+    private Chunk generateChunkIfNull(int x, int z, Chunk loadedChunk) {
+        long hashXZ = HashUtils.hashXZ(x, z);
+        if (loadedChunk == null) {
+            loadedChunk = generateChunkImmediately(x, z);
+        }
+        setChunk(x, z, loadedChunk);
+        loadingChunks.remove(hashXZ);
+        return loadedChunk;
+    }
+
+    private AllayChunk generateChunkImmediately(int x, int z) {
+        var unsafeChunk = AllayUnsafeChunk.builder().emptyChunk(x, z, getWorldStorage().getWorldDataCache().getDimensionInfo());
+        var chunkGenerateContext = new ChunkGenerateContext(unsafeChunk, world);
+        world.getWorldGenerator().generate(chunkGenerateContext);
+        return new AllayChunk(unsafeChunk);
+    }
+
     public void unloadChunk(int x, int z) {
         unloadChunk(HashUtils.hashXZ(x, z));
     }