Don't inject world load context for main worlds to prevent incompatibilities

nms/v1_20_R1/src/main/java/com/volmit/iris/core/nms/v1_20_R1/NMSBinding.java

@@ -645,37 +645,44 @@ public class NMSBinding implements INMSBinding {
     }
 
     @Override
+    @SneakyThrows
     public AutoClosing injectLevelStems() {
-        return inject(this::supplier);
+        if (!dataContextLock.tryLock()) throw new IllegalStateException("Failed to inject data context!");
+
+        var server = ((CraftServer) Bukkit.getServer());
+        var field = getField(MinecraftServer.class, WorldLoader.DataLoadContext.class);
+        var nmsServer = server.getServer();
+        var old = nmsServer.worldLoader;
+
+        field.setAccessible(true);
+        field.set(nmsServer, dataLoadContext.aquire(() -> new WorldLoader.DataLoadContext(
+                old.resources(),
+                old.dataConfiguration(),
+                old.datapackWorldgen(),
+                createRegistryAccess(old.datapackDimensions(), false, true, true, true)
+        )));
+
+        return new AutoClosing(() -> {
+            field.set(nmsServer, old);
+            dataContextLock.unlock();
+        });
     }
 
     @Override
     @SneakyThrows
-    public com.volmit.iris.core.nms.container.Pair<Integer, AutoClosing> injectUncached(boolean overworld, boolean nether, boolean end) {
+    public AutoClosing injectUncached(boolean overworld, boolean nether, boolean end) {
         var reg = registry();
         var field = getField(RegistryAccess.ImmutableRegistryAccess.class, Map.class);
         field.setAccessible(true);
 
-        AutoClosing closing = inject(old -> new WorldLoader.DataLoadContext(
-                        old.resources(),
-                        old.dataConfiguration(),
-                        old.datapackWorldgen(),
-                        createRegistryAccess(old.datapackDimensions(), true, overworld, nether, end)
-                )
-        );
-
-        var injected = ((CraftServer) Bukkit.getServer()).getServer().worldLoader.datapackDimensions().registryOrThrow(Registries.LEVEL_STEM);
+        var access = createRegistryAccess(((CraftServer) Bukkit.getServer()).getServer().worldLoader.datapackDimensions(), true, overworld, nether, end);
+        var injected = access.registryOrThrow(Registries.LEVEL_STEM);
         var old = (Map<ResourceKey<? extends Registry<?>>, Registry<?>>) field.get(reg);
         var fake = new HashMap<>(old);
         fake.put(Registries.LEVEL_STEM, injected);
         field.set(reg, fake);
 
-        return new com.volmit.iris.core.nms.container.Pair<>(
-                injected.size(),
-                new AutoClosing(() -> {
-                    closing.close();
-                    field.set(reg, old);
-                }));
+        return new AutoClosing(() -> field.set(reg, old));
     }
 
     @Override
@@ -692,33 +699,6 @@ public class NMSBinding implements INMSBinding {
         ((CraftWorld) world).getHandle().K.customDimensions = null;
     }
 
-    private WorldLoader.DataLoadContext supplier(WorldLoader.DataLoadContext old) {
-        return dataLoadContext.aquire(() -> new WorldLoader.DataLoadContext(
-                old.resources(),
-                old.dataConfiguration(),
-                old.datapackWorldgen(),
-                createRegistryAccess(old.datapackDimensions(), false, true, true, true)
-        ));
-    }
-
-    @SneakyThrows
-    private AutoClosing inject(Function<WorldLoader.DataLoadContext, WorldLoader.DataLoadContext> transformer) {
-        if (!dataContextLock.tryLock()) throw new IllegalStateException("Failed to inject data context!");
-
-        var server = ((CraftServer) Bukkit.getServer());
-        var field = getField(MinecraftServer.class, WorldLoader.DataLoadContext.class);
-        var nmsServer = server.getServer();
-        var old = nmsServer.worldLoader;
-
-        field.setAccessible(true);
-        field.set(nmsServer, transformer.apply(old));
-
-        return new AutoClosing(() -> {
-            field.set(nmsServer, old);
-            dataContextLock.unlock();
-        });
-    }
-
     private RegistryAccess.Frozen createRegistryAccess(RegistryAccess.Frozen datapack, boolean copy, boolean overworld, boolean nether, boolean end) {
         var access = registry();
         var dimensions = access.registryOrThrow(Registries.DIMENSION_TYPE);
@@ -740,7 +720,7 @@ public class NMSBinding implements INMSBinding {
 
         if (copy) copy(fake, access.registryOrThrow(Registries.LEVEL_STEM));
 
-        return new RegistryAccess.Frozen.ImmutableRegistryAccess(List.of(fake.freeze())).freeze();
+        return new RegistryAccess.Frozen.ImmutableRegistryAccess(List.of(fake)).freeze();
     }
 
     private void register(MappedRegistry<LevelStem> target, Registry<DimensionType> dimensions, FlatLevelSource source, ResourceKey<LevelStem> key) {

nms/v1_20_R2/src/main/java/com/volmit/iris/core/nms/v1_20_R2/NMSBinding.java

@@ -646,37 +646,44 @@ public class NMSBinding implements INMSBinding {
     }
 
     @Override
+    @SneakyThrows
     public AutoClosing injectLevelStems() {
-        return inject(this::supplier);
+        if (!dataContextLock.tryLock()) throw new IllegalStateException("Failed to inject data context!");
+
+        var server = ((CraftServer) Bukkit.getServer());
+        var field = getField(MinecraftServer.class, WorldLoader.DataLoadContext.class);
+        var nmsServer = server.getServer();
+        var old = nmsServer.worldLoader;
+
+        field.setAccessible(true);
+        field.set(nmsServer, dataLoadContext.aquire(() -> new WorldLoader.DataLoadContext(
+                old.resources(),
+                old.dataConfiguration(),
+                old.datapackWorldgen(),
+                createRegistryAccess(old.datapackDimensions(), false, true, true, true)
+        )));
+
+        return new AutoClosing(() -> {
+            field.set(nmsServer, old);
+            dataContextLock.unlock();
+        });
     }
 
     @Override
     @SneakyThrows
-    public com.volmit.iris.core.nms.container.Pair<Integer, AutoClosing> injectUncached(boolean overworld, boolean nether, boolean end) {
+    public AutoClosing injectUncached(boolean overworld, boolean nether, boolean end) {
         var reg = registry();
         var field = getField(RegistryAccess.ImmutableRegistryAccess.class, Map.class);
         field.setAccessible(true);
 
-        AutoClosing closing = inject(old -> new WorldLoader.DataLoadContext(
-                        old.resources(),
-                        old.dataConfiguration(),
-                        old.datapackWorldgen(),
-                        createRegistryAccess(old.datapackDimensions(), true, overworld, nether, end)
-                )
-        );
-
-        var injected = ((CraftServer) Bukkit.getServer()).getServer().worldLoader.datapackDimensions().registryOrThrow(Registries.LEVEL_STEM);
+        var access = createRegistryAccess(((CraftServer) Bukkit.getServer()).getServer().worldLoader.datapackDimensions(), true, overworld, nether, end);
+        var injected = access.registryOrThrow(Registries.LEVEL_STEM);
         var old = (Map<ResourceKey<? extends Registry<?>>, Registry<?>>) field.get(reg);
         var fake = new HashMap<>(old);
         fake.put(Registries.LEVEL_STEM, injected);
         field.set(reg, fake);
 
-        return new com.volmit.iris.core.nms.container.Pair<>(
-                injected.size(),
-                new AutoClosing(() -> {
-                    closing.close();
-                    field.set(reg, old);
-                }));
+        return new AutoClosing(() -> field.set(reg, old));
     }
 
     @Override
@@ -693,33 +700,6 @@ public class NMSBinding implements INMSBinding {
         ((CraftWorld) world).getHandle().K.customDimensions = null;
     }
 
-    private WorldLoader.DataLoadContext supplier(WorldLoader.DataLoadContext old) {
-        return dataLoadContext.aquire(() -> new WorldLoader.DataLoadContext(
-                old.resources(),
-                old.dataConfiguration(),
-                old.datapackWorldgen(),
-                createRegistryAccess(old.datapackDimensions(), false, true, true, true)
-        ));
-    }
-
-    @SneakyThrows
-    private AutoClosing inject(Function<WorldLoader.DataLoadContext, WorldLoader.DataLoadContext> transformer) {
-        if (!dataContextLock.tryLock()) throw new IllegalStateException("Failed to inject data context!");
-
-        var server = ((CraftServer) Bukkit.getServer());
-        var field = getField(MinecraftServer.class, WorldLoader.DataLoadContext.class);
-        var nmsServer = server.getServer();
-        var old = nmsServer.worldLoader;
-
-        field.setAccessible(true);
-        field.set(nmsServer, transformer.apply(old));
-
-        return new AutoClosing(() -> {
-            field.set(nmsServer, old);
-            dataContextLock.unlock();
-        });
-    }
-
     private RegistryAccess.Frozen createRegistryAccess(RegistryAccess.Frozen datapack, boolean copy, boolean overworld, boolean nether, boolean end) {
         var access = registry();
         var dimensions = access.registryOrThrow(Registries.DIMENSION_TYPE);
@@ -741,7 +721,7 @@ public class NMSBinding implements INMSBinding {
 
         if (copy) copy(fake, access.registryOrThrow(Registries.LEVEL_STEM));
 
-        return new RegistryAccess.Frozen.ImmutableRegistryAccess(List.of(fake.freeze())).freeze();
+        return new RegistryAccess.Frozen.ImmutableRegistryAccess(List.of(fake)).freeze();
     }
 
     private void register(MappedRegistry<LevelStem> target, Registry<DimensionType> dimensions, FlatLevelSource source, ResourceKey<LevelStem> key) {

nms/v1_20_R3/src/main/java/com/volmit/iris/core/nms/v1_20_R3/NMSBinding.java

@@ -647,37 +647,44 @@ public class NMSBinding implements INMSBinding {
     }
 
     @Override
+    @SneakyThrows
     public AutoClosing injectLevelStems() {
-        return inject(this::supplier);
+        if (!dataContextLock.tryLock()) throw new IllegalStateException("Failed to inject data context!");
+
+        var server = ((CraftServer) Bukkit.getServer());
+        var field = getField(MinecraftServer.class, WorldLoader.DataLoadContext.class);
+        var nmsServer = server.getServer();
+        var old = nmsServer.worldLoader;
+
+        field.setAccessible(true);
+        field.set(nmsServer, dataLoadContext.aquire(() -> new WorldLoader.DataLoadContext(
+                old.resources(),
+                old.dataConfiguration(),
+                old.datapackWorldgen(),
+                createRegistryAccess(old.datapackDimensions(), false, true, true, true)
+        )));
+
+        return new AutoClosing(() -> {
+            field.set(nmsServer, old);
+            dataContextLock.unlock();
+        });
     }
 
     @Override
     @SneakyThrows
-    public com.volmit.iris.core.nms.container.Pair<Integer, AutoClosing> injectUncached(boolean overworld, boolean nether, boolean end) {
+    public AutoClosing injectUncached(boolean overworld, boolean nether, boolean end) {
         var reg = registry();
         var field = getField(RegistryAccess.ImmutableRegistryAccess.class, Map.class);
         field.setAccessible(true);
 
-        AutoClosing closing = inject(old -> new WorldLoader.DataLoadContext(
-                        old.resources(),
-                        old.dataConfiguration(),
-                        old.datapackWorldgen(),
-                        createRegistryAccess(old.datapackDimensions(), true, overworld, nether, end)
-                )
-        );
-
-        var injected = ((CraftServer) Bukkit.getServer()).getServer().worldLoader.datapackDimensions().registryOrThrow(Registries.LEVEL_STEM);
+        var access = createRegistryAccess(((CraftServer) Bukkit.getServer()).getServer().worldLoader.datapackDimensions(), true, overworld, nether, end);
+        var injected = access.registryOrThrow(Registries.LEVEL_STEM);
         var old = (Map<ResourceKey<? extends Registry<?>>, Registry<?>>) field.get(reg);
         var fake = new HashMap<>(old);
         fake.put(Registries.LEVEL_STEM, injected);
         field.set(reg, fake);
 
-        return new com.volmit.iris.core.nms.container.Pair<>(
-                injected.size(),
-                new AutoClosing(() -> {
-                    closing.close();
-                    field.set(reg, old);
-                }));
+        return new AutoClosing(() -> field.set(reg, old));
     }
 
     @Override
@@ -694,33 +701,6 @@ public class NMSBinding implements INMSBinding {
         ((CraftWorld) world).getHandle().K.customDimensions = null;
     }
 
-    private WorldLoader.DataLoadContext supplier(WorldLoader.DataLoadContext old) {
-        return dataLoadContext.aquire(() -> new WorldLoader.DataLoadContext(
-                old.resources(),
-                old.dataConfiguration(),
-                old.datapackWorldgen(),
-                createRegistryAccess(old.datapackDimensions(), false, true, true, true)
-        ));
-    }
-
-    @SneakyThrows
-    private AutoClosing inject(Function<WorldLoader.DataLoadContext, WorldLoader.DataLoadContext> transformer) {
-        if (!dataContextLock.tryLock()) throw new IllegalStateException("Failed to inject data context!");
-
-        var server = ((CraftServer) Bukkit.getServer());
-        var field = getField(MinecraftServer.class, WorldLoader.DataLoadContext.class);
-        var nmsServer = server.getServer();
-        var old = nmsServer.worldLoader;
-
-        field.setAccessible(true);
-        field.set(nmsServer, transformer.apply(old));
-
-        return new AutoClosing(() -> {
-            field.set(nmsServer, old);
-            dataContextLock.unlock();
-        });
-    }
-
     private RegistryAccess.Frozen createRegistryAccess(RegistryAccess.Frozen datapack, boolean copy, boolean overworld, boolean nether, boolean end) {
         var access = registry();
         var dimensions = access.registryOrThrow(Registries.DIMENSION_TYPE);
@@ -742,7 +722,7 @@ public class NMSBinding implements INMSBinding {
 
         if (copy) copy(fake, access.registryOrThrow(Registries.LEVEL_STEM));
 
-        return new RegistryAccess.Frozen.ImmutableRegistryAccess(List.of(fake.freeze())).freeze();
+        return new RegistryAccess.Frozen.ImmutableRegistryAccess(List.of(fake)).freeze();
     }
 
     private void register(MappedRegistry<LevelStem> target, Registry<DimensionType> dimensions, FlatLevelSource source, ResourceKey<LevelStem> key) {

nms/v1_20_R4/src/main/java/com/volmit/iris/core/nms/v1_20_R4/NMSBinding.java

@@ -672,37 +672,44 @@ public class NMSBinding implements INMSBinding {
     }
 
     @Override
+    @SneakyThrows
     public AutoClosing injectLevelStems() {
-        return inject(this::supplier);
+        if (!dataContextLock.tryLock()) throw new IllegalStateException("Failed to inject data context!");
+
+        var server = ((CraftServer) Bukkit.getServer());
+        var field = getField(MinecraftServer.class, WorldLoader.DataLoadContext.class);
+        var nmsServer = server.getServer();
+        var old = nmsServer.worldLoader;
+
+        field.setAccessible(true);
+        field.set(nmsServer, dataLoadContext.aquire(() -> new WorldLoader.DataLoadContext(
+                old.resources(),
+                old.dataConfiguration(),
+                old.datapackWorldgen(),
+                createRegistryAccess(old.datapackDimensions(), false, true, true, true)
+        )));
+
+        return new AutoClosing(() -> {
+            field.set(nmsServer, old);
+            dataContextLock.unlock();
+        });
     }
 
     @Override
     @SneakyThrows
-    public com.volmit.iris.core.nms.container.Pair<Integer, AutoClosing> injectUncached(boolean overworld, boolean nether, boolean end) {
+    public AutoClosing injectUncached(boolean overworld, boolean nether, boolean end) {
         var reg = registry();
         var field = getField(RegistryAccess.ImmutableRegistryAccess.class, Map.class);
         field.setAccessible(true);
 
-        AutoClosing closing = inject(old -> new WorldLoader.DataLoadContext(
-                        old.resources(),
-                        old.dataConfiguration(),
-                        old.datapackWorldgen(),
-                        createRegistryAccess(old.datapackDimensions(), true, overworld, nether, end)
-                )
-        );
-
-        var injected = ((CraftServer) Bukkit.getServer()).getServer().worldLoader.datapackDimensions().registryOrThrow(Registries.LEVEL_STEM);
+        var access = createRegistryAccess(((CraftServer) Bukkit.getServer()).getServer().worldLoader.datapackDimensions(), true, overworld, nether, end);
+        var injected = access.registryOrThrow(Registries.LEVEL_STEM);
         var old = (Map<ResourceKey<? extends Registry<?>>, Registry<?>>) field.get(reg);
         var fake = new HashMap<>(old);
         fake.put(Registries.LEVEL_STEM, injected);
         field.set(reg, fake);
 
-        return new com.volmit.iris.core.nms.container.Pair<>(
-                injected.size(),
-                new AutoClosing(() -> {
-                    closing.close();
-                    field.set(reg, old);
-                }));
+        return new AutoClosing(() -> field.set(reg, old));
     }
 
     @Override
@@ -719,33 +726,6 @@ public class NMSBinding implements INMSBinding {
         ((CraftWorld) world).getHandle().K.customDimensions = null;
     }
 
-    private WorldLoader.DataLoadContext supplier(WorldLoader.DataLoadContext old) {
-        return dataLoadContext.aquire(() -> new WorldLoader.DataLoadContext(
-                old.resources(),
-                old.dataConfiguration(),
-                old.datapackWorldgen(),
-                createRegistryAccess(old.datapackDimensions(), false, true, true, true)
-        ));
-    }
-
-    @SneakyThrows
-    private AutoClosing inject(Function<WorldLoader.DataLoadContext, WorldLoader.DataLoadContext> transformer) {
-        if (!dataContextLock.tryLock()) throw new IllegalStateException("Failed to inject data context!");
-
-        var server = ((CraftServer) Bukkit.getServer());
-        var field = getField(MinecraftServer.class, WorldLoader.DataLoadContext.class);
-        var nmsServer = server.getServer();
-        var old = nmsServer.worldLoader;
-
-        field.setAccessible(true);
-        field.set(nmsServer, transformer.apply(old));
-
-        return new AutoClosing(() -> {
-            field.set(nmsServer, old);
-            dataContextLock.unlock();
-        });
-    }
-
     private RegistryAccess.Frozen createRegistryAccess(RegistryAccess.Frozen datapack, boolean copy, boolean overworld, boolean nether, boolean end) {
         var access = registry();
         var dimensions = access.registryOrThrow(Registries.DIMENSION_TYPE);
@@ -767,7 +747,7 @@ public class NMSBinding implements INMSBinding {
 
         if (copy) copy(fake, access.registryOrThrow(Registries.LEVEL_STEM));
 
-        return new RegistryAccess.Frozen.ImmutableRegistryAccess(List.of(fake.freeze())).freeze();
+        return new RegistryAccess.Frozen.ImmutableRegistryAccess(List.of(fake)).freeze();
     }
 
     private void register(MappedRegistry<LevelStem> target, Registry<DimensionType> dimensions, FlatLevelSource source, ResourceKey<LevelStem> key) {

nms/v1_21_R1/src/main/java/com/volmit/iris/core/nms/v1_21_R1/NMSBinding.java

@@ -676,37 +676,44 @@ public class NMSBinding implements INMSBinding {
     }
 
     @Override
+    @SneakyThrows
     public AutoClosing injectLevelStems() {
-        return inject(this::supplier);
+        if (!dataContextLock.tryLock()) throw new IllegalStateException("Failed to inject data context!");
+
+        var server = ((CraftServer) Bukkit.getServer());
+        var field = getField(MinecraftServer.class, WorldLoader.DataLoadContext.class);
+        var nmsServer = server.getServer();
+        var old = nmsServer.worldLoader;
+
+        field.setAccessible(true);
+        field.set(nmsServer, dataLoadContext.aquire(() -> new WorldLoader.DataLoadContext(
+                old.resources(),
+                old.dataConfiguration(),
+                old.datapackWorldgen(),
+                createRegistryAccess(old.datapackDimensions(), false, true, true, true)
+        )));
+
+        return new AutoClosing(() -> {
+            field.set(nmsServer, old);
+            dataContextLock.unlock();
+        });
     }
 
     @Override
     @SneakyThrows
-    public com.volmit.iris.core.nms.container.Pair<Integer, AutoClosing> injectUncached(boolean overworld, boolean nether, boolean end) {
+    public AutoClosing injectUncached(boolean overworld, boolean nether, boolean end) {
         var reg = registry();
         var field = getField(RegistryAccess.ImmutableRegistryAccess.class, Map.class);
         field.setAccessible(true);
 
-        AutoClosing closing = inject(old -> new WorldLoader.DataLoadContext(
-                        old.resources(),
-                        old.dataConfiguration(),
-                        old.datapackWorldgen(),
-                        createRegistryAccess(old.datapackDimensions(), true, overworld, nether, end)
-                )
-        );
-
-        var injected = ((CraftServer) Bukkit.getServer()).getServer().worldLoader.datapackDimensions().registryOrThrow(Registries.LEVEL_STEM);
+        var access = createRegistryAccess(((CraftServer) Bukkit.getServer()).getServer().worldLoader.datapackDimensions(), true, overworld, nether, end);
+        var injected = access.registryOrThrow(Registries.LEVEL_STEM);
         var old = (Map<ResourceKey<? extends Registry<?>>, Registry<?>>) field.get(reg);
         var fake = new HashMap<>(old);
         fake.put(Registries.LEVEL_STEM, injected);
         field.set(reg, fake);
 
-        return new com.volmit.iris.core.nms.container.Pair<>(
-                injected.size(),
-                new AutoClosing(() -> {
-                    closing.close();
-                    field.set(reg, old);
-                }));
+        return new AutoClosing(() -> field.set(reg, old));
     }
 
     @Override
@@ -723,33 +730,6 @@ public class NMSBinding implements INMSBinding {
         ((CraftWorld) world).getHandle().K.customDimensions = null;
     }
 
-    private WorldLoader.DataLoadContext supplier(WorldLoader.DataLoadContext old) {
-        return dataLoadContext.aquire(() -> new WorldLoader.DataLoadContext(
-                old.resources(),
-                old.dataConfiguration(),
-                old.datapackWorldgen(),
-                createRegistryAccess(old.datapackDimensions(), false, true, true, true)
-        ));
-    }
-
-    @SneakyThrows
-    private AutoClosing inject(Function<WorldLoader.DataLoadContext, WorldLoader.DataLoadContext> transformer) {
-        if (!dataContextLock.tryLock()) throw new IllegalStateException("Failed to inject data context!");
-
-        var server = ((CraftServer) Bukkit.getServer());
-        var field = getField(MinecraftServer.class, WorldLoader.DataLoadContext.class);
-        var nmsServer = server.getServer();
-        var old = nmsServer.worldLoader;
-
-        field.setAccessible(true);
-        field.set(nmsServer, transformer.apply(old));
-
-        return new AutoClosing(() -> {
-            field.set(nmsServer, old);
-            dataContextLock.unlock();
-        });
-    }
-
     private RegistryAccess.Frozen createRegistryAccess(RegistryAccess.Frozen datapack, boolean copy, boolean overworld, boolean nether, boolean end) {
         var access = registry();
         var dimensions = access.registryOrThrow(Registries.DIMENSION_TYPE);
@@ -771,7 +751,7 @@ public class NMSBinding implements INMSBinding {
 
         if (copy) copy(fake, access.registryOrThrow(Registries.LEVEL_STEM));
 
-        return new RegistryAccess.Frozen.ImmutableRegistryAccess(List.of(fake.freeze())).freeze();
+        return new RegistryAccess.Frozen.ImmutableRegistryAccess(List.of(fake)).freeze();
     }
 
     private void register(MappedRegistry<LevelStem> target, Registry<DimensionType> dimensions, FlatLevelSource source, ResourceKey<LevelStem> key) {

nms/v1_21_R2/src/main/java/com/volmit/iris/core/nms/v1_21_R2/NMSBinding.java

@@ -666,37 +666,44 @@ public class NMSBinding implements INMSBinding {
     }
 
     @Override
+    @SneakyThrows
     public AutoClosing injectLevelStems() {
-        return inject(this::supplier);
+        if (!dataContextLock.tryLock()) throw new IllegalStateException("Failed to inject data context!");
+
+        var server = ((CraftServer) Bukkit.getServer());
+        var field = getField(MinecraftServer.class, WorldLoader.DataLoadContext.class);
+        var nmsServer = server.getServer();
+        var old = nmsServer.worldLoader;
+
+        field.setAccessible(true);
+        field.set(nmsServer, dataLoadContext.aquire(() -> new WorldLoader.DataLoadContext(
+                old.resources(),
+                old.dataConfiguration(),
+                old.datapackWorldgen(),
+                createRegistryAccess(old.datapackDimensions(), false, true, true, true)
+        )));
+
+        return new AutoClosing(() -> {
+            field.set(nmsServer, old);
+            dataContextLock.unlock();
+        });
     }
 
     @Override
     @SneakyThrows
-    public Pair<Integer, AutoClosing> injectUncached(boolean overworld, boolean nether, boolean end) {
+    public AutoClosing injectUncached(boolean overworld, boolean nether, boolean end) {
         var reg = registry();
         var field = getField(RegistryAccess.ImmutableRegistryAccess.class, Map.class);
         field.setAccessible(true);
 
-        AutoClosing closing = inject(old -> new WorldLoader.DataLoadContext(
-                        old.resources(),
-                        old.dataConfiguration(),
-                        old.datapackWorldgen(),
-                        createRegistryAccess(old.datapackDimensions(), true, overworld, nether, end)
-                )
-        );
-
-        var injected = ((CraftServer) Bukkit.getServer()).getServer().worldLoader.datapackDimensions().lookupOrThrow(Registries.LEVEL_STEM);
+        var access = createRegistryAccess(((CraftServer) Bukkit.getServer()).getServer().worldLoader.datapackDimensions(), true, overworld, nether, end);
+        var injected = access.lookupOrThrow(Registries.LEVEL_STEM);
         var old = (Map<ResourceKey<? extends Registry<?>>, Registry<?>>) field.get(reg);
         var fake = new HashMap<>(old);
         fake.put(Registries.LEVEL_STEM, injected);
         field.set(reg, fake);
 
-        return new Pair<>(
-                injected.size(),
-                new AutoClosing(() -> {
-                    closing.close();
-                    field.set(reg, old);
-                }));
+        return new AutoClosing(() -> field.set(reg, old));
     }
 
     @Override
@@ -713,33 +720,6 @@ public class NMSBinding implements INMSBinding {
         ((CraftWorld) world).getHandle().L.customDimensions = null;
     }
 
-    private WorldLoader.DataLoadContext supplier(WorldLoader.DataLoadContext old) {
-        return dataLoadContext.aquire(() -> new WorldLoader.DataLoadContext(
-                old.resources(),
-                old.dataConfiguration(),
-                old.datapackWorldgen(),
-                createRegistryAccess(old.datapackDimensions(), false, true, true, true)
-        ));
-    }
-
-    @SneakyThrows
-    private AutoClosing inject(Function<WorldLoader.DataLoadContext, WorldLoader.DataLoadContext> transformer) {
-        if (!dataContextLock.tryLock()) throw new IllegalStateException("Failed to inject data context!");
-
-        var server = ((CraftServer) Bukkit.getServer());
-        var field = getField(MinecraftServer.class, WorldLoader.DataLoadContext.class);
-        var nmsServer = server.getServer();
-        var old = nmsServer.worldLoader;
-
-        field.setAccessible(true);
-        field.set(nmsServer, transformer.apply(old));
-
-        return new AutoClosing(() -> {
-            field.set(nmsServer, old);
-            dataContextLock.unlock();
-        });
-    }
-
     private RegistryAccess.Frozen createRegistryAccess(RegistryAccess.Frozen datapack, boolean copy, boolean overworld, boolean nether, boolean end) {
         var access = registry();
         var dimensions = access.lookupOrThrow(Registries.DIMENSION_TYPE);
@@ -761,7 +741,7 @@ public class NMSBinding implements INMSBinding {
 
         if (copy) copy(fake, access.lookupOrThrow(Registries.LEVEL_STEM));
 
-        return new RegistryAccess.Frozen.ImmutableRegistryAccess(List.of(fake.freeze())).freeze();
+        return new RegistryAccess.Frozen.ImmutableRegistryAccess(List.of(fake)).freeze();
     }
 
     private void register(MappedRegistry<LevelStem> target, Registry<DimensionType> dimensions, FlatLevelSource source, ResourceKey<LevelStem> key) {

nms/v1_21_R3/src/main/java/com/volmit/iris/core/nms/v1_21_R3/NMSBinding.java

@@ -665,37 +665,44 @@ public class NMSBinding implements INMSBinding {
     }
 
     @Override
+    @SneakyThrows
     public AutoClosing injectLevelStems() {
-        return inject(this::supplier);
+        if (!dataContextLock.tryLock()) throw new IllegalStateException("Failed to inject data context!");
+
+        var server = ((CraftServer) Bukkit.getServer());
+        var field = getField(MinecraftServer.class, WorldLoader.DataLoadContext.class);
+        var nmsServer = server.getServer();
+        var old = nmsServer.worldLoader;
+
+        field.setAccessible(true);
+        field.set(nmsServer, dataLoadContext.aquire(() -> new WorldLoader.DataLoadContext(
+                old.resources(),
+                old.dataConfiguration(),
+                old.datapackWorldgen(),
+                createRegistryAccess(old.datapackDimensions(), false, true, true, true)
+        )));
+
+        return new AutoClosing(() -> {
+            field.set(nmsServer, old);
+            dataContextLock.unlock();
+        });
     }
 
     @Override
     @SneakyThrows
-    public Pair<Integer, AutoClosing> injectUncached(boolean overworld, boolean nether, boolean end) {
+    public AutoClosing injectUncached(boolean overworld, boolean nether, boolean end) {
         var reg = registry();
         var field = getField(RegistryAccess.ImmutableRegistryAccess.class, Map.class);
         field.setAccessible(true);
 
-        AutoClosing closing = inject(old -> new WorldLoader.DataLoadContext(
-                    old.resources(),
-                    old.dataConfiguration(),
-                    old.datapackWorldgen(),
-                    createRegistryAccess(old.datapackDimensions(), true, overworld, nether, end)
-                )
-        );
-
-        var injected = ((CraftServer) Bukkit.getServer()).getServer().worldLoader.datapackDimensions().lookupOrThrow(Registries.LEVEL_STEM);
+        var access = createRegistryAccess(((CraftServer) Bukkit.getServer()).getServer().worldLoader.datapackDimensions(), true, overworld, nether, end);
+        var injected = access.lookupOrThrow(Registries.LEVEL_STEM);
         var old = (Map<ResourceKey<? extends Registry<?>>, Registry<?>>) field.get(reg);
         var fake = new HashMap<>(old);
         fake.put(Registries.LEVEL_STEM, injected);
         field.set(reg, fake);
 
-        return new Pair<>(
-                injected.size(),
-                new AutoClosing(() -> {
-                    closing.close();
-                    field.set(reg, old);
-                }));
+        return new AutoClosing(() -> field.set(reg, old));
     }
 
     @Override
@@ -712,33 +719,6 @@ public class NMSBinding implements INMSBinding {
         ((CraftWorld) world).getHandle().L.customDimensions = null;
     }
 
-    private WorldLoader.DataLoadContext supplier(WorldLoader.DataLoadContext old) {
-        return dataLoadContext.aquire(() -> new WorldLoader.DataLoadContext(
-                old.resources(),
-                old.dataConfiguration(),
-                old.datapackWorldgen(),
-                createRegistryAccess(old.datapackDimensions(), false, true, true, true)
-        ));
-    }
-
-    @SneakyThrows
-    private AutoClosing inject(Function<WorldLoader.DataLoadContext, WorldLoader.DataLoadContext> transformer) {
-        if (!dataContextLock.tryLock()) throw new IllegalStateException("Failed to inject data context!");
-
-        var server = ((CraftServer) Bukkit.getServer());
-        var field = getField(MinecraftServer.class, WorldLoader.DataLoadContext.class);
-        var nmsServer = server.getServer();
-        var old = nmsServer.worldLoader;
-
-        field.setAccessible(true);
-        field.set(nmsServer, transformer.apply(old));
-
-        return new AutoClosing(() -> {
-            field.set(nmsServer, old);
-            dataContextLock.unlock();
-        });
-    }
-
     private RegistryAccess.Frozen createRegistryAccess(RegistryAccess.Frozen datapack, boolean copy, boolean overworld, boolean nether, boolean end) {
         var access = registry();
         var dimensions = access.lookupOrThrow(Registries.DIMENSION_TYPE);
@@ -760,7 +740,7 @@ public class NMSBinding implements INMSBinding {
 
         if (copy) copy(fake, access.lookupOrThrow(Registries.LEVEL_STEM));
 
-        return new RegistryAccess.Frozen.ImmutableRegistryAccess(List.of(fake.freeze())).freeze();
+        return new RegistryAccess.Frozen.ImmutableRegistryAccess(List.of(fake)).freeze();
     }
 
     private void register(MappedRegistry<LevelStem> target, Registry<DimensionType> dimensions, FlatLevelSource source, ResourceKey<LevelStem> key) {