update generator

8 files changed, 133 insertions, 78 deletions

diff --git a/paper-api-generator/src/main/java/io/papermc/generator/Main.java b/paper-api-generator/src/main/java/io/papermc/generator/Main.java
index e02f55871f..fa2b05e94c 100644
--- a/paper-api-generator/src/main/java/io/papermc/generator/Main.java
+++ b/paper-api-generator/src/main/java/io/papermc/generator/Main.java
@@ -1,20 +1,14 @@
 package io.papermc.generator;
 
 import com.mojang.logging.LogUtils;
-import io.papermc.generator.types.GeneratedKeyType;
 import io.papermc.generator.types.SourceGenerator;
-import io.papermc.paper.registry.RegistryKey;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
-import java.util.List;
 import net.minecraft.SharedConstants;
 import net.minecraft.core.LayeredRegistryAccess;
-import net.minecraft.core.Registry;
 import net.minecraft.core.RegistryAccess;
-import net.minecraft.core.registries.Registries;
 import net.minecraft.resources.RegistryDataLoader;
-import net.minecraft.resources.ResourceKey;
 import net.minecraft.server.Bootstrap;
 import net.minecraft.server.RegistryLayer;
 import net.minecraft.server.WorldLoader;
@@ -24,12 +18,6 @@ import net.minecraft.server.packs.repository.PackRepository;
 import net.minecraft.server.packs.repository.ServerPacksSource;
 import net.minecraft.server.packs.resources.MultiPackResourceManager;
 import org.apache.commons.io.file.PathUtils;
