implement dimension type injection/caching

platforms/forge/src/main/java/com/dfsek/terra/forge/TerraForgePlugin.java

@@ -52,6 +52,8 @@ import com.dfsek.terra.world.TerraWorld;
 import net.minecraft.util.ResourceLocation;
 import net.minecraft.util.registry.Registry;
 import net.minecraft.util.registry.WorldGenRegistries;
+import net.minecraft.world.DimensionType;
+import net.minecraft.world.IWorld;
 import net.minecraft.world.biome.Biome;
 import net.minecraft.world.gen.feature.ConfiguredFeature;
 import net.minecraft.world.gen.feature.Features;
@@ -59,6 +61,7 @@ import net.minecraft.world.gen.feature.IFeatureConfig;
 import net.minecraft.world.gen.feature.NoFeatureConfig;
 import net.minecraft.world.gen.placement.DecoratedPlacement;
 import net.minecraft.world.gen.placement.NoPlacementConfig;
+import net.minecraft.world.server.ServerWorld;
 import net.minecraftforge.eventbus.api.SubscribeEvent;
 import net.minecraftforge.fml.common.Mod;
 import net.minecraftforge.fml.event.lifecycle.FMLCommonSetupEvent;
@@ -86,7 +89,7 @@ public class TerraForgePlugin implements TerraPlugin {
     public static final ConfiguredFeature<?, ?> POPULATOR_CONFIGURED_FEATURE = POPULATOR_FEATURE.configured(IFeatureConfig.NONE).decorated(DecoratedPlacement.NOPE.configured(NoPlacementConfig.INSTANCE));
 
     private static TerraForgePlugin INSTANCE;
-    private final Map<Long, TerraWorld> worldMap = new HashMap<>();
+    private final Map<DimensionType, Pair<ServerWorld, TerraWorld>> worldMap = new HashMap<>();
     private final EventManager eventManager = new TerraEventManager(this);
     private final GenericLoaders genericLoaders = new GenericLoaders(this);
     private final Profiler profiler = new ProfilerImpl();
@@ -178,10 +181,13 @@ public class TerraForgePlugin implements TerraPlugin {
 
     @Override
     public TerraWorld getWorld(World world) {
-        return worldMap.computeIfAbsent(world.getSeed(), w -> {
-            logger.info("Loading world " + w);
-            return new TerraWorld(world, ((ForgeChunkGeneratorWrapper) world.getGenerator()).getPack(), this);
-        });
+        return getWorld(((IWorld) world).dimensionType());
+    }
+
+    public TerraWorld getWorld(DimensionType type) {
+        TerraWorld world = worldMap.get(type).getRight();
+        if(world == null) throw new IllegalArgumentException("No world exists with dimension type " + type);
+        return world;
     }
 
     /**
@@ -207,12 +213,6 @@ public class TerraForgePlugin implements TerraPlugin {
         return JarUtil.getJarFile();
     }
 
-    public TerraWorld getWorld(long seed) {
-        TerraWorld world = worldMap.get(seed);
-        if(world == null) throw new IllegalArgumentException("No world exists with seed " + seed);
-        return world;
-    }
-
     @Override
     public com.dfsek.terra.api.util.logging.Logger logger() {
         return logger;
@@ -253,14 +253,11 @@ public class TerraForgePlugin implements TerraPlugin {
         config.load(this);
         LangUtil.load(config.getLanguage(), this); // Load language.
         boolean succeed = registry.loadAll(this);
-        Map<Long, TerraWorld> newMap = new HashMap<>();
-        worldMap.forEach((seed, tw) -> {
-            tw.getConfig().getSamplerCache().clear();
-            String packID = tw.getConfig().getTemplate().getID();
-            newMap.put(seed, new TerraWorld(tw.getWorld(), registry.get(packID), this));
+        worldMap.forEach((seed, pair) -> {
+            pair.getRight().getConfig().getSamplerCache().clear();
+            String packID = pair.getRight().getConfig().getTemplate().getID();
+            pair.setRight(new TerraWorld(pair.getRight().getWorld(), registry.get(packID), this));
         });
-        worldMap.clear();
-        worldMap.putAll(newMap);
         return succeed;
     }
 
@@ -316,6 +313,10 @@ public class TerraForgePlugin implements TerraPlugin {
         return manager;
     }
 
+    public Map<DimensionType, Pair<ServerWorld, TerraWorld>> getWorldMap() {
+        return worldMap;
+    }
+
     @Addon("Terra-Forge")
     @Author("Terra")
     @Version("1.0.0")

platforms/forge/src/main/java/com/dfsek/terra/forge/generation/ForgeChunkGeneratorWrapper.java

@@ -23,6 +23,7 @@ import net.minecraft.util.math.ChunkPos;
 import net.minecraft.util.registry.DynamicRegistries;
 import net.minecraft.util.registry.Registry;
 import net.minecraft.world.Blockreader;
+import net.minecraft.world.DimensionType;
 import net.minecraft.world.IBlockReader;
 import net.minecraft.world.IWorld;
 import net.minecraft.world.biome.BiomeManager;
@@ -61,6 +62,8 @@ public class ForgeChunkGeneratorWrapper extends ChunkGenerator implements Genera
         return pack;
     }
 
+    private DimensionType dimensionType;
+
     public ForgeChunkGeneratorWrapper(TerraBiomeSource biomeSource, long seed, ConfigPack configPack) {
         super(biomeSource, new DimensionStructuresSettings(false));
         this.pack = configPack;
@@ -133,7 +136,7 @@ public class ForgeChunkGeneratorWrapper extends ChunkGenerator implements Genera
 
     @Override
     public int getBaseHeight(int x, int z, Heightmap.@NotNull Type p_222529_3_) {
-        TerraWorld world = TerraForgePlugin.getInstance().getWorld(seed);
+        TerraWorld world = TerraForgePlugin.getInstance().getWorld(dimensionType);
         Sampler sampler = world.getConfig().getSamplerCache().getChunk(FastMath.floorDiv(x, 16), FastMath.floorDiv(z, 16));
         int cx = FastMath.floorMod(x, 16);
         int cz = FastMath.floorMod(z, 16);
@@ -171,4 +174,8 @@ public class ForgeChunkGeneratorWrapper extends ChunkGenerator implements Genera
     public TerraChunkGenerator getHandle() {
         return delegate;
     }
+
+    public void setDimensionType(DimensionType dimensionType) {
+        this.dimensionType = dimensionType;
+    }
 }

platforms/forge/src/main/java/com/dfsek/terra/forge/mixin/ServerWorldMixin.java

@@ -0,0 +1,37 @@
+package com.dfsek.terra.forge.mixin;
+
+import com.dfsek.terra.api.util.generic.pair.Pair;
+import com.dfsek.terra.forge.TerraForgePlugin;
+import com.dfsek.terra.forge.generation.ForgeChunkGeneratorWrapper;
+import com.dfsek.terra.world.TerraWorld;
+import net.minecraft.world.DimensionType;
+import net.minecraft.world.World;
+import net.minecraft.world.server.ServerChunkProvider;
+import net.minecraft.world.server.ServerWorld;
+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;
+import org.spongepowered.asm.mixin.injection.Redirect;
+
+@Mixin(ServerWorld.class)
+public abstract class ServerWorldMixin {
+    @Shadow
+    @Final
+    private ServerChunkProvider chunkSource;
+
+    @Shadow
+    protected abstract void initCapabilities();
+
+    @Redirect(method = "<init>", at = @At(value = "INVOKE", target = "Lnet/minecraft/world/server/ServerWorld;initCapabilities()V"))
+    public void injectConstructor(ServerWorld serverWorld) {
+        if(chunkSource.getGenerator() instanceof ForgeChunkGeneratorWrapper) {
+            ForgeChunkGeneratorWrapper chunkGeneratorWrapper = (ForgeChunkGeneratorWrapper) chunkSource.getGenerator();
+            DimensionType dimensionType = ((World) (Object) this).dimensionType();
+            TerraForgePlugin.getInstance().getWorldMap().put(dimensionType, Pair.of((ServerWorld) (Object) this, new TerraWorld((com.dfsek.terra.api.platform.world.World) this, chunkGeneratorWrapper.getPack(), TerraForgePlugin.getInstance())));
+            chunkGeneratorWrapper.setDimensionType(dimensionType);
+            TerraForgePlugin.getInstance().logger().info("Registered world " + this + " to dimension type " + dimensionType);
+        }
+        initCapabilities();
+    }
+}

platforms/forge/src/main/resources/terra.mixins.json

@@ -5,6 +5,7 @@
   "refmap": "terra-refmap.json",
   "mixins": [
     "DimensionGeneratorSettingsMixin",
+    "ServerWorldMixin",
     "access.AbstractSpawnerAccessor",
     "access.BiomeGeneratorTypeScreensAccessor",
     "implementations.BiomeMixin",
@@ -31,7 +32,6 @@
     "implementations.world.ServerWorldMixin",
     "implementations.world.WorldGenRegionMixin"
   ],
-
   "client": [
     "init.MinecraftClientMixin"
   ],