Improve the PaperChunkSnapshotProvider

this measurably speeds up render performance
jpenilla · Jun 21, 2022 · 472f561 · 472f561
1 parent 9efd3cd
commit 472f561
Showing 1 changed file with 23 additions and 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
@@ -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()));
     }
 }