Heavily optimize the BlockColumn impl used during surface rule evaluation

src/main/java/org/embeddedt/modernfix/common/mixin/perf/optimize_surface_rules/SurfaceSystemMixin.java

@@ -1,11 +1,14 @@
 package org.embeddedt.modernfix.common.mixin.perf.optimize_surface_rules;
 
 import com.llamalad7.mixinextras.sugar.Local;
+import com.llamalad7.mixinextras.sugar.Share;
+import com.llamalad7.mixinextras.sugar.ref.LocalRef;
 import net.minecraft.core.BlockPos;
 import net.minecraft.core.Holder;
 import net.minecraft.core.Registry;
 import net.minecraft.world.level.biome.Biome;
 import net.minecraft.world.level.biome.BiomeManager;
+import net.minecraft.world.level.chunk.BlockColumn;
 import net.minecraft.world.level.chunk.ChunkAccess;
 import net.minecraft.world.level.levelgen.NoiseChunk;
 import net.minecraft.world.level.levelgen.RandomState;
@@ -13,10 +16,12 @@ import net.minecraft.world.level.levelgen.SurfaceRules;
 import net.minecraft.world.level.levelgen.SurfaceSystem;
 import net.minecraft.world.level.levelgen.WorldGenerationContext;
 import org.embeddedt.modernfix.world.gen.ChunkBiomeLookup;
+import org.embeddedt.modernfix.world.gen.PrefetchingBlockColumn;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.injection.At;
 import org.spongepowered.asm.mixin.injection.Inject;
 import org.spongepowered.asm.mixin.injection.ModifyArg;
+import org.spongepowered.asm.mixin.injection.Redirect;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
 
 import java.util.function.Function;
