add bytebuddy binding for 1.21.5

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

@@ -1,16 +1,17 @@
 package com.volmit.iris.core.nms.v1_21_R4;
 
 import com.mojang.brigadier.exceptions.CommandSyntaxException;
-import com.mojang.serialization.Lifecycle;
 import com.volmit.iris.Iris;
 import com.volmit.iris.core.nms.INMSBinding;
-import com.volmit.iris.core.nms.container.AutoClosing;
 import com.volmit.iris.core.nms.container.BiomeColor;
 import com.volmit.iris.core.nms.datapack.DataVersion;
 import com.volmit.iris.engine.data.cache.AtomicCache;
 import com.volmit.iris.engine.framework.Engine;
+import com.volmit.iris.engine.platform.PlatformChunkGenerator;
+import com.volmit.iris.util.agent.Agent;
 import com.volmit.iris.util.collection.KList;
 import com.volmit.iris.util.collection.KMap;
+import com.volmit.iris.util.format.C;
 import com.volmit.iris.util.hunk.Hunk;
 import com.volmit.iris.util.json.JSONObject;
 import com.volmit.iris.util.mantle.Mantle;
@@ -21,20 +22,24 @@ import com.volmit.iris.util.nbt.mca.palette.*;
 import com.volmit.iris.util.nbt.tag.CompoundTag;
 import com.volmit.iris.util.scheduling.J;
 import it.unimi.dsi.fastutil.objects.Object2IntMap;
-import lombok.SneakyThrows;
+import it.unimi.dsi.fastutil.shorts.ShortList;
-import net.minecraft.core.Registry;
+import net.bytebuddy.ByteBuddy;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.matcher.ElementMatchers;
 import net.minecraft.core.*;
+import net.minecraft.core.Registry;
 import net.minecraft.core.component.DataComponents;
 import net.minecraft.core.registries.Registries;
-import net.minecraft.nbt.Tag;
 import net.minecraft.nbt.*;
+import net.minecraft.nbt.Tag;
 import net.minecraft.resources.ResourceKey;
 import net.minecraft.resources.ResourceLocation;
 import net.minecraft.server.MinecraftServer;
-import net.minecraft.server.WorldLoader;
 import net.minecraft.server.commands.data.BlockDataAccessor;
 import net.minecraft.server.level.ServerLevel;
+import net.minecraft.server.level.progress.ChunkProgressListener;
 import net.minecraft.tags.TagKey;
+import net.minecraft.world.RandomSequences;
 import net.minecraft.world.entity.EntityType;
 import net.minecraft.world.item.component.CustomData;
 import net.minecraft.world.level.LevelReader;
@@ -46,13 +51,16 @@ import net.minecraft.world.level.block.entity.BlockEntity;
 import net.minecraft.world.level.block.state.BlockState;
 import net.minecraft.world.level.chunk.ChunkAccess;
 import net.minecraft.world.level.chunk.LevelChunk;
+import net.minecraft.world.level.chunk.ProtoChunk;
 import net.minecraft.world.level.chunk.status.ChunkStatus;
 import net.minecraft.world.level.chunk.status.WorldGenContext;
-import net.minecraft.world.level.dimension.DimensionType;
+import net.minecraft.world.level.chunk.storage.SerializableChunkData;
 import net.minecraft.world.level.dimension.LevelStem;
 import net.minecraft.world.level.levelgen.FlatLevelSource;
 import net.minecraft.world.level.levelgen.flat.FlatLayerInfo;
 import net.minecraft.world.level.levelgen.flat.FlatLevelGeneratorSettings;
+import net.minecraft.world.level.storage.LevelStorageSource;
+import net.minecraft.world.level.storage.PrimaryLevelData;
 import org.bukkit.*;
 import org.bukkit.block.Biome;
 import org.bukkit.block.data.BlockData;
@@ -66,6 +74,7 @@ import org.bukkit.craftbukkit.v1_21_R4.inventory.CraftItemStack;
 import org.bukkit.craftbukkit.v1_21_R4.util.CraftNamespacedKey;
 import org.bukkit.entity.Entity;
 import org.bukkit.event.entity.CreatureSpawnEvent;
+import org.bukkit.generator.BiomeProvider;
 import org.bukkit.generator.ChunkGenerator;
 import org.bukkit.inventory.ItemStack;
 import org.jetbrains.annotations.Contract;
