Fancy wrapping logic to not make level files incompatible

lukebemishprojects · Jun 13, 2024 · 600a955 · 600a955
1 parent 17fd513
commit 600a955
Show file tree

Hide file tree

Showing 8 changed files with 161 additions and 20 deletions.
diff --git a/src/main/java/dev/lukebemish/biomesquisher/impl/BiomeSquisher.java b/src/main/java/dev/lukebemish/biomesquisher/impl/BiomeSquisher.java
@@ -1,6 +1,7 @@
 package dev.lukebemish.biomesquisher.impl;
 
 import dev.lukebemish.biomesquisher.BiomeSquisherRegistries;
+import dev.lukebemish.biomesquisher.impl.injected.KnowsOriginalKey;
 import dev.lukebemish.biomesquisher.impl.injected.Squishable;
 import dev.lukebemish.biomesquisher.impl.mixin.MultiNoiseBiomeSourceAccessor;
 import dev.lukebemish.biomesquisher.impl.mixin.NoiseBasedChunkGeneratorAccessor;
@@ -59,22 +60,25 @@ public static void squishBiomeSource(ResourceManager resourceManager, @Nullable
         }
     }
 
-    public static void setupSurfaceRuleModification(NoiseGeneratorSettings generator, ResourceKey<NoiseGeneratorSettings> key) {
+    public static void setupOriginalKeyAwareGenerators(NoiseGeneratorSettings generator, ResourceKey<NoiseGeneratorSettings> key) {
+        //noinspection DataFlowIssue
+        ((KnowsOriginalKey) (Object) generator).biomesquisher_generatorKey(key);
+    }
+
+    public static void modifySurfaceRules(NoiseGeneratorSettings generator, RegistryAccess access, ResourceKey<NoiseGeneratorSettings> backupKey) {
         var surfaceRulesSource = generator.surfaceRule();
-        WrappingRuleSource.NotifyingOps.NotifyingJsonOps ops = WrappingRuleSource.NotifyingOps.NotifyingJsonOps.create(wrapped -> wrapped.generator(key));
+        //noinspection DataFlowIssue
+        var key = ((KnowsOriginalKey) (Object) generator).biomesquisher_generatorKey();
+        WrappingRuleSource.NotifyingOps.NotifyingJsonOps ops = WrappingRuleSource.NotifyingOps.NotifyingJsonOps.create(wrapped -> wrapped.modifiers(loadRuleModifiers(key == null ? wrapped.generator() : key, access)));
         SurfaceRules.RuleSource.CODEC.encodeStart(ops, surfaceRulesSource);
         if (!ops.isWrapped()) {
-            var wrappedSource = WrappingRuleSource.create(surfaceRulesSource, key);
+            var realKey = key == null ? backupKey : key;
+            var wrappedSource = WrappingRuleSource.create(surfaceRulesSource, realKey);
+            wrappedSource.modifiers(loadRuleModifiers(realKey, access));
             setSurfaceRule(generator, wrappedSource);
         }
     }
 
-    public static void modifySurfaceRules(NoiseGeneratorSettings generator, RegistryAccess access) {
-        var surfaceRulesSource = generator.surfaceRule();
-        WrappingRuleSource.NotifyingOps.NotifyingJsonOps ops = WrappingRuleSource.NotifyingOps.NotifyingJsonOps.create(wrapped -> wrapped.modifiers(loadRuleModifiers(wrapped.generator(), access)));
-        SurfaceRules.RuleSource.CODEC.encodeStart(ops, surfaceRulesSource);
-    }
-
     @SuppressWarnings("unused")
     @Open(
         targetClass = NoiseGeneratorSettings.class,

diff --git a/src/main/java/dev/lukebemish/biomesquisher/impl/WrappingRuleSource.java b/src/main/java/dev/lukebemish/biomesquisher/impl/WrappingRuleSource.java
@@ -1,12 +1,15 @@
 package dev.lukebemish.biomesquisher.impl;
 
+import com.mojang.datafixers.util.Pair;
+import com.mojang.serialization.Codec;
 import com.mojang.serialization.DataResult;
 import com.mojang.serialization.DynamicOps;
 import com.mojang.serialization.JsonOps;
 import com.mojang.serialization.MapCodec;
 import com.mojang.serialization.MapLike;
 import com.mojang.serialization.RecordBuilder;
 import com.mojang.serialization.codecs.RecordCodecBuilder;
+import dev.lukebemish.biomesquisher.impl.injected.KnowsOriginalKey;
 import dev.lukebemish.biomesquisher.surface.RuleModifier;
 import dev.lukebemish.biomesquisher.surface.SurfaceRuleInjection;
 import dev.lukebemish.opensesame.annotations.Coerce;
@@ -62,6 +65,51 @@ public <T> RecordBuilder<T> encode(WrappingRuleSource input, DynamicOps<T> ops,
         }
     });
 
