From 472f561181b2b17218c2c71ca08ce85e8040a6ed Mon Sep 17 00:00:00 2001
From: Jason Penilla <11360596+jpenilla@users.noreply.github.com>
Date: Mon, 20 Jun 2022 22:00:59 -0700
Subject: [PATCH] Improve the PaperChunkSnapshotProvider

this measurably speeds up render performance
---
 .../util/PaperChunkSnapshotProvider.java      | 35 ++++++++++++-------
 1 file changed, 23 insertions(+), 12 deletions(-)

diff --git a/paper/src/main/java/xyz/jpenilla/squaremap/paper/util/PaperChunkSnapshotProvider.java b/paper/src/main/java/xyz/jpenilla/squaremap/paper/util/PaperChunkSnapshotProvider.java
index af6442cd..306a5eb7 100644
--- a/paper/src/main/java/xyz/jpenilla/squaremap/paper/util/PaperChunkSnapshotProvider.java
+++ b/paper/src/main/java/xyz/jpenilla/squaremap/paper/util/PaperChunkSnapshotProvider.java
@@ -3,9 +3,10 @@
 import com.google.inject.Inject;
 import com.google.inject.Singleton;
 import java.util.concurrent.CompletableFuture;
-import java.util.function.Function;
-import java.util.function.Supplier;
 import net.minecraft.server.level.ServerLevel;
+import net.minecraft.world.level.chunk.ChunkAccess;
+import net.minecraft.world.level.chunk.ChunkStatus;
+import net.minecraft.world.level.chunk.ImposterProtoChunk;
 import net.minecraft.world.level.chunk.LevelChunk;
 import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
@@ -27,15 +28,25 @@ private PaperChunkSnapshotProvider() {
         final int z,
         final boolean biomesOnly
     ) {
-        final Supplier<CompletableFuture<@Nullable ChunkSnapshot>> futureSupplier = () -> level.getChunkSource()
-            .getChunkAtAsynchronously(x, z, false, false)
-            .thenApply(either -> {
-                final @Nullable LevelChunk chunk = (LevelChunk) either.left().orElse(null);
-                if (chunk == null || chunk.isEmpty()) {
-                    return null;
-                }
-                return ChunkSnapshot.snapshot(chunk, biomesOnly);
-            });
-        return CompletableFuture.supplyAsync(futureSupplier, level.getServer()).thenCompose(Function.identity());
+        return CompletableFuture.supplyAsync(() -> {
+            @Nullable ChunkAccess existing = level.getChunkIfLoadedImmediately(x, z);
+            if (existing == null) {
+                existing = level.getChunkSource().chunkMap.getUnloadingChunk(x, z);
+            }
+            if (existing != null && existing.getStatus().isOrAfter(ChunkStatus.FULL)) {
+                return CompletableFuture.completedFuture(existing);
+            }
+            final CompletableFuture<@Nullable ChunkAccess> load = new CompletableFuture<>();
+            level.getChunkSource().getChunkAtAsynchronously(x, z, ChunkStatus.EMPTY, false, false, load::complete);
+            return load;
+        }, level.getServer()).thenCompose(chunkFuture -> chunkFuture.thenApplyAsync(chunk -> {
+            if (chunk == null || !chunk.getStatus().isOrAfter(ChunkStatus.FULL)) {
+                return null;
+            }
+            if (chunk instanceof ImposterProtoChunk imposter) {
+                chunk = imposter.getWrapped();
+            }
+            return ChunkSnapshot.snapshot((LevelChunk) chunk, biomesOnly);
+        }, level.getServer()));
     }
 }