Fix missing translation keys in entity extractor

extractor/src/main/java/dev/_00a/valence_extractor/extractors/Entities.java

@@ -21,6 +21,7 @@ import net.minecraft.util.math.GlobalPos;
 import net.minecraft.util.registry.Registry;
 import net.minecraft.util.registry.RegistryEntry;
 import net.minecraft.village.VillagerData;
+import org.jetbrains.annotations.Nullable;
 
 import java.lang.reflect.ParameterizedType;
 import java.util.*;
@@ -31,12 +32,12 @@ public class Entities implements Main.Extractor {
             bits(
                     "flags",
                     bit("on_fire", 0),
-                    bit("crouching", 1),
+                    bit("sneaking", 1),
                     bit("sprinting", 3),
                     bit("swimming", 4),
                     bit("invisible", 5),
                     bit("glowing", 6),
-                    bit("elytra_flying", 7)
+                    bit("fall_flying", 7)
             ),
             bits(
                     "projectile_flags",
@@ -253,27 +254,40 @@ public class Entities implements Main.Extractor {
     @Override
     @SuppressWarnings("unchecked")
     public JsonElement extract() throws IllegalAccessException, NoSuchFieldException {
-        final var entitiesJson = new JsonObject();
-        final var entityClasses = new HashSet<Class<? extends Entity>>();
-
-        final var dummyWorld = DummyWorld.INSTANCE;
 
+        final var entityClassToType = new HashMap<Class<? extends Entity>, EntityType<?>>();
         for (var f : EntityType.class.getFields()) {
             if (f.getType().equals(EntityType.class)) {
-                var entityType = (EntityType<?>) f.get(null);
                 var entityClass = (Class<? extends Entity>) ((ParameterizedType) f.getGenericType()).getActualTypeArguments()[0];
+                var entityType = (EntityType<?>) f.get(null);
 
-                // While we can use the tracked data registry and reflection to get the tracked fields on entities, we won't know what their default values are because they are assigned in the entity's constructor.
+                entityClassToType.put(entityClass, entityType);
-                // To obtain this, we create a dummy world to spawn the entities into and then read the data tracker field from the base entity class.
+            }
-                // We also handle player entities specially since they cannot be spawned with EntityType#create.
+        }
-                final var entityInstance = entityType.equals(EntityType.PLAYER) ? DummyPlayerEntity.INSTANCE : entityType.create(dummyWorld);
 
-                var dataTrackerField = Entity.class.getDeclaredField("dataTracker");
+        final var dataTrackerField = Entity.class.getDeclaredField("dataTracker");
         dataTrackerField.setAccessible(true);
 
-                while (entityClasses.add(entityClass)) {
+        var entitiesJson = new JsonObject();
+
+        for (var entry : entityClassToType.entrySet()) {
+            var entityClass = entry.getKey();
+            @Nullable var entityType = entry.getValue();
+            assert entityType != null;
+
+            // While we can use the tracked data registry and reflection to get the tracked fields on entities, we won't know what their default values are because they are assigned in the entity's constructor.
+            // To obtain this, we create a dummy world to spawn the entities into and read the data tracker field from the base entity class.
+            // We also handle player entities specially since they cannot be spawned with EntityType#create.
+            final var entityInstance = entityType.equals(EntityType.PLAYER) ? DummyPlayerEntity.INSTANCE : entityType.create(DummyWorld.INSTANCE);
+            final var dataTracker = (DataTracker) dataTrackerField.get(entityInstance);
+
+            while (entitiesJson.get(entityClass.getSimpleName()) == null) {
                 var entityJson = new JsonObject();
 
+                var parent = entityClass.getSuperclass();
+                var hasParent = parent != null && Entity.class.isAssignableFrom(parent);
+                entityJson.addProperty("parent", hasParent ? parent.getSimpleName() : null);
+
                 entityJson.add("translation_key", entityType != null ? new JsonPrimitive(entityType.getTranslationKey()) : null);
 
                 var fieldsJson = new JsonArray();
@@ -281,17 +295,16 @@ public class Entities implements Main.Extractor {
                     if (entityField.getType().equals(TrackedData.class)) {
                         entityField.setAccessible(true);
 
-                            var data = (TrackedData<?>) entityField.get(null);
+                        var trackedData = (TrackedData<?>) entityField.get(null);
 
                         var fieldJson = new JsonObject();
                         var fieldName = entityField.getName().toLowerCase(Locale.ROOT);
                         fieldJson.addProperty("name", fieldName);
-                            fieldJson.addProperty("index", data.getId());
+                        fieldJson.addProperty("index", trackedData.getId());
 
-                            var dataTracker = (DataTracker) dataTrackerField.get(entityInstance);
+                        var data = Entities.trackedDataToJson(trackedData, dataTracker);
-                            var res = Entities.trackedDataToJson(data, dataTracker);
+                        fieldJson.addProperty("type", data.left());
-                            fieldJson.addProperty("type", res.left());
+                        fieldJson.add("default_value", data.right());
-                            fieldJson.add("default_value", res.right());
 
                         var bitsJson = new JsonArray();
                         for (var bit : BIT_FIELDS.getOrDefault(fieldName, new Bit[]{})) {
@@ -307,19 +320,12 @@ public class Entities implements Main.Extractor {
                 }
                 entityJson.add("fields", fieldsJson);
 
-                    var parent = entityClass.getSuperclass();
-                    if (parent == null || !Entity.class.isAssignableFrom(parent)) {
-                        entityJson.add("parent", null);
                 entitiesJson.add(entityClass.getSimpleName(), entityJson);
-                        break;
-                    }
 
-                    entityJson.addProperty("parent", parent.getSimpleName());
+                if (!hasParent) break;
-                    entitiesJson.add(entityClass.getSimpleName(), entityJson);
 
                 entityClass = (Class<? extends Entity>) parent;
-                    entityType = null;
+                entityType = entityClassToType.get(entityClass);
-                }
             }
         }