From 2e8c00357239827b05c9afd06f6a000e524ac3d9 Mon Sep 17 00:00:00 2001
From: embeddedt <42941056+embeddedt@users.noreply.github.com>
Date: Tue, 30 Apr 2024 17:11:34 -0400
Subject: [PATCH] Improve speed of searching for structures

---
 .../ChunkGeneratorAccessor.java               | 17 +++++++
 .../StructureCheckMixin.java                  | 45 +++++++++++++++++++
 2 files changed, 62 insertions(+)
 create mode 100644 common/src/main/java/org/embeddedt/modernfix/common/mixin/perf/faster_structure_location/ChunkGeneratorAccessor.java
 create mode 100644 common/src/main/java/org/embeddedt/modernfix/common/mixin/perf/faster_structure_location/StructureCheckMixin.java

diff --git a/common/src/main/java/org/embeddedt/modernfix/common/mixin/perf/faster_structure_location/ChunkGeneratorAccessor.java b/common/src/main/java/org/embeddedt/modernfix/common/mixin/perf/faster_structure_location/ChunkGeneratorAccessor.java
new file mode 100644
index 00000000..32fc1811
--- /dev/null
+++ b/common/src/main/java/org/embeddedt/modernfix/common/mixin/perf/faster_structure_location/ChunkGeneratorAccessor.java
@@ -0,0 +1,17 @@
+package org.embeddedt.modernfix.common.mixin.perf.faster_structure_location;
+
+import net.minecraft.core.Holder;
+import net.minecraft.world.level.chunk.ChunkGenerator;
+import net.minecraft.world.level.levelgen.RandomState;
+import net.minecraft.world.level.levelgen.structure.Structure;
+import net.minecraft.world.level.levelgen.structure.placement.StructurePlacement;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.gen.Invoker;
+
+import java.util.List;
+
+@Mixin(ChunkGenerator.class)
+public interface ChunkGeneratorAccessor {
+    @Invoker("getPlacementsForStructure")
+    List<StructurePlacement> invokeGetPlacementsForStructure(Holder<Structure> structure, RandomState random);
+}
diff --git a/common/src/main/java/org/embeddedt/modernfix/common/mixin/perf/faster_structure_location/StructureCheckMixin.java b/common/src/main/java/org/embeddedt/modernfix/common/mixin/perf/faster_structure_location/StructureCheckMixin.java
new file mode 100644
index 00000000..c99749dd
--- /dev/null
+++ b/common/src/main/java/org/embeddedt/modernfix/common/mixin/perf/faster_structure_location/StructureCheckMixin.java
@@ -0,0 +1,45 @@
+package org.embeddedt.modernfix.common.mixin.perf.faster_structure_location;
+
+import com.llamalad7.mixinextras.injector.ModifyExpressionValue;
+import net.minecraft.core.Registry;
+import net.minecraft.world.level.ChunkPos;
+import net.minecraft.world.level.chunk.ChunkGenerator;
+import net.minecraft.world.level.levelgen.RandomState;
+import net.minecraft.world.level.levelgen.structure.Structure;
+import net.minecraft.world.level.levelgen.structure.StructureCheck;
+import net.minecraft.world.level.levelgen.structure.StructureCheckResult;
+import org.spongepowered.asm.mixin.Final;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Shadow;
+import org.spongepowered.asm.mixin.injection.At;
+
+@Mixin(StructureCheck.class)
+public class StructureCheckMixin {
+    @Shadow @Final private ChunkGenerator chunkGenerator;
+    @Shadow @Final private long seed;
+    @Shadow @Final private Registry<Structure> structureConfigs;
+    @Shadow @Final private RandomState randomState;
+
+    /**
+     * @author embeddedt (inspired by 24w04a and Bytzo's comment on https://bugs.mojang.com/browse/MC-249136)
+     * @reason Avoid running the canCreateStructure method (which can be expensive) if the structure placement already
+     * forbids placing the structure in this chunk.
+     */
+    @ModifyExpressionValue(method = "checkStart", at = @At(value = "INVOKE", target = "Lnet/minecraft/world/level/levelgen/structure/StructureCheck;tryLoadFromStorage(Lnet/minecraft/world/level/ChunkPos;Lnet/minecraft/world/level/levelgen/structure/Structure;ZJ)Lnet/minecraft/world/level/levelgen/structure/StructureCheckResult;"))
+    private StructureCheckResult mfix$checkForValidPosition(StructureCheckResult storageResult, ChunkPos chunkPos, Structure structure, boolean skipKnownStructures) {
+        if (storageResult != null) {
+            return storageResult;
+        } else {
+            // Check if any of the placements allow for this structure to be in this chunk
+            var structureHolder = this.structureConfigs.getHolder(this.structureConfigs.getId(structure)).orElseThrow();
+            for (var placement : ((ChunkGeneratorAccessor)this.chunkGenerator).invokeGetPlacementsForStructure(structureHolder, this.randomState)) {
+                if (placement.isStructureChunk(this.chunkGenerator, this.randomState, this.seed, chunkPos.x, chunkPos.z)) {
+                    // Allowed - return null so regular check runs
+                    return null;
+                }
+            }
+            // Not allowed - early exit by returning a non-null value
+            return StructureCheckResult.START_NOT_PRESENT;
+        }
+    }
+}