-import org.bukkit.GameEvent;
-import org.bukkit.block.Biome;
-import org.bukkit.generator.structure.Structure;
-import org.bukkit.generator.structure.StructureType;
-import org.bukkit.inventory.meta.trim.TrimMaterial;
-import org.bukkit.inventory.meta.trim.TrimPattern;
 import org.slf4j.Logger;
 
 public final class Main {
@@ -40,6 +28,8 @@ public final class Main {
     static {
         SharedConstants.tryDetectVersion();
         Bootstrap.bootStrap();
+        Bootstrap.validate();
+
         final PackRepository resourceRepository = ServerPacksSource.createVanillaTrustedRepository();
         resourceRepository.reload();
         final MultiPackResourceManager resourceManager = new MultiPackResourceManager(PackType.SERVER_DATA, resourceRepository.getAvailablePacks().stream().map(Pack::open).toList());
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 b5ad821c1c..ba27aa640d 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
@@ -10,6 +10,7 @@ import com.squareup.javapoet.TypeSpec;
 import io.papermc.generator.Main;
 import io.papermc.generator.utils.Annotations;
 import io.papermc.generator.utils.CollectingContext;
+import io.papermc.generator.utils.Formatting;
 import io.papermc.generator.utils.Javadocs;
 import io.papermc.paper.registry.RegistryKey;
 import io.papermc.paper.registry.TypedKey;
@@ -19,14 +20,15 @@ import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
-import java.util.Locale;
 import java.util.Map;
 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.resources.ResourceKey;
+import org.bukkit.MinecraftExperimental;
 import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.checkerframework.framework.qual.DefaultQualifier;
@@ -115,30 +117,31 @@ public class GeneratedKeyType<T, A> extends SimpleGenerator {
         final TypeName typedKey = ParameterizedTypeName.get(TypedKey.class, this.apiType);
 
         final TypeSpec.Builder typeBuilder = this.keyHolderType();
-        typeBuilder.addAnnotation(EXPERIMENTAL_API_ANNOTATION); // TODO experimental API
         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);
 
         boolean allExperimental = true;
-        for (final T value : registry) {
-            final ResourceKey<T> key = registry.getResourceKey(value).orElseThrow();
+        for (final Holder.Reference<T> reference : registry.holders().sorted(Formatting.alphabeticKeyOrder(reference -> reference.key().location().getPath())).toList()) {
+            final ResourceKey<T> key = reference.key();
             final String keyPath = key.location().getPath();
-            final String fieldName = keyPath.toUpperCase(Locale.ENGLISH).replaceAll("[.-/]", "_"); // replace invalid field name chars
+            final String fieldName = Formatting.formatKeyAsField(keyPath);
             final FieldSpec.Builder fieldBuilder = FieldSpec.builder(typedKey, fieldName, PUBLIC, STATIC, FINAL)
                 .initializer("$N(key($S))", createMethod.build(), keyPath)
                 .addJavadoc(Javadocs.getVersionDependentField("{@code $L}"), key.location().toString());
             if (experimental.contains(key)) {
-                fieldBuilder.addAnnotations(experimentalAnnotations("MinecraftExperimental.Requires.UPDATE_1_21"));
+                fieldBuilder.addAnnotations(experimentalAnnotations(MinecraftExperimental.Requires.UPDATE_1_21));
             } else {
                 allExperimental = false;
             }
             typeBuilder.addField(fieldBuilder.build());
         }
         if (allExperimental) {
-            typeBuilder.addAnnotations(experimentalAnnotations("MinecraftExperimental.Requires.UPDATE_1_21"));
-            createMethod.addAnnotations(experimentalAnnotations("MinecraftExperimental.Requires.UPDATE_1_21"));
+            typeBuilder.addAnnotations(experimentalAnnotations(MinecraftExperimental.Requires.UPDATE_1_21));
+            createMethod.addAnnotations(experimentalAnnotations(MinecraftExperimental.Requires.UPDATE_1_21));
+        } else {
+            typeBuilder.addAnnotation(EXPERIMENTAL_API_ANNOTATION); // TODO experimental API
         }
         return typeBuilder.addMethod(createMethod.build()).build();
     }
diff --git a/paper-api-generator/src/main/java/io/papermc/generator/types/SimpleGenerator.java b/paper-api-generator/src/main/java/io/papermc/generator/types/SimpleGenerator.java
index 76edecd9ff..b4c1095c84 100644
--- a/paper-api-generator/src/main/java/io/papermc/generator/types/SimpleGenerator.java
+++ b/paper-api-generator/src/main/java/io/papermc/generator/types/SimpleGenerator.java
@@ -5,7 +5,6 @@ import com.squareup.javapoet.TypeSpec;
 
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
-import java.nio.file.Files;
 import java.nio.file.Path;
 
 public abstract class SimpleGenerator implements SourceGenerator {
@@ -24,14 +23,13 @@ public abstract class SimpleGenerator implements SourceGenerator {
 
     @Override
     public void writeToFile(Path parent) throws IOException {
-        Path packagePath = parent.resolve(this.packageName.replace('.', '/'));
-        Files.createDirectories(packagePath);
 
-        JavaFile.Builder builder = JavaFile.builder(this.packageName, this.getTypeSpec())
-            .indent("    ");
-        this.file(builder);
+        JavaFile.Builder builder = JavaFile.builder(this.packageName, this.getTypeSpec());
+        this.file(builder)
+            .indent("    ")
+            .skipJavaLangImports(true);
 
-        Files.writeString(packagePath.resolve(this.className + ".java"), this.file(builder).build().toString(), StandardCharsets.UTF_8);
+        builder.build().writeTo(parent, StandardCharsets.UTF_8);
     }
 
 }
diff --git a/paper-api-generator/src/main/java/io/papermc/generator/types/goal/MobGoalGenerator.java b/paper-api-generator/src/main/java/io/papermc/generator/types/goal/MobGoalGenerator.java
index 649490b5b0..7f94034424 100644
--- a/paper-api-generator/src/main/java/io/papermc/generator/types/goal/MobGoalGenerator.java
+++ b/paper-api-generator/src/main/java/io/papermc/generator/types/goal/MobGoalGenerator.java
@@ -2,7 +2,6 @@ package io.papermc.generator.types.goal;
 
 import com.destroystokyo.paper.entity.RangedEntity;
 import com.destroystokyo.paper.entity.ai.GoalKey;
-import com.squareup.javapoet.AnnotationSpec;
 import com.squareup.javapoet.ClassName;
 import com.squareup.javapoet.FieldSpec;
 import com.squareup.javapoet.JavaFile;
@@ -19,12 +18,7 @@ import io.papermc.generator.utils.Annotations;
 import io.papermc.generator.utils.Formatting;
 import io.papermc.generator.utils.Javadocs;
 import java.util.Comparator;
-import java.util.LinkedHashMap;
 import java.util.List;
-import java.util.Map;
-import java.util.TreeMap;
-import java.util.function.Supplier;
-import java.util.stream.Collectors;
 import javax.lang.model.element.Modifier;
 import net.minecraft.world.entity.ai.goal.Goal;
 import net.minecraft.world.entity.ai.goal.WrappedGoal;
@@ -65,9 +59,13 @@ import org.bukkit.entity.Vindicator;
 import org.bukkit.entity.WanderingTrader;
 import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.framework.qual.DefaultQualifier;
-import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.Nullable;
 
+import static javax.lang.model.element.Modifier.FINAL;
+import static javax.lang.model.element.Modifier.PRIVATE;
+import static javax.lang.model.element.Modifier.PUBLIC;
+import static javax.lang.model.element.Modifier.STATIC;
+
 @DefaultQualifier(NonNull.class)
 public class MobGoalGenerator extends SimpleGenerator {
 
@@ -168,57 +166,53 @@ public class MobGoalGenerator extends SimpleGenerator {
 
     @Override
     protected TypeSpec getTypeSpec() {
-        TypeName clazzType = TypeName.get(Class.class)
-            .annotated(Annotations.NOT_NULL);
-        TypeName keyType = TypeName.get(String.class)
-            .annotated(Annotations.NOT_NULL);
-
-        MethodSpec.Builder createMethod = MethodSpec.methodBuilder("create")
-            .addModifiers(Modifier.PRIVATE, Modifier.STATIC)
-            .addParameter(ParameterSpec.builder(keyType, "key", Modifier.FINAL)
-                .build()
-            )
-            .addParameter(ParameterSpec.builder(clazzType, "clazz", Modifier.FINAL)
-                .build()
-            )
-            .addCode("return $T.of(clazz, $T.minecraft(key));", GoalKey.class, NamespacedKey.class)
-            .returns(ParameterizedTypeName.get(GoalKey.class).annotated(Annotations.NOT_NULL));
-
-
         TypeVariableName type = TypeVariableName.get("T", Mob.class);
         TypeSpec.Builder typeBuilder = TypeSpec.interfaceBuilder(this.className)
             .addSuperinterface(ParameterizedTypeName.get(ClassName.get(com.destroystokyo.paper.entity.ai.Goal.class), type))
-            .addModifiers(Modifier.PUBLIC)
+            .addModifiers(PUBLIC)
             .addTypeVariable(type)
             .addAnnotations(Annotations.CLASS_HEADER)
             .addJavadoc(CLASS_HEADER);
 
+        TypeName mobType = ParameterizedTypeName.get(ClassName.get(Class.class), type)
+            .annotated(Annotations.NOT_NULL);
+        TypeName keyType = TypeName.get(String.class)
+            .annotated(Annotations.NOT_NULL);
 
-        List<Class<?>> classes;
+        ParameterSpec keyParam = ParameterSpec.builder(keyType, "key", FINAL).build();
+        ParameterSpec typeParam = ParameterSpec.builder(mobType, "type", FINAL).build();
+        MethodSpec.Builder createMethod = MethodSpec.methodBuilder("create")
+            .addModifiers(PRIVATE, STATIC)
+            .addParameter(keyParam)
+            .addParameter(typeParam)
+            .addCode("return $T.of($N, $T.minecraft($N));", GoalKey.class, typeParam, NamespacedKey.class, keyParam)
+            .addTypeVariable(type)
+            .returns(ParameterizedTypeName.get(ClassName.get(GoalKey.class), type).annotated(Annotations.NOT_NULL));
+
+        List<Class<Goal>> classes;
         try (ScanResult scanResult = new ClassGraph().enableAllInfo().whitelistPackages("net.minecraft").scan()) {
-            classes = scanResult.getSubclasses(net.minecraft.world.entity.ai.goal.Goal.class.getName()).loadClasses();
+            classes = scanResult.getSubclasses(Goal.class.getName()).loadClasses(Goal.class);
         }
 
-        List<VanillaGoalKey> vanillaNames = classes.stream()
+        List<GoalKey<Mob>> vanillaNames = classes.stream()
             .filter(clazz -> !java.lang.reflect.Modifier.isAbstract(clazz.getModifiers()))
             .filter(clazz -> !WrappedGoal.class.equals(clazz)) // TODO - properly fix
-            .map(goalClass -> new VanillaGoalKey(goalClass, MobGoalNames.getKey(goalClass.getName(), (Class<? extends Goal>) goalClass)))
-            .filter((key) -> !MobGoalNames.isIgnored(key.key().getNamespacedKey().getKey()))
-            .sorted(Comparator.<VanillaGoalKey, String>comparing(o -> o.key().getEntityClass().getSimpleName())
-                .thenComparing(vanillaGoalKey -> vanillaGoalKey.key.getNamespacedKey().getKey())
+            .map(goalClass -> MobGoalNames.getKey(goalClass.getName(), goalClass))
+            .filter((key) -> !MobGoalNames.isIgnored(key.getNamespacedKey().getKey()))
+            .sorted(Comparator.<GoalKey<?>, String>comparing(o -> o.getEntityClass().getSimpleName())
+                .thenComparing(vanillaGoalKey -> vanillaGoalKey.getNamespacedKey().getKey())
             )
             .toList();
 
 
-        for (final VanillaGoalKey vanillaGoalKey : vanillaNames) {
-            GoalKey<?> value = vanillaGoalKey.key();
-            TypeName typedKey = ParameterizedTypeName.get(GoalKey.class, value.getEntityClass());
-            NamespacedKey key = value.getNamespacedKey();
+        for (final GoalKey<?> goalKey : vanillaNames) {
+            TypeName typedKey = ParameterizedTypeName.get(GoalKey.class, goalKey.getEntityClass());
+            NamespacedKey key = goalKey.getNamespacedKey();
 
             String keyPath = key.getKey();
-            String fieldName = Formatting.formatKeyAsField(key);
-            FieldSpec.Builder fieldBuilder = FieldSpec.builder(typedKey, fieldName, Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL)
-                .initializer("$N($S, $T.class)", createMethod.build(), keyPath, value.getEntityClass());
+            String fieldName = Formatting.formatKeyAsField(keyPath);
+            FieldSpec.Builder fieldBuilder = FieldSpec.builder(typedKey, fieldName, PUBLIC, STATIC, FINAL)
+                .initializer("$N($S, $T.class)", createMethod.build(), keyPath, goalKey.getEntityClass());
             typeBuilder.addField(fieldBuilder.build());
         }
 
@@ -227,8 +221,8 @@ public class MobGoalGenerator extends SimpleGenerator {
             NamespacedKey key = NamespacedKey.minecraft(value.entryName);
 
             String keyPath = key.getKey();
-            String fieldName = Formatting.formatKeyAsField(key);
-            FieldSpec.Builder fieldBuilder = FieldSpec.builder(typedKey, fieldName, Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL)
+            String fieldName = Formatting.formatKeyAsField(keyPath);
+            FieldSpec.Builder fieldBuilder = FieldSpec.builder(typedKey, fieldName, PUBLIC, STATIC, FINAL)
                 .addAnnotation(Annotations.deprecatedVersioned(value.removedVersion, value.removalVersion != null))
                 .initializer("$N($S, $T.class)", createMethod.build(), keyPath, value.entity);
 
@@ -251,9 +245,6 @@ public class MobGoalGenerator extends SimpleGenerator {
             .skipJavaLangImports(true);
     }
 
-    record VanillaGoalKey(Class<?> clazz, GoalKey<?> key) {
-    }
-
     record DeprecatedEntry(Class<?> entity, String entryName, @Nullable String removalVersion,
                            @Nullable String removedVersion) {
 
diff --git a/paper-api-generator/src/main/java/io/papermc/generator/types/goal/MobGoalNames.java b/paper-api-generator/src/main/java/io/papermc/generator/types/goal/MobGoalNames.java
index 37506e3f84..2ad5e0ac06 100644
--- a/paper-api-generator/src/main/java/io/papermc/generator/types/goal/MobGoalNames.java
+++ b/paper-api-generator/src/main/java/io/papermc/generator/types/goal/MobGoalNames.java
@@ -126,7 +126,7 @@ import java.util.Set;
 public class MobGoalNames {
 
     private static final Map<Class<? extends Goal>, Class<? extends Mob>> entityClassCache = new HashMap<>();
-    private static final Map<Class<? extends net.minecraft.world.entity.Mob>, Class<? extends Mob>> bukkitMap = new HashMap<>();
+    public static final Map<Class<? extends net.minecraft.world.entity.Mob>, Class<? extends Mob>> bukkitMap = new HashMap<>();
 
 
     static {
@@ -233,6 +233,8 @@ public class MobGoalNames {
         bukkitMap.put(net.minecraft.world.entity.animal.allay.Allay.class, org.bukkit.entity.Allay.class);
         bukkitMap.put(net.minecraft.world.entity.animal.sniffer.Sniffer.class, org.bukkit.entity.Sniffer.class);
         bukkitMap.put(net.minecraft.world.entity.monster.breeze.Breeze.class, org.bukkit.entity.Breeze.class);
+        bukkitMap.put(net.minecraft.world.entity.animal.armadillo.Armadillo.class, org.bukkit.entity.Armadillo.class);
+        bukkitMap.put(net.minecraft.world.entity.monster.Bogged.class, org.bukkit.entity.Bogged.class);
         //</editor-fold>
     }
 
diff --git a/paper-api-generator/src/main/java/io/papermc/generator/utils/Annotations.java b/paper-api-generator/src/main/java/io/papermc/generator/utils/Annotations.java
index 5c82a94e7a..a00a682330 100644
--- a/paper-api-generator/src/main/java/io/papermc/generator/utils/Annotations.java
+++ b/paper-api-generator/src/main/java/io/papermc/generator/utils/Annotations.java
@@ -12,11 +12,11 @@ import org.jetbrains.annotations.Nullable;
 
 public final class Annotations {
 
-    public static List<AnnotationSpec> experimentalAnnotations(final String version) {
+    public static List<AnnotationSpec> experimentalAnnotations(final MinecraftExperimental.Requires requiredFeatureFlag) {
         return List.of(
             AnnotationSpec.builder(ApiStatus.Experimental.class).build(),
             AnnotationSpec.builder(MinecraftExperimental.class)
-                .addMember("value", "$L", version)
+                .addMember("value", "$T.$L", MinecraftExperimental.Requires.class, requiredFeatureFlag.name())
                 .build()
         );
     }
diff --git a/paper-api-generator/src/main/java/io/papermc/generator/utils/Formatting.java b/paper-api-generator/src/main/java/io/papermc/generator/utils/Formatting.java
index f4b07411a4..a6c88a8a14 100644
--- a/paper-api-generator/src/main/java/io/papermc/generator/utils/Formatting.java
+++ b/paper-api-generator/src/main/java/io/papermc/generator/utils/Formatting.java
@@ -1,13 +1,47 @@
 package io.papermc.generator.utils;
 
-import net.kyori.adventure.key.Key;
-
+import org.apache.commons.lang3.math.NumberUtils;
+import java.util.Comparator;
 import java.util.Locale;
+import java.util.OptionalInt;
+import java.util.function.Function;
+import java.util.regex.Pattern;
 
 public final class Formatting {
 
-    public static String formatKeyAsField(Key key) {
-       return key.value().toUpperCase(Locale.ENGLISH).replaceAll("[.-/]", "_"); // replace invalid field name chars
+    private static final Pattern ILLEGAL_FIELD_CHARACTERS = Pattern.compile("[.-/]");
+
+    public static String formatKeyAsField(String path) {
+        return ILLEGAL_FIELD_CHARACTERS.matcher(path.toUpperCase(Locale.ENGLISH)).replaceAll("_");
+    }
+
+    public static Comparator<String> ALPHABETIC_KEY_ORDER = alphabeticKeyOrder(path -> path);
+
+    public static <T> Comparator<T> alphabeticKeyOrder(Function<T, String> mapper) {
+        return (o1, o2) -> {
+            String path1 = mapper.apply(o1);
+            String path2 = mapper.apply(o2);
+
+            OptionalInt trailingInt1 = tryParseTrailingInt(path1);
+            OptionalInt trailingInt2 = tryParseTrailingInt(path2);
+
+            if (trailingInt1.isPresent() && trailingInt2.isPresent()) {
+                return Integer.compare(trailingInt1.getAsInt(), trailingInt2.getAsInt());
+            }
+
+            return path1.compareTo(path2);
+        };
+    }
+
+    private static OptionalInt tryParseTrailingInt(String path) {
+        int delimiterIndex = path.lastIndexOf('_');
+        if (delimiterIndex != -1) {
+            String score = path.substring(delimiterIndex + 1);
+            if (NumberUtils.isDigits(score)) {
+                return OptionalInt.of(Integer.parseInt(score));
+            }
+        }
+        return OptionalInt.empty();
     }
 
     private Formatting() {
diff --git a/paper-api-generator/src/test/java/io/papermc/generator/MobGoalConverterTest.java b/paper-api-generator/src/test/java/io/papermc/generator/MobGoalConverterTest.java
new file mode 100644
index 0000000000..9bc9f6b400
--- /dev/null
+++ b/paper-api-generator/src/test/java/io/papermc/generator/MobGoalConverterTest.java
@@ -0,0 +1,37 @@
+package io.papermc.generator;
+
+import io.github.classgraph.ClassGraph;
+import io.github.classgraph.ScanResult;
+import io.papermc.generator.types.goal.MobGoalNames;
+import net.minecraft.world.entity.Entity;
+import net.minecraft.world.entity.Mob;
+import org.junit.jupiter.api.Test;
+import java.util.ArrayList;
+import java.util.List;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.fail;
+
+public class MobGoalConverterTest {
+
+    @Test
+    public void testBukkitMap() {
+        final List<Class<Mob>> classes;
+        try (ScanResult scanResult = new ClassGraph().enableAllInfo().whitelistPackages(Entity.class.getPackageName()).scan()) {
+            classes = scanResult.getSubclasses(Mob.class.getName()).loadClasses(Mob.class);
+        }
+
+        assertFalse(classes.isEmpty(), "There are supposed to be more than 0 mob classes!");
+
+        List<String> missingClasses = new ArrayList<>();
+        for (Class<Mob> nmsClass : classes) {
+            if (!MobGoalNames.bukkitMap.containsKey(nmsClass)) {
+                missingClasses.add(nmsClass.getCanonicalName());
+            }
+        }
+
+        if (!missingClasses.isEmpty()) {
+            fail("Missing some entity classes in the bukkit map: " + String.join(", ", missingClasses));
+        }
+    }
+}