@@ -24,17 +29,52 @@ import java.util.function.Function;
 @Mixin(SurfaceSystem.class)
 public class SurfaceSystemMixin {
     private static final ThreadLocal<ChunkBiomeLookup> MFIX_LOOKUP_CACHE = ThreadLocal.withInitial(ChunkBiomeLookup::new);
+    private static final ThreadLocal<PrefetchingBlockColumn> MFIX_BLOCK_COLUMN = new ThreadLocal<>();
 
     @ModifyArg(method = "buildSurface", at = @At(value = "INVOKE", target = "Lnet/minecraft/world/level/levelgen/SurfaceRules$Context;<init>(Lnet/minecraft/world/level/levelgen/SurfaceSystem;Lnet/minecraft/world/level/levelgen/RandomState;Lnet/minecraft/world/level/chunk/ChunkAccess;Lnet/minecraft/world/level/levelgen/NoiseChunk;Ljava/util/function/Function;Lnet/minecraft/core/Registry;Lnet/minecraft/world/level/levelgen/WorldGenerationContext;)V"), index = 4)
-    private Function<BlockPos, Holder<Biome>> useFasterLookup(Function<BlockPos, Holder<Biome>> biomeGetter, @Local(ordinal = 0, argsOnly = true) BiomeManager manager, @Local(ordinal = 0, argsOnly = true) ChunkAccess chunk) {
+    private Function<BlockPos, Holder<Biome>> useFasterLookup(Function<BlockPos, Holder<Biome>> biomeGetter,
+                                                              @Local(ordinal = 0, argsOnly = true) BiomeManager manager,
+                                                              @Local(ordinal = 0, argsOnly = true) ChunkAccess chunk,
+                                                              @Share("chunkBiomeLookup") LocalRef<ChunkBiomeLookup> lookupRef) {
         var lookup = MFIX_LOOKUP_CACHE.get();
         BiomeManagerAccessor accessor = (BiomeManagerAccessor)manager;
         lookup.prepare(accessor.mfix$getBiomeSource(), accessor.mfix$getZoomSeed(), chunk, manager);
+        lookupRef.set(lookup);
         return lookup;
     }
 
     @Inject(method = "buildSurface", at = @At("RETURN"))
-    private void disposeLookup(RandomState randomState, BiomeManager biomeManager, Registry<Biome> biomes, boolean p_224652_, WorldGenerationContext context, ChunkAccess chunk, NoiseChunk noiseChunk, SurfaceRules.RuleSource ruleSource, CallbackInfo ci) {
+    private void finishAndDisposeLookups(RandomState randomState, BiomeManager biomeManager, Registry<Biome> biomes, boolean p_224652_, WorldGenerationContext context, ChunkAccess chunk, NoiseChunk noiseChunk, SurfaceRules.RuleSource ruleSource, CallbackInfo ci) {
         MFIX_LOOKUP_CACHE.get().dispose();
+        var column = MFIX_BLOCK_COLUMN.get();
+        if (column != null) {
+            column.dispose();
+        }
+    }
+
+    @Redirect(method = "buildSurface", at = @At(value = "INVOKE", target = "Lnet/minecraft/world/level/biome/BiomeManager;getBiome(Lnet/minecraft/core/BlockPos;)Lnet/minecraft/core/Holder;"))
+    private Holder<Biome> useFasterLookup(BiomeManager instance, BlockPos pos, @Share("chunkBiomeLookup") LocalRef<ChunkBiomeLookup> lookupRef) {
+        return lookupRef.get().apply(pos);
+    }
+
+    @Inject(method = "buildSurface", at = @At(value = "INVOKE", target = "Lnet/minecraft/world/level/levelgen/SurfaceRules$Context;<init>(Lnet/minecraft/world/level/levelgen/SurfaceSystem;Lnet/minecraft/world/level/levelgen/RandomState;Lnet/minecraft/world/level/chunk/ChunkAccess;Lnet/minecraft/world/level/levelgen/NoiseChunk;Ljava/util/function/Function;Lnet/minecraft/core/Registry;Lnet/minecraft/world/level/levelgen/WorldGenerationContext;)V"))
+    private void captureRealBlockColumn(CallbackInfo ci, @Local(ordinal = 0) LocalRef<BlockColumn> column,
+                                        @Local(ordinal = 0, argsOnly = true) ChunkAccess chunk,
+                                        @Share("prefetchColumn") LocalRef<PrefetchingBlockColumn> prefetchRef) {
+        var prefetchingBlockColumn = MFIX_BLOCK_COLUMN.get();
+        if (prefetchingBlockColumn == null || prefetchingBlockColumn.getExpectedHeight() != chunk.getHeight()) {
+            prefetchingBlockColumn = new PrefetchingBlockColumn(chunk.getHeight());
+            MFIX_BLOCK_COLUMN.set(prefetchingBlockColumn);
+        }
+        column.set(prefetchingBlockColumn);
+        prefetchRef.set(prefetchingBlockColumn);
+    }
+
+    @Inject(method = "buildSurface", at = @At(value = "INVOKE", target = "Lnet/minecraft/core/BlockPos$MutableBlockPos;setZ(I)Lnet/minecraft/core/BlockPos$MutableBlockPos;", ordinal = 0, shift = At.Shift.AFTER))
+    private void prefetchBlockArray(RandomState randomState, BiomeManager biomeManager, Registry<Biome> biomes, boolean p_224652_,
+                                    WorldGenerationContext context, ChunkAccess chunk, NoiseChunk noiseChunk, SurfaceRules.RuleSource ruleSource, CallbackInfo ci,
+                                    @Local(ordinal = 0) BlockColumn column,
+                                    @Local(ordinal = 0) BlockPos.MutableBlockPos cursor) {
+        ((PrefetchingBlockColumn)column).prefetch(chunk, cursor.getX() & 15, cursor.getZ() & 15);
     }
 }

src/main/java/org/embeddedt/modernfix/forge/init/ModernFixForge.java

@@ -138,7 +138,7 @@ public class ModernFixForge {
     public void onServerStarted(ServerStartedEvent event) {
         commonMod.onServerStarted();
         if (Boolean.getBoolean("modernfix.runWorldgenBenchmark")) {
-            int iterations = Integer.getInteger("modernfix.worldgenIterations", 100);
+            int iterations = Integer.getInteger("modernfix.worldgenIterations", 15);
             int testRadius = Integer.getInteger("modernfix.worldgenTestRadius", 10);
             var level = event.getServer().overworld();
             ModernFix.LOGGER.info("Worldgen results: {}", WorldgenBenchmark.run(level, new ChunkPos(0, 0), testRadius, iterations,

src/main/java/org/embeddedt/modernfix/world/gen/PrefetchingBlockColumn.java

@@ -0,0 +1,126 @@
+package org.embeddedt.modernfix.world.gen;
+
+import net.minecraft.core.BlockPos;
+import net.minecraft.core.SectionPos;
+import net.minecraft.world.level.block.Blocks;
+import net.minecraft.world.level.block.state.BlockState;
+import net.minecraft.world.level.chunk.BlockColumn;
+import net.minecraft.world.level.chunk.ChunkAccess;
+import net.minecraft.world.level.chunk.LevelChunkSection;
+import net.minecraft.world.level.levelgen.Heightmap;
+
+/**
+ * Wraps a BlockColumn and prefetches all block states for a column into an array.
+ *
+ * <p>Writes bypass {@link ChunkAccess#setBlockState} and go directly to the section,
+ * skipping heightmap and light updates that are unnecessary during the surface building
+ * stage (which runs before {@code INITIALIZE_LIGHT}).
+ */
+public class PrefetchingBlockColumn implements BlockColumn {
+    private static final BlockState AIR = Blocks.AIR.defaultBlockState();
+    private static final BlockState VOID_AIR = Blocks.VOID_AIR.defaultBlockState();
+
+    private final BlockState[] states;
+    private final BlockPos.MutableBlockPos cursor = new BlockPos.MutableBlockPos();
+
+    private ChunkAccess chunk;
+    private LevelChunkSection[] sections;
+    private int minBuildHeight;
+    private int localX, localZ;
+
+    public PrefetchingBlockColumn(int height) {
+        this.states = new BlockState[height];
+    }
+
+    public int getExpectedHeight() {
+        return this.states.length;
+    }
+
+    /**
+     * Prefetch all block states for the column at the given local XZ coordinates.
+     * Must be called before any getBlock/setBlock calls for this column.
+     */
+    public void prefetch(ChunkAccess chunk, int localX, int localZ) {
+        if (chunk.getHeight() != this.states.length) {
+            throw new IllegalStateException();
+        }
+        this.chunk = chunk;
+        this.sections = chunk.getSections();
+        this.minBuildHeight = chunk.getMinBuildHeight();
+        this.localX = localX;
+        this.localZ = localZ;
+        var sections = this.sections;
+        var states = this.states;
+        int offset = 0;
+        for (LevelChunkSection section : sections) {
+            if (section.hasOnlyAir()) {
+                for (int y = 0; y < 16; y++) {
+                    states[offset + y] = AIR;
+                }
+            } else {
+                var container = section.getStates();
+                for (int y = 0; y < 16; y++) {
+                    states[offset + y] = container.get(localX, y, localZ);
+                }
+            }
+            offset += 16;
+        }
+    }
+
+    /**
+     * Clear cached references to allow GC of the chunk between uses.
+     */
+    public void dispose() {
+        this.chunk = null;
+        this.sections = null;
+    }
+
+    @Override
+    public BlockState getBlock(int y) {
+        int idx = y - minBuildHeight;
+        if (idx >= 0 && idx < states.length) {
+            return states[idx];
+        }
+        return VOID_AIR;
+    }
+
+    private void markPostprocessing(int y) {
+        cursor.set(
+                SectionPos.sectionToBlockCoord(chunk.getPos().x, localX),
+                y,
+                SectionPos.sectionToBlockCoord(chunk.getPos().z, localZ)
+        );
+        chunk.markPosForPostprocessing(cursor);
+    }
+
+    private void updateHeightmap(Heightmap.Types type, int y, BlockState newState) {
+        chunk.getOrCreateHeightmapUnprimed(type).update(localX, y, localZ, newState);
+    }
+
+    @Override
+    public void setBlock(int y, BlockState state) {
+        int idx = y - minBuildHeight;
+        var states = this.states;
+        if (idx >= 0 && idx < states.length) {
+            BlockState oldState = states[idx];
+            if (oldState == state) {
+                return;
+            }
+            states[idx] = state;
+            // Write directly to the section, bypassing ProtoChunk.setBlockState which
+            // does expensive heightmap updates that are not needed during surface building.
+            int sectionIdx = idx >> 4;
+            sections[sectionIdx].setBlockState(localX, y & 15, localZ, state, false);
+            if (!state.getFluidState().isEmpty()) {
+                markPostprocessing(y);
+            }
+            // Update heightmaps if the air/motion-blocking properties changed.
+            if (oldState.isAir() != state.isAir()) {
+                updateHeightmap(Heightmap.Types.WORLD_SURFACE_WG, y, state);
+            }
+            if (oldState.blocksMotion() != state.blocksMotion()) {
+                updateHeightmap(Heightmap.Types.OCEAN_FLOOR_WG, y, state);
+            }
+        }
+    }
+}