Clipped pregen prefetch

- Limited chunk prefetch rings to the windows a bounded pregeneration still needs
- Forwarded pregeneration region bounds from the async, medieval and modded methods
- Restored unclipped prefetch when the pregeneration closes
- Added a runtime test for clipped and restored prefetch

adapters/modded-common/src/main/java/art/arcane/iris/modded/command/ModdedPregenMethod.java

@@ -348,6 +348,11 @@ public final class ModdedPregenMethod implements PregeneratorMethod {
         return false;
     }
 
+    @Override
+    public void onRegionBounds(int minRegionX, int minRegionZ, int maxRegionX, int maxRegionZ) {
+        hydrology.bound(minRegionX, minRegionZ, maxRegionX, maxRegionZ);
+    }
+
     @Override
     public void onRegionsAhead(List<long[]> regions) {
         if (!closing.get()) {

core/src/main/java/art/arcane/iris/generation/hydrology/engine/HydrologyPregeneration.java

@@ -18,6 +18,12 @@ public final class HydrologyPregeneration implements AutoCloseable {
         }
     }
 
+    public synchronized void bound(int minRegionX, int minRegionZ, int maxRegionX, int maxRegionZ) {
+        if (!closed) {
+            runtime.boundPregeneration(epoch, minRegionX, minRegionZ, maxRegionX, maxRegionZ);
+        }
+    }
+
     @Override
     public synchronized void close() {
         if (!closed) {

core/src/main/java/art/arcane/iris/generation/hydrology/engine/HydrologyRuntime.java

@@ -36,6 +36,7 @@ import java.io.UncheckedIOException;
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.List;
@@ -77,6 +78,7 @@ public final class HydrologyRuntime {
     private final RasterFailures regionPreviewFailures;
     private final Set<Object> previews;
     private final LinkedHashMap<WindowKey, WindowPrefetch> prefetches;
+    private final Map<Object, PregenArea> pregenAreas;
     private final HydrologyBinding binding;
     private final Function<PreparedHydrologyCacheIdentity, HydrologyPlanStore> stores;
     private final boolean caveFluidPolicies;
@@ -99,6 +101,7 @@ public final class HydrologyRuntime {
         this.regionPreviewFailures = null;
         this.previews = null;
         this.prefetches = null;
+        this.pregenAreas = null;
         this.binding = HydrologyBinding.UNBOUND;
         this.stores = identity -> null;
         this.caveFluidPolicies = false;
@@ -125,6 +128,7 @@ public final class HydrologyRuntime {
         this.regionPreviewFailures = new RasterFailures(ticker);
         this.previews = new HashSet<>();
         this.prefetches = new LinkedHashMap<>(PREFETCH_CAPACITY * 2, 0.75f, true);
+        this.pregenAreas = new HashMap<>();
         this.binding = Objects.requireNonNull(binding, "binding");
         this.stores = Objects.requireNonNull(stores, "stores");
         this.caveFluidPolicies = caveFluidPolicies;
@@ -354,8 +358,28 @@ public final class HydrologyRuntime {
         }
     }
 
+    void boundPregeneration(Object epoch, int minRegionX, int minRegionZ, int maxRegionX, int maxRegionZ) {
+        Objects.requireNonNull(epoch, "epoch");
+        if (maxRegionX < minRegionX || maxRegionZ < minRegionZ) {
+            throw new IllegalArgumentException("Invalid pregeneration bounds " + minRegionX + "," + minRegionZ + " to " + maxRegionX + ","
+                    + maxRegionZ);
+        }
+        if (cache == null) {
+            return;
+        }
+        synchronized (prefetches) {
+            pregenAreas.put(epoch, new PregenArea(minRegionX - 1, minRegionZ - 1, maxRegionX + 1, maxRegionZ + 1));
+            prefetches.clear();
+        }
+    }
+
     void cancelPregeneration(Object epoch) {
         if (cache != null) {
+            synchronized (prefetches) {
+                if (pregenAreas.remove(epoch) != null) {
+                    prefetches.clear();
+                }
+            }
             cache.cancel(epoch);
         }
     }
@@ -667,7 +691,7 @@ public final class HydrologyRuntime {
                     && (!known.prepared().isDone() || known.evictions() == evictions)) {
                 return;
             }
-            CompletableFuture<Void> prepared = cache.prepareWindowsAround(chunkX, chunkZ, chunkX, chunkZ, 1 + settings.prefetchRings());
+            CompletableFuture<Void> prepared = prepareAround(chunkX, chunkZ, owner, 1 + settings.prefetchRings());
             prefetches.put(owner, new WindowPrefetch(prepared, evictions));
             if (prefetches.size() > PREFETCH_CAPACITY) {
                 prefetches.pollFirstEntry();
@@ -675,6 +699,24 @@ public final class HydrologyRuntime {
         }
     }
 
+    private CompletableFuture<Void> prepareAround(int chunkX, int chunkZ, WindowKey owner, int rings) {
+        PregenArea covering = null;
+        for (PregenArea area : pregenAreas.values()) {
+            if (area.contains(owner)) {
+                covering = covering == null ? area : covering.union(area);
+            }
+        }
+        if (covering == null) {
+            return cache.prepareWindowsAround(chunkX, chunkZ, chunkX, chunkZ, rings);
+        }
+        int minX = Math.max(owner.wx() - rings, covering.minX());
+        int minZ = Math.max(owner.wz() - rings, covering.minZ());
+        int maxX = Math.min(owner.wx() + rings, covering.maxX());
+        int maxZ = Math.min(owner.wz() + rings, covering.maxZ());
+        return cache.prepareWindowsAround(minX * CHUNKS_PER_WINDOW, minZ * CHUNKS_PER_WINDOW, (maxX + 1) * CHUNKS_PER_WINDOW - 1,
+                (maxZ + 1) * CHUNKS_PER_WINDOW - 1, 0);
+    }
+
     private static List<WindowKey> neighbourhood(int chunkX, int chunkZ) {
         WindowKey own = WindowKey.containing(chunkX << 4, chunkZ << 4);
         List<WindowKey> keys = new ArrayList<>(9);
@@ -754,6 +796,16 @@ public final class HydrologyRuntime {
     private record WindowPrefetch(CompletableFuture<Void> prepared, long evictions) {
     }
 
+    private record PregenArea(int minX, int minZ, int maxX, int maxZ) {
+        private boolean contains(WindowKey key) {
+            return key.wx() > minX && key.wx() < maxX && key.wz() > minZ && key.wz() < maxZ;
+        }
+
+        private PregenArea union(PregenArea other) {
+            return new PregenArea(Math.min(minX, other.minX), Math.min(minZ, other.minZ), Math.max(maxX, other.maxX), Math.max(maxZ, other.maxZ));
+        }
+    }
+
     private static final class LastRaster {
         private long key;
         private ChunkWaterRaster raster;

core/src/main/java/art/arcane/iris/world/pregen/AsyncPregenMethod.java

@@ -268,6 +268,7 @@ public class AsyncPregenMethod implements PregeneratorMethod {
     @Override
     public void onRegionBounds(int minRegionX, int minRegionZ, int maxRegionX, int maxRegionZ) {
         retention = createRetention(minRegionX, minRegionZ, maxRegionX, maxRegionZ);
+        hydrology.bound(minRegionX, minRegionZ, maxRegionX, maxRegionZ);
     }
 
     private PregenChunkRetention createRetention(int minRegionX, int minRegionZ, int maxRegionX, int maxRegionZ) {

core/src/main/java/art/arcane/iris/world/pregen/MedievalPregenMethod.java

@@ -219,6 +219,11 @@ public class MedievalPregenMethod implements PregeneratorMethod {
         return false;
     }
 
+    @Override
+    public void onRegionBounds(int minRegionX, int minRegionZ, int maxRegionX, int maxRegionZ) {
+        hydrology.bound(minRegionX, minRegionZ, maxRegionX, maxRegionZ);
+    }
+
     @Override
     public void onRegionsAhead(List<long[]> regions) {
         if (!closing.get()) {

core/src/main/java/art/arcane/iris/world/pregen/PregenHydrology.java

@@ -13,6 +13,7 @@ public final class PregenHydrology implements AutoCloseable {
     private final Supplier<Engine> engine;
     private HydrologyRuntime owner;
     private HydrologyPregeneration epoch;
+    private int[] bounds;
     private boolean closed;
 
     public PregenHydrology(Supplier<Engine> engine) {
@@ -28,12 +29,21 @@ public final class PregenHydrology implements AutoCloseable {
             withdraw();
             owner = runtime;
             epoch = runtime == null ? null : runtime.openPregeneration();
+            applyBounds();
         }
         if (epoch != null) {
             epoch.prepare(regions);
         }
     }
 
+    public synchronized void bound(int minRegionX, int minRegionZ, int maxRegionX, int maxRegionZ) {
+        if (closed) {
+            return;
+        }
+        bounds = new int[]{minRegionX, minRegionZ, maxRegionX, maxRegionZ};
+        applyBounds();
+    }
+
     @Override
     public synchronized void close() {
         closed = true;
@@ -41,6 +51,12 @@ public final class PregenHydrology implements AutoCloseable {
         owner = null;
     }
 
+    private void applyBounds() {
+        if (epoch != null && bounds != null) {
+            epoch.bound(bounds[0], bounds[1], bounds[2], bounds[3]);
+        }
+    }
+
     private void withdraw() {
         if (epoch != null) {
             epoch.close();

core/src/test/java/art/arcane/iris/generation/hydrology/engine/HydrologyRuntimeTest.java

@@ -67,6 +67,7 @@ import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BooleanSupplier;
 import java.util.stream.Stream;
 
+import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
@@ -414,6 +415,35 @@ public class HydrologyRuntimeTest {
         assertEquals(3, prefetches.get());
     }
 
+    @Test(timeout = 180_000L)
+    public void boundedPregenerationPrefetchesOnlyWindowsItsChunksNeed() {
+        FakePlanners planners = new FakePlanners(settings());
+        HydrologyPlanCache cache = spy(new HydrologyPlanCache(planners.settings, planners::region, planners::window,
+                pool(planners, 2), () -> false));
+        List<int[]> prepared = Collections.synchronizedList(new ArrayList<>());
+        doAnswer(invocation -> {
+            prepared.add(new int[]{invocation.getArgument(0), invocation.getArgument(1), invocation.getArgument(2), invocation.getArgument(3),
+                    invocation.getArgument(4)});
+            return CompletableFuture.completedFuture(null);
+        }).when(cache).prepareWindowsAround(anyInt(), anyInt(), anyInt(), anyInt(), anyInt());
+        HydrologyRuntime runtime = new HydrologyRuntime(planners.settings, cache, planners::natural, () -> false, planners.ticker::get);
+        runtimes.add(runtime);
+        int chunks = WindowKey.SIZE >> 4;
+        HydrologyPregeneration pregeneration = runtime.openPregeneration();
+        pregeneration.bound(0, 0, 1, 1);
+
+        runtime.raster(0, 0);
+        assertArrayEquals(new int[]{-chunks, -chunks, 3 * chunks - 1, 3 * chunks - 1, 0}, prepared.getLast());
+        runtime.raster(chunks + 5, chunks + 5);
+        assertArrayEquals(new int[]{-chunks, -chunks, 3 * chunks - 1, 3 * chunks - 1, 0}, prepared.getLast());
+        runtime.raster(-3 * chunks, 0);
+        assertArrayEquals(new int[]{-3 * chunks, 0, -3 * chunks, 0, 2}, prepared.getLast());
+
+        pregeneration.close();
+        runtime.raster(1, 1);
+        assertArrayEquals(new int[]{1, 1, 1, 1, 2}, prepared.getLast());
+    }
+
     @Test(timeout = 180_000L)
     public void preparePregenerationPlansRegionsInOrder() throws Exception {
         FakePlanners planners = new FakePlanners(settings());