properly check for experimental stuff in generator

4 files changed, 27 insertions, 17 deletions

diff --git a/paper-api-generator/generated/io/papermc/paper/registry/keys/BiomeKeys.java b/paper-api-generator/generated/io/papermc/paper/registry/keys/BiomeKeys.java
index fc62be44d6..3bd15016b0 100644
--- a/paper-api-generator/generated/io/papermc/paper/registry/keys/BiomeKeys.java
+++ b/paper-api-generator/generated/io/papermc/paper/registry/keys/BiomeKeys.java
@@ -6,7 +6,6 @@ import io.papermc.paper.generated.GeneratedFrom;
 import io.papermc.paper.registry.RegistryKey;
 import io.papermc.paper.registry.TypedKey;
 import net.kyori.adventure.key.Key;
-import org.bukkit.MinecraftExperimental;
 import org.bukkit.block.Biome;
 import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.NotNull;
@@ -249,8 +248,6 @@ public final class BiomeKeys {
      *
      * @apiNote This field is version-dependant and may be removed in future Minecraft versions
      */
-    @ApiStatus.Experimental
-    @MinecraftExperimental(MinecraftExperimental.Requires.UPDATE_1_21)
     public static final TypedKey<Biome> MANGROVE_SWAMP = create(key("mangrove_swamp"));
 
     /**
@@ -405,8 +402,6 @@ public final class BiomeKeys {
      *
      * @apiNote This field is version-dependant and may be removed in future Minecraft versions
      */
-    @ApiStatus.Experimental
-    @MinecraftExperimental(MinecraftExperimental.Requires.UPDATE_1_21)
     public static final TypedKey<Biome> SWAMP = create(key("swamp"));
 
     /**
diff --git a/paper-api-generator/src/main/java/io/papermc/generator/types/GeneratedKeyType.java b/paper-api-generator/src/main/java/io/papermc/generator/types/GeneratedKeyType.java
index ba27aa640d..66c4a4b14f 100644
--- a/paper-api-generator/src/main/java/io/papermc/generator/types/GeneratedKeyType.java
+++ b/paper-api-generator/src/main/java/io/papermc/generator/types/GeneratedKeyType.java
@@ -1,5 +1,6 @@
 package io.papermc.generator.types;
 
+import com.google.common.collect.Sets;
 import com.squareup.javapoet.FieldSpec;
 import com.squareup.javapoet.JavaFile;
 import com.squareup.javapoet.MethodSpec;
@@ -16,17 +17,18 @@ import io.papermc.paper.registry.RegistryKey;
 import io.papermc.paper.registry.TypedKey;
 import java.lang.reflect.Field;
 import java.lang.reflect.Modifier;
-import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.List;
+import java.util.IdentityHashMap;
 import java.util.Map;
+import java.util.Set;
 import java.util.stream.Collectors;
 import net.kyori.adventure.key.Key;
 import net.minecraft.core.Holder;
 import net.minecraft.core.Registry;
 import net.minecraft.core.RegistrySetBuilder;
 import net.minecraft.data.registries.UpdateOneTwentyOneRegistries;
+import net.minecraft.data.registries.VanillaRegistries;
 import net.minecraft.resources.ResourceKey;
 import org.bukkit.MinecraftExperimental;
 import org.checkerframework.checker.nullness.qual.NonNull;
@@ -46,6 +48,9 @@ import static javax.lang.model.element.Modifier.STATIC;
 @DefaultQualifier(NonNull.class)
 public class GeneratedKeyType<T, A> extends SimpleGenerator {
 
+    private static final Map<ResourceKey<? extends Registry<?>>, RegistrySetBuilder.RegistryBootstrap<?>> VANILLA_REGISTRY_ENTRIES = VanillaRegistries.BUILDER.entries.stream()
+            .collect(Collectors.toMap(RegistrySetBuilder.RegistryStub::key, RegistrySetBuilder.RegistryStub::bootstrap));
+
     private static final Map<ResourceKey<? extends Registry<?>>, RegistrySetBuilder.RegistryBootstrap<?>> EXPERIMENTAL_REGISTRY_ENTRIES = UpdateOneTwentyOneRegistries.BUILDER.entries.stream()
             .collect(Collectors.toMap(RegistrySetBuilder.RegistryStub::key, RegistrySetBuilder.RegistryStub::bootstrap));
 
@@ -120,7 +125,7 @@ public class GeneratedKeyType<T, A> extends SimpleGenerator {
         final MethodSpec.Builder createMethod = this.createMethod(typedKey);
 
         final Registry<T> registry = Main.REGISTRY_ACCESS.registryOrThrow(this.registryKey);
-        final List<ResourceKey<T>> experimental = this.collectExperimentalKeys(registry);
+        final Set<ResourceKey<T>> experimental = this.collectExperimentalKeys(registry);
 
         boolean allExperimental = true;
         for (final Holder.Reference<T> reference : registry.holders().sorted(Formatting.alphabeticKeyOrder(reference -> reference.key().location().getPath())).toList()) {
@@ -147,19 +152,27 @@ public class GeneratedKeyType<T, A> extends SimpleGenerator {
     }
 
     @SuppressWarnings("unchecked")
-    private List<ResourceKey<T>> collectExperimentalKeys(final Registry<T> registry) {
-        final RegistrySetBuilder.@Nullable RegistryBootstrap<T> registryBootstrap = (RegistrySetBuilder.RegistryBootstrap<T>) EXPERIMENTAL_REGISTRY_ENTRIES.get(this.registryKey);
-        if (registryBootstrap == null) {
-            return Collections.emptyList();
+    private Set<ResourceKey<T>> collectExperimentalKeys(final Registry<T> registry) {
+        final RegistrySetBuilder.@Nullable RegistryBootstrap<T> experimentalBootstrap = (RegistrySetBuilder.RegistryBootstrap<T>) EXPERIMENTAL_REGISTRY_ENTRIES.get(this.registryKey);
+        if (experimentalBootstrap == null) {
+            return Collections.emptySet();
+        }
+        final Set<ResourceKey<T>> experimental = Collections.newSetFromMap(new IdentityHashMap<>());
+        final CollectingContext<T> experimentalCollector = new CollectingContext<>(experimental, registry);
+        experimentalBootstrap.run(experimentalCollector);
+
+        final RegistrySetBuilder.@Nullable RegistryBootstrap<T> vanillaBootstrap = (RegistrySetBuilder.RegistryBootstrap<T>) VANILLA_REGISTRY_ENTRIES.get(this.registryKey);
+        if (vanillaBootstrap != null) {
+            final Set<ResourceKey<T>> vanilla = Collections.newSetFromMap(new IdentityHashMap<>());
+            final CollectingContext<T> vanillaCollector = new CollectingContext<>(vanilla, registry);
+            vanillaBootstrap.run(vanillaCollector);
+            return Sets.difference(experimental, vanilla);
         }
-        final List<ResourceKey<T>> experimental = new ArrayList<>();
-        final CollectingContext<T> context = new CollectingContext<>(experimental, registry);
-        registryBootstrap.run(context);
         return experimental;
     }
 
     @Override
-    protected JavaFile.Builder file(JavaFile.Builder builder) {
+    protected JavaFile.Builder file(final JavaFile.Builder builder) {
         return builder
             .skipJavaLangImports(true)
             .addStaticImport(Key.class, "key")
diff --git a/paper-api-generator/src/main/java/io/papermc/generator/utils/CollectingContext.java b/paper-api-generator/src/main/java/io/papermc/generator/utils/CollectingContext.java
index 614babf9e0..15cb4ac3e7 100644
--- a/paper-api-generator/src/main/java/io/papermc/generator/utils/CollectingContext.java
+++ b/paper-api-generator/src/main/java/io/papermc/generator/utils/CollectingContext.java
@@ -3,6 +3,7 @@ package io.papermc.generator.utils;
 import com.mojang.serialization.Lifecycle;
 import io.papermc.generator.Main;
 import java.util.List;
+import java.util.Set;
 import net.minecraft.core.Holder;
 import net.minecraft.core.HolderGetter;
 import net.minecraft.core.Registry;
@@ -12,7 +13,7 @@ import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.framework.qual.DefaultQualifier;
 
 @DefaultQualifier(NonNull.class)
-public record CollectingContext<T>(List<ResourceKey<T>> registered,
+public record CollectingContext<T>(Set<ResourceKey<T>> registered,
                                    Registry<T> registry) implements BootstrapContext<T> {
 
     @Override
diff --git a/paper-api-generator/wideners.at b/paper-api-generator/wideners.at
index 54524a165e..6928316973 100644
--- a/paper-api-generator/wideners.at
+++ b/paper-api-generator/wideners.at
@@ -4,3 +4,4 @@ public net/minecraft/server/WorldLoader loadAndReplaceLayer(Lnet/minecraft/serve
 public net/minecraft/core/RegistrySetBuilder entries
 public net/minecraft/core/RegistrySetBuilder$RegistryStub
 public net/minecraft/data/registries/UpdateOneTwentyOneRegistries BUILDER
+public net/minecraft/data/registries/VanillaRegistries BUILDER