@@ -73,22 +82,23 @@ import org.jetbrains.annotations.NotNull;
 
 import java.awt.Color;
 import java.lang.reflect.Field;
+import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
-import java.util.List;
 import java.util.*;
+import java.util.List;
+import java.util.concurrent.Executor;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
-import java.util.concurrent.locks.ReentrantLock;
 
 public class NMSBinding implements INMSBinding {
     private final KMap<Biome, Object> baseBiomeCache = new KMap<>();
     private final BlockData AIR = Material.AIR.createBlockData();
     private final AtomicCache<MCAIdMap<net.minecraft.world.level.biome.Biome>> biomeMapCache = new AtomicCache<>();
-    private final AtomicCache<WorldLoader.DataLoadContext> dataLoadContext = new AtomicCache<>();
+    private final AtomicBoolean injected = new AtomicBoolean();
     private final AtomicCache<MCAIdMapper<BlockState>> registryCache = new AtomicCache<>();
     private final AtomicCache<MCAPalette<BlockState>> globalCache = new AtomicCache<>();
     private final AtomicCache<RegistryAccess> registryAccess = new AtomicCache<>();
-    private final ReentrantLock dataContextLock = new ReentrantLock(true);
     private final AtomicCache<Method> byIdRef = new AtomicCache<>();
     private Field biomeStorageCache = null;
 
@@ -663,103 +673,78 @@ public class NMSBinding implements INMSBinding {
     }
 
     @Override
-    @SneakyThrows
+    public boolean missingDimensionTypes(String... keys) {
-    public AutoClosing injectLevelStems() {
+        var type = registry().lookupOrThrow(Registries.DIMENSION_TYPE);
-        if (!dataContextLock.tryLock()) throw new IllegalStateException("Failed to inject data context!");
+        return !Arrays.stream(keys)
-
+                .map(key -> ResourceLocation.fromNamespaceAndPath("iris", key))
-        var server = ((CraftServer) Bukkit.getServer());
+                .allMatch(type::containsKey);
-        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 boolean injectBukkit() {
-    public AutoClosing injectUncached(boolean overworld, boolean nether, boolean end) {
+        if (injected.getAndSet(true))
-        var reg = registry();
+            return true;
-        var field = getField(RegistryAccess.ImmutableRegistryAccess.class, Map.class);
+        try {
-        field.setAccessible(true);
+            Iris.info("Injecting Bukkit");
+            var buddy = new ByteBuddy();
+            buddy.redefine(ServerLevel.class)
+                    .visit(Advice.to(ServerLevelAdvice.class).on(ElementMatchers.isConstructor().and(ElementMatchers.takesArguments(
+                            MinecraftServer.class, Executor.class, LevelStorageSource.LevelStorageAccess.class, PrimaryLevelData.class,
+                            ResourceKey.class, LevelStem.class, ChunkProgressListener.class, boolean.class, long.class, List.class,
+                            boolean.class, RandomSequences.class, World.Environment.class, ChunkGenerator.class, BiomeProvider.class))))
+                    .make()
+                    .load(ServerLevel.class.getClassLoader(), Agent.installed());
 
-        var access = createRegistryAccess(((CraftServer) Bukkit.getServer()).getServer().worldLoader.datapackDimensions(), true, overworld, nether, end);
+            return true;
-        var injected = access.lookupOrThrow(Registries.LEVEL_STEM);
+        } catch (Throwable e) {
-        var old = (Map<ResourceKey<? extends Registry<?>>, Registry<?>>) field.get(reg);
+            Iris.error(C.RED + "Failed to inject Bukkit");
-        var fake = new HashMap<>(old);
+            e.printStackTrace();
-        fake.put(Registries.LEVEL_STEM, injected);
+        }
-        field.set(reg, fake);
+        return false;
-
-        return new AutoClosing(() -> field.set(reg, old));
     }
 
-    @Override
+    public LevelStem levelStem(RegistryAccess access, ChunkGenerator raw) {
-    public boolean missingDimensionTypes(boolean overworld, boolean nether, boolean end) {
+        if (!(raw instanceof PlatformChunkGenerator gen))
-        var registry = registry().lookupOrThrow(Registries.DIMENSION_TYPE);
+            throw new IllegalStateException("Generator is not platform chunk generator!");
-        if (overworld) overworld = !registry.containsKey(createIrisKey(LevelStem.OVERWORLD));
+
-        if (nether) nether = !registry.containsKey(createIrisKey(LevelStem.NETHER));
+        var dimensionKey = ResourceLocation.fromNamespaceAndPath("iris", gen.getTarget().getDimension().getDimensionTypeKey());
-        if (end) end = !registry.containsKey(createIrisKey(LevelStem.END));
+        var dimensionType = access.lookupOrThrow(Registries.DIMENSION_TYPE).getOrThrow(ResourceKey.create(Registries.DIMENSION_TYPE, dimensionKey));
-        return overworld || nether || end;
+        return new LevelStem(dimensionType, chunkGenerator(access));
     }
 
-    @Override
+    private net.minecraft.world.level.chunk.ChunkGenerator chunkGenerator(RegistryAccess access) {
-    public void removeCustomDimensions(World world) {
+        var settings = new FlatLevelGeneratorSettings(Optional.empty(), access.lookupOrThrow(Registries.BIOME).getOrThrow(Biomes.THE_VOID), List.of());
-        ((CraftWorld) world).getHandle().L.customDimensions = null;
-    }
-
-    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);
-
-        var settings = new FlatLevelGeneratorSettings(
-                Optional.empty(),
-                access.lookupOrThrow(Registries.BIOME).getOrThrow(Biomes.THE_VOID),
-                List.of()
-        );
         settings.getLayersInfo().add(new FlatLayerInfo(1, Blocks.AIR));
         settings.updateLayers();
-
+        return new FlatLevelSource(settings);
-        var source = new FlatLevelSource(settings);
-        var fake = new MappedRegistry<>(Registries.LEVEL_STEM, Lifecycle.experimental());
-        if (overworld) register(fake, dimensions, source, LevelStem.OVERWORLD);
-        if (nether) register(fake, dimensions, source, LevelStem.NETHER);
-        if (end) register(fake, dimensions, source, LevelStem.END);
-        copy(fake, datapack.lookup(Registries.LEVEL_STEM).orElse(null));
-
-        if (copy) copy(fake, access.lookupOrThrow(Registries.LEVEL_STEM));
-
-        return new RegistryAccess.Frozen.ImmutableRegistryAccess(List.of(fake)).freeze();
     }
 
-    private void register(MappedRegistry<LevelStem> target, Registry<DimensionType> dimensions, FlatLevelSource source, ResourceKey<LevelStem> key) {
+    private static class ServerLevelAdvice {
-        var loc = createIrisKey(key);
+        @Advice.OnMethodEnter
-        target.register(key, new LevelStem(
+        static void enter(
-                dimensions.get(loc).orElseThrow(() -> new IllegalStateException("Missing dimension type " + loc + " in " + dimensions.keySet())),
+                @Advice.Argument(0) MinecraftServer server,
-                source
+                @Advice.Argument(3) PrimaryLevelData levelData,
-        ), RegistrationInfo.BUILT_IN);
+                @Advice.Argument(value = 5, readOnly = false) LevelStem levelStem,
-    }
+                @Advice.Argument(12) World.Environment env,
+                @Advice.Argument(value = 13) ChunkGenerator gen
+        ) {
+            if (gen == null || !gen.getClass().getPackageName().startsWith("com.volmit.iris"))
+                return;
 
-    private void copy(MappedRegistry<LevelStem> target, Registry<LevelStem> source) {
+            try {
-        if (source == null) return;
+                Object bindings = Class.forName("com.volmit.iris.core.nms.INMS", true, Bukkit.getPluginManager().getPlugin("Iris")
-        source.listElementIds().forEach(key -> {
+                                .getClass()
-            var value = source.getValue(key);
+                                .getClassLoader())
-            var info = source.registrationInfo(key).orElse(null);
+                        .getDeclaredMethod("get")
-            if (value != null && info != null && !target.containsKey(key))
+                        .invoke(null);
-                target.register(key, value, info);
+                levelStem = (LevelStem) bindings.getClass()
-        });
+                        .getDeclaredMethod("levelStem", RegistryAccess.class, ChunkGenerator.class)
-    }
+                        .invoke(bindings, server.registryAccess(), gen);
 
-    private ResourceLocation createIrisKey(ResourceKey<LevelStem> key) {
+                levelData.customDimensions = null;
-        return ResourceLocation.fromNamespaceAndPath("iris", key.location().getPath());
+            } catch (Throwable e) {
+                throw new RuntimeException("Iris failed to replace the levelStem", e instanceof InvocationTargetException ex ? ex.getCause() : e);
+            }
+        }
     }
 }