+    static Codec<NoiseGeneratorSettings> wrap(Codec<NoiseGeneratorSettings> original) {
+        return new Codec<>() {
+            private static final String BIOMESQUISHER_GENERATOR_KEY = "biome_squisher_generator_key";
+            private static final Codec<ResourceKey<NoiseGeneratorSettings>> RESOURCE_KEY_CODEC = ResourceKey.codec(Registries.NOISE_SETTINGS);
+
+            @SuppressWarnings("DataFlowIssue")
+            @Override
+            public <T> DataResult<Pair<NoiseGeneratorSettings, T>> decode(DynamicOps<T> ops, T input) {
+                return original.decode(ops, input).flatMap(pair -> {
+                    ops.getMap(input).result().ifPresent(mapLike -> {
+                        var generatorKey = mapLike.get(BIOMESQUISHER_GENERATOR_KEY);
+                        if (generatorKey != null) {
+                            var key = RESOURCE_KEY_CODEC.parse(ops, generatorKey).result();
+                            key.ifPresent(s -> ((KnowsOriginalKey) (Object) pair.getFirst()).biomesquisher_generatorKey(s));
+                        }
+                    });
+
+                    return DataResult.success(pair);
+                });
+            }
+
+            @SuppressWarnings("DataFlowIssue")
+            @Override
+            public <T> DataResult<T> encode(NoiseGeneratorSettings input, DynamicOps<T> ops, T prefix) {
+                return original.encode(input, ops, prefix).flatMap(it -> {
+                    var key = ((KnowsOriginalKey) (Object) input).biomesquisher_generatorKey();
+
+                    if (key == null) {
+                        // capture the key from a wrapped codec if it's null
+                        //noinspection unchecked
+                        ResourceKey<NoiseGeneratorSettings>[] keyHolder = new ResourceKey[1];
+                        WrappingRuleSource.NotifyingOps.NotifyingJsonOps notifyingOps = WrappingRuleSource.NotifyingOps.NotifyingJsonOps.create(wrapped -> keyHolder[0] = wrapped.generator());
+                        SurfaceRules.RuleSource.CODEC.encodeStart(notifyingOps, input.surfaceRule());
+                        key = keyHolder[0];
+                    }
+
+                    if (key != null) {
+                        return RESOURCE_KEY_CODEC.encode(key, ops, ops.empty()).flatMap(rk -> ops.mergeToMap(it, ops.createString(BIOMESQUISHER_GENERATOR_KEY), rk));
+                    }
+                    return DataResult.success(it);
+                });
+            }
+        };
+    }
+
     interface NotifyingOps {
         void wrapped(WrappingRuleSource source);
 
@@ -76,8 +124,10 @@ protected NotifyingJsonOps(boolean compressed, Consumer<WrappingRuleSource> acti
 
             @Override
             public void wrapped(WrappingRuleSource source) {
-                this.action.accept(source);
-                this.wrapped = true;
+                if (!this.wrapped) {
+                    this.action.accept(source);
+                    this.wrapped = true;
+                }
             }
 
             public boolean isWrapped() {
@@ -92,18 +142,14 @@ public static NotifyingJsonOps create(Consumer<WrappingRuleSource> consumer) {
 
     @Constructor
     @SuppressWarnings("unused")
-    static WrappingRuleSource create(@Field("delegate") @Field.Final SurfaceRules.RuleSource delegate, @Field("generator") ResourceKey<NoiseGeneratorSettings> generator) {
+    static WrappingRuleSource create(@Field("delegate") @Field.Final SurfaceRules.RuleSource delegate, @Field("generator") @Field.Final ResourceKey<NoiseGeneratorSettings> generator) {
         throw new UnsupportedOperationException("Replaced by OpenSesame at compile time");
     }
 
     @Field("generator")
     ResourceKey<NoiseGeneratorSettings> generator();
 
-    @Field("generator")
-    void generator(ResourceKey<NoiseGeneratorSettings> generator);
-
     @Field("delegate")
-    @Field.Final
     SurfaceRules.RuleSource delegate();
 
     @Field("modifiers")

diff --git a/src/main/java/dev/lukebemish/biomesquisher/impl/injected/KnowsOriginalKey.java b/src/main/java/dev/lukebemish/biomesquisher/impl/injected/KnowsOriginalKey.java
@@ -0,0 +1,9 @@
+package dev.lukebemish.biomesquisher.impl.injected;
+
+import net.minecraft.resources.ResourceKey;
+import net.minecraft.world.level.levelgen.NoiseGeneratorSettings;
+
+public interface KnowsOriginalKey {
+    ResourceKey<NoiseGeneratorSettings> biomesquisher_generatorKey();
+    void biomesquisher_generatorKey(ResourceKey<NoiseGeneratorSettings> key);
+}
diff --git a/src/main/java/dev/lukebemish/biomesquisher/impl/mixin/MinecraftServerMixin.java b/src/main/java/dev/lukebemish/biomesquisher/impl/mixin/MinecraftServerMixin.java
@@ -53,8 +53,9 @@ private void biomesquisher_load(
                     Utils.LOGGER.info("Not squishing {}; not a MultiNoiseBiomeSource", key.location());
                 }
 
+                var backupKey = ResourceKey.create(Registries.NOISE_SETTINGS, key.location());
                 var settings = generator.generatorSettings().value();
-                BiomeSquisher.modifySurfaceRules(settings, access);
+                BiomeSquisher.modifySurfaceRules(settings, access, backupKey);
             } else {
                 Utils.LOGGER.info("Not squishing {}; not a NoiseBasedChunkGenerator", key.location());
             }

diff --git a/src/main/java/dev/lukebemish/biomesquisher/impl/mixin/NoiseGeneratorSettingsMixin.java b/src/main/java/dev/lukebemish/biomesquisher/impl/mixin/NoiseGeneratorSettingsMixin.java
@@ -0,0 +1,38 @@
+package dev.lukebemish.biomesquisher.impl.mixin;
+
+import com.llamalad7.mixinextras.injector.ModifyExpressionValue;
+import com.mojang.serialization.Codec;
+import dev.lukebemish.biomesquisher.impl.WrappingRuleSource;
+import dev.lukebemish.biomesquisher.impl.injected.KnowsOriginalKey;
+import net.minecraft.resources.ResourceKey;
+import net.minecraft.world.level.levelgen.NoiseGeneratorSettings;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Unique;
+import org.spongepowered.asm.mixin.injection.At;
+
+@Mixin(NoiseGeneratorSettings.class)
+public class NoiseGeneratorSettingsMixin implements KnowsOriginalKey {
+    @Unique
+    private ResourceKey<NoiseGeneratorSettings> biomesquisher_generatorKey;
+
+    @Override
+    public ResourceKey<NoiseGeneratorSettings> biomesquisher_generatorKey() {
+        return biomesquisher_generatorKey;
+    }
+
+    @Override
+    public synchronized void biomesquisher_generatorKey(ResourceKey<NoiseGeneratorSettings> key) {
+        this.biomesquisher_generatorKey = key;
+    }
+
+    @ModifyExpressionValue(
+        method = "<clinit>()V",
+        at = @At(
+            value = "INVOKE",
+            target = "Lcom/mojang/serialization/codecs/RecordCodecBuilder;create(Ljava/util/function/Function;)Lcom/mojang/serialization/Codec;"
+        )
+    )
+    private static Codec<NoiseGeneratorSettings> biome_squisher$wrapCodec(Codec<NoiseGeneratorSettings> original) {
+        return WrappingRuleSource.wrap(original);
+    }
+}
diff --git a/src/main/java/dev/lukebemish/biomesquisher/impl/mixin/RuleSourceMixin.java b/src/main/java/dev/lukebemish/biomesquisher/impl/mixin/RuleSourceMixin.java
@@ -0,0 +1,41 @@
+package dev.lukebemish.biomesquisher.impl.mixin;
+
+import com.llamalad7.mixinextras.injector.ModifyExpressionValue;
+import com.mojang.datafixers.util.Pair;
+import com.mojang.serialization.Codec;
+import com.mojang.serialization.DataResult;
+import com.mojang.serialization.DynamicOps;
+import dev.lukebemish.biomesquisher.impl.WrappingRuleSource;
+import net.minecraft.world.level.levelgen.SurfaceRules;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.injection.At;
+
+@Mixin(SurfaceRules.RuleSource.class)
+public interface RuleSourceMixin {
+    @ModifyExpressionValue(
+        method = "<clinit>()V",
+        at = @At(
+            value = "INVOKE",
+            target = "Lcom/mojang/serialization/Codec;dispatch(Ljava/util/function/Function;Ljava/util/function/Function;)Lcom/mojang/serialization/Codec;"
+        )
+    )
+    private static Codec<SurfaceRules.RuleSource> biome_squisher$wrapCodec(Codec<SurfaceRules.RuleSource> original) {
+        return new Codec<>() {
+            @Override
+            public <T> DataResult<Pair<SurfaceRules.RuleSource, T>> decode(DynamicOps<T> ops, T input) {
+                return original.decode(ops, input);
+            }
+
+            @Override
+            public <T> DataResult<T> encode(SurfaceRules.RuleSource input, DynamicOps<T> ops, T prefix) {
+                if (input instanceof WrappingRuleSource wrapped) {
+                    if (ops instanceof WrappingRuleSource.NotifyingOps notifying) {
+                        notifying.wrapped(wrapped);
+                    }
+                    return encode(wrapped.delegate(), ops, prefix);
+                }
+                return original.encode(input, ops, prefix);
+            }
+        };
+    }
+}
diff --git a/src/main/java/dev/lukebemish/biomesquisher/impl/mixin/WorldLoaderMixin.java b/src/main/java/dev/lukebemish/biomesquisher/impl/mixin/WorldLoaderMixin.java
@@ -29,8 +29,8 @@ public abstract class WorldLoaderMixin {
         noiseSettingsRegistry.forEach(value -> {
             ResourceKey<NoiseGeneratorSettings> key = noiseSettingsRegistry.getResourceKey(value).orElseThrow();
             Utils.LOGGER.info("Modifying surface rules in {}", key.location());
-            BiomeSquisher.setupSurfaceRuleModification(value, key);
-            BiomeSquisher.modifySurfaceRules(value, access);
+            BiomeSquisher.setupOriginalKeyAwareGenerators(value, key);
+            BiomeSquisher.modifySurfaceRules(value, access, key);
         });
 
         return original;

diff --git a/src/main/resources/biomesquisher.mixins.json b/src/main/resources/biomesquisher.mixins.json
@@ -8,7 +8,9 @@
     "ParameterListMixin",
     "MinecraftServerMixin",
     "MultiNoiseBiomeSourceMixin",
-    "WorldLoaderMixin"
+    "WorldLoaderMixin",
+    "RuleSourceMixin",
+    "NoiseGeneratorSettingsMixin"
   ],
 	"injectors": {
 		"defaultRequire": 1