Work on abstracting entity variants to reduce code duplication

core/src/main/java/org/geysermc/geyser/entity/type/living/animal/FrogEntity.java

@@ -25,7 +25,6 @@
 
 package org.geysermc.geyser.entity.type.living.animal;
 
-import net.kyori.adventure.key.Key;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.cloudburstmc.math.vector.Vector3f;
 import org.cloudburstmc.protocol.bedrock.data.entity.EntityDataTypes;
@@ -35,20 +34,17 @@ import org.geysermc.geyser.entity.type.Entity;
 import org.geysermc.geyser.item.type.Item;
 import org.geysermc.geyser.session.GeyserSession;
 import org.geysermc.geyser.session.cache.registry.JavaRegistries;
-import org.geysermc.geyser.session.cache.registry.RegistryEntryContext;
+import org.geysermc.geyser.session.cache.registry.JavaRegistryKey;
 import org.geysermc.geyser.session.cache.tags.ItemTag;
 import org.geysermc.geyser.session.cache.tags.Tag;
-import org.geysermc.geyser.util.MinecraftKey;
 import org.geysermc.mcprotocollib.protocol.data.game.entity.metadata.Pose;
-import org.geysermc.mcprotocollib.protocol.data.game.entity.metadata.type.IntEntityMetadata;
 import org.geysermc.mcprotocollib.protocol.data.game.entity.metadata.type.ObjectEntityMetadata;
 
-import java.util.Locale;
 import java.util.OptionalInt;
 import java.util.UUID;
-import java.util.function.Function;
 
-public class FrogEntity extends AnimalEntity {
+// TODO this is implementing VariantHolder<Object> until MCPL updates
+public class FrogEntity extends AnimalEntity implements VariantHolder<Object> {
     public FrogEntity(GeyserSession session, int entityId, long geyserId, UUID uuid, EntityDefinition<?> definition, Vector3f position, Vector3f motion, float yaw, float pitch, float headYaw) {
         super(session, entityId, geyserId, uuid, definition, position, motion, yaw, pitch, headYaw);
     }
@@ -62,15 +58,19 @@ public class FrogEntity extends AnimalEntity {
         super.setPose(pose);
     }
 
-    // TODO this is a holder when MCPL updates
-    // TODO also check if this works
-    public void setFrogVariant(IntEntityMetadata entityMetadata) {
-        BuiltInVariant variant = JavaRegistries.FROG_VARIANT.fromNetworkId(session, entityMetadata.getPrimitiveValue());
-        if (variant == null) {
-            variant = BuiltInVariant.TEMPERATE;
-        }
+    @Override
+    public JavaRegistryKey<BuiltInVariant> variantRegistry() {
+        return JavaRegistries.FROG_VARIANT;
+    }
 
-        dirtyMetadata.put(EntityDataTypes.VARIANT, variant.ordinal());
+    @Override
+    public void setBedrockVariant(int bedrockId) {
+        dirtyMetadata.put(EntityDataTypes.VARIANT, bedrockId);
+    }
+
+    @Override
+    public BuiltIn<Object> defaultVariant() {
+        return BuiltInVariant.TEMPERATE;
     }
 
     public void setTongueTarget(ObjectEntityMetadata<OptionalInt> entityMetadata) {
@@ -93,26 +93,9 @@ public class FrogEntity extends AnimalEntity {
 
     // Ordered by bedrock id
     // TODO: are these ordered correctly?
-    public enum BuiltInVariant {
+    public enum BuiltInVariant implements BuiltIn<Object> {
         TEMPERATE,
         COLD,
-        WARM;
-
-        public static final Function<RegistryEntryContext, BuiltInVariant> READER = context -> getByJavaIdentifier(context.id());
-
-        private final Key javaIdentifier;
-
-        BuiltInVariant() {
-            javaIdentifier = MinecraftKey.key(name().toLowerCase(Locale.ROOT));
-        }
-
-        public static @Nullable BuiltInVariant getByJavaIdentifier(Key identifier) {
-            for (BuiltInVariant variant : values()) {
-                if (variant.javaIdentifier.equals(identifier)) {
-                    return variant;
-                }
-            }
-            return null;
-        }
+        WARM
     }
 }

core/src/main/java/org/geysermc/geyser/entity/type/living/animal/VariantHolder.java

@@ -0,0 +1,126 @@
+/*
+ * Copyright (c) 2025 GeyserMC. http://geysermc.org
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ *
+ * @author GeyserMC
+ * @link https://github.com/GeyserMC/Geyser
+ */
+
+package org.geysermc.geyser.entity.type.living.animal;
+
+import net.kyori.adventure.key.Key;
+import org.geysermc.geyser.session.GeyserSession;
+import org.geysermc.geyser.session.cache.registry.JavaRegistryKey;
+import org.geysermc.geyser.session.cache.registry.RegistryEntryContext;
+import org.geysermc.geyser.util.MinecraftKey;
+import org.geysermc.mcprotocollib.protocol.data.game.Holder;
+import org.geysermc.mcprotocollib.protocol.data.game.entity.metadata.EntityMetadata;
+import org.geysermc.mcprotocollib.protocol.data.game.entity.metadata.MetadataType;
+
+import java.util.Locale;
+import java.util.function.Function;
+
+/**
+ * Utility interface to help set up data-driven entity variants for mobs.
+ *
+ * <p>This interface is designed for mobs that have their variant wrapped in a {@link Holder}. Implementations usually have to
+ * implement {@link VariantHolder#variantRegistry()}, {@link VariantHolder#setBedrockVariant(int)}, and {@link VariantHolder#defaultVariant()}, and should also
+ * have an enum with built-in variants on bedrock (implementing {@link BuiltIn}).</p>
+ *
+ * @param <Variant> the MCPL variant class that a {@link Holder} wraps.
+ */
+public interface VariantHolder<Variant> {
+
+    default void setVariant(EntityMetadata<Holder<Variant>, ? extends MetadataType<Holder<Variant>>> variant) {
+        setVariant(variant.getValue());
+    }
+
+    /**
+     * Sets the variant of the entity. Defaults to {@link VariantHolder#defaultVariant()} for custom holders and non-vanilla IDs.
+     */
+    default void setVariant(Holder<Variant> variant) {
+        BuiltIn<Variant> builtInVariant;
+        if (variant.isId()) {
+            builtInVariant = variantRegistry().fromNetworkId(getSession(), variant.id());
+            if (builtInVariant == null) {
+                builtInVariant = defaultVariant();
+            }
+        } else {
+            builtInVariant = defaultVariant();
+        }
+        setBedrockVariant(builtInVariant.ordinal());
+    }
+
+    GeyserSession getSession();
+
+    /**
+     * The registry in {@link org.geysermc.geyser.session.cache.registry.JavaRegistries} for this mob's variants. The registry can utilise the {@link VariantHolder#reader(Class)} method
+     * to create a reader to be used in {@link org.geysermc.geyser.session.cache.RegistryCache}.
+     */
+    JavaRegistryKey<? extends BuiltIn<Variant>> variantRegistry();
+
+    /**
+     * Should set the variant on bedrock's metadata (or however the variant is set for the mob). The bedrock ID has already been checked and is always valid.
+     */
+    void setBedrockVariant(int bedrockId);
+
+    /**
+     * Should return the default variant, that is to be used when this mob's variant is a custom or non-vanilla one.
+     */
+    BuiltIn<Variant> defaultVariant();
+
+    /**
+     * Creates a registry reader for this mob's variants.
+     *
+     * <p>This reader simply matches the identifiers of registry entries with built-in variants. If no built-in variant matches, null is returned.</p>
+     */
+    static <BuiltInVariant extends Enum<? extends BuiltIn<?>>> Function<RegistryEntryContext, BuiltInVariant> reader(Class<BuiltInVariant> clazz) {
+        BuiltInVariant[] variants = clazz.getEnumConstants();
+        if (variants == null) {
+            throw new IllegalArgumentException("Class is not an enum");
+        }
+        return context -> {
+            for (BuiltInVariant variant : variants) {
+                if (((BuiltIn<?>) variant).javaIdentifier().equals(context.id())) {
+                    return variant;
+                }
+            }
+            return null;
+        };
+    }
+
+    /**
+     * Should be implemented on an enum within the entity class. The enum lists vanilla variants that can appear on bedrock, in the order of their bedrock network ID.
+     *
+     * <p>The enum constants should be named the same as their Java identifiers.</p>
+     *
+     * @param <Variant> the same as the parent entity class. Used for type checking.
+     */
+    interface BuiltIn<Variant> {
+
+        String name();
+
+        int ordinal();
+
+        default Key javaIdentifier() {
+            return MinecraftKey.key(name().toLowerCase(Locale.ROOT));
+        }
+    }
+}

core/src/main/java/org/geysermc/geyser/entity/type/living/animal/farm/TemperatureVariantAnimal.java

@@ -31,6 +31,7 @@ import org.cloudburstmc.math.vector.Vector3f;
 import org.cloudburstmc.protocol.bedrock.data.entity.EntityDataTypes;
 import org.geysermc.geyser.entity.EntityDefinition;
 import org.geysermc.geyser.entity.type.living.animal.AnimalEntity;
+import org.geysermc.geyser.entity.type.living.animal.VariantHolder;
 import org.geysermc.geyser.session.GeyserSession;
 import org.geysermc.geyser.session.cache.registry.JavaRegistryKey;
 import org.geysermc.geyser.session.cache.registry.RegistryEntryContext;
@@ -43,7 +44,7 @@ import java.util.Locale;
 import java.util.UUID;
 import java.util.function.Function;
 
-public abstract class TemperatureVariantAnimal<Variant> extends AnimalEntity {
+public abstract class TemperatureVariantAnimal<Variant> extends AnimalEntity implements VariantHolder<Variant> {
 
     public TemperatureVariantAnimal(GeyserSession session, int entityId, long geyserId, UUID uuid, EntityDefinition<?> definition,
                                     Vector3f position, Vector3f motion, float yaw, float pitch, float headYaw) {
@@ -68,7 +69,7 @@ public abstract class TemperatureVariantAnimal<Variant> extends AnimalEntity {
 
     // Ordered by bedrock id
     // TODO: are these ordered correctly? Does the order differ for mobs?
-    public enum BuiltInVariant {
+    public enum BuiltInVariant implements BuiltIn<?> {
         COLD,
         TEMPERATE,
         WARM;

core/src/main/java/org/geysermc/geyser/entity/type/living/animal/tameable/CatEntity.java

@@ -25,33 +25,31 @@
 
 package org.geysermc.geyser.entity.type.living.animal.tameable;
 
-import net.kyori.adventure.key.Key;
 import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.cloudburstmc.math.vector.Vector3f;
 import org.cloudburstmc.protocol.bedrock.data.entity.EntityDataTypes;
 import org.cloudburstmc.protocol.bedrock.data.entity.EntityFlag;
 import org.geysermc.geyser.entity.EntityDefinition;
+import org.geysermc.geyser.entity.type.living.animal.VariantHolder;
 import org.geysermc.geyser.inventory.GeyserItemStack;
 import org.geysermc.geyser.item.type.Item;
 import org.geysermc.geyser.session.GeyserSession;
 import org.geysermc.geyser.session.cache.registry.JavaRegistries;
-import org.geysermc.geyser.session.cache.registry.RegistryEntryContext;
+import org.geysermc.geyser.session.cache.registry.JavaRegistryKey;
 import org.geysermc.geyser.session.cache.tags.ItemTag;
 import org.geysermc.geyser.session.cache.tags.Tag;
 import org.geysermc.geyser.util.InteractionResult;
 import org.geysermc.geyser.util.InteractiveTag;
-import org.geysermc.geyser.util.MinecraftKey;
 import org.geysermc.mcprotocollib.protocol.data.game.entity.metadata.type.BooleanEntityMetadata;
 import org.geysermc.mcprotocollib.protocol.data.game.entity.metadata.type.ByteEntityMetadata;
 import org.geysermc.mcprotocollib.protocol.data.game.entity.metadata.type.IntEntityMetadata;
 import org.geysermc.mcprotocollib.protocol.data.game.entity.player.Hand;
 
-import java.util.Locale;
 import java.util.UUID;
-import java.util.function.Function;
 
-public class CatEntity extends TameableEntity {
+// TODO this is implementing VariantHolder<Object> until MCPL updates
+public class CatEntity extends TameableEntity implements VariantHolder<Object> {
 
     private byte collarColor = 14; // Red - default
 
@@ -87,17 +85,19 @@ public class CatEntity extends TameableEntity {
         updateCollarColor();
     }
 
-    // TODO this is a holder when MCPL updates
-    // TODO also checks if this works
-    public void setCatVariant(IntEntityMetadata entityMetadata) {
-        // Different colors in Java and Bedrock for some reason
-        int metadataValue = entityMetadata.getPrimitiveValue();
+    @Override
+    public JavaRegistryKey<BuiltInVariant> variantRegistry() {
+        return JavaRegistries.CAT_VARIANT;
+    }
 
-        BuiltInVariant variant = JavaRegistries.CAT_VARIANT.fromNetworkId(session, metadataValue);
-        if (variant == null) {
-            variant = BuiltInVariant.BLACK; // Default variant on Java
-        }
-        dirtyMetadata.put(EntityDataTypes.VARIANT, variant.ordinal());
+    @Override
+    public void setBedrockVariant(int bedrockId) {
+        dirtyMetadata.put(EntityDataTypes.VARIANT, bedrockId);
+    }
+
+    @Override
+    public BuiltIn<Object> defaultVariant() {
+        return BuiltInVariant.BLACK; // Default variant on Java
     }
 
     public void setResting(BooleanEntityMetadata entityMetadata) {
@@ -147,8 +147,7 @@ public class CatEntity extends TameableEntity {
 
     // Ordered by bedrock id
     // TODO: are these ordered correctly?
-    // TODO lessen code duplication with other variant mobs
-    public enum BuiltInVariant {
+    public enum BuiltInVariant implements BuiltIn<Object> {
         WHITE,
         BLACK,
         RED,
@@ -159,23 +158,6 @@ public class CatEntity extends TameableEntity {
         RAGDOLL,
         TABBY,
         ALL_BLACK,
-        JELLIE;
-
-        public static final Function<RegistryEntryContext, BuiltInVariant> READER = context -> getByJavaIdentifier(context.id());
-
-        private final Key javaIdentifier;
-
-        BuiltInVariant() {
-            javaIdentifier = MinecraftKey.key(name().toLowerCase(Locale.ROOT));
-        }
-
-        public static @Nullable BuiltInVariant getByJavaIdentifier(Key identifier) {
-            for (BuiltInVariant variant : values()) {
-                if (variant.javaIdentifier.equals(identifier)) {
-                    return variant;
-                }
-            }
-            return null;
-        }
+        JELLIE
     }
 }

core/src/main/java/org/geysermc/geyser/entity/type/living/animal/tameable/WolfEntity.java

@@ -25,7 +25,6 @@
 
 package org.geysermc.geyser.entity.type.living.animal.tameable;
 
-import net.kyori.adventure.key.Key;
 import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.cloudburstmc.math.vector.Vector3f;
@@ -34,6 +33,7 @@ import org.cloudburstmc.protocol.bedrock.data.entity.EntityFlag;
 import org.cloudburstmc.protocol.bedrock.packet.AddEntityPacket;
 import org.cloudburstmc.protocol.bedrock.packet.UpdateAttributesPacket;
 import org.geysermc.geyser.entity.EntityDefinition;
+import org.geysermc.geyser.entity.type.living.animal.VariantHolder;
 import org.geysermc.geyser.inventory.GeyserItemStack;
 import org.geysermc.geyser.item.Items;
 import org.geysermc.geyser.item.enchantment.EnchantmentComponent;
@@ -41,18 +41,15 @@ import org.geysermc.geyser.item.type.DyeItem;
 import org.geysermc.geyser.item.type.Item;
 import org.geysermc.geyser.session.GeyserSession;
 import org.geysermc.geyser.session.cache.registry.JavaRegistries;
-import org.geysermc.geyser.session.cache.registry.RegistryEntryContext;
+import org.geysermc.geyser.session.cache.registry.JavaRegistryKey;
 import org.geysermc.geyser.session.cache.tags.ItemTag;
 import org.geysermc.geyser.session.cache.tags.Tag;
 import org.geysermc.geyser.util.InteractionResult;
 import org.geysermc.geyser.util.InteractiveTag;
 import org.geysermc.geyser.util.ItemUtils;
-import org.geysermc.geyser.util.MinecraftKey;
-import org.geysermc.mcprotocollib.protocol.data.game.Holder;
 import org.geysermc.mcprotocollib.protocol.data.game.entity.metadata.WolfVariant;
 import org.geysermc.mcprotocollib.protocol.data.game.entity.metadata.type.ByteEntityMetadata;
 import org.geysermc.mcprotocollib.protocol.data.game.entity.metadata.type.IntEntityMetadata;
-import org.geysermc.mcprotocollib.protocol.data.game.entity.metadata.type.ObjectEntityMetadata;
 import org.geysermc.mcprotocollib.protocol.data.game.entity.player.GameMode;
 import org.geysermc.mcprotocollib.protocol.data.game.entity.player.Hand;
 import org.geysermc.mcprotocollib.protocol.data.game.item.ItemStack;
@@ -60,11 +57,9 @@ import org.geysermc.mcprotocollib.protocol.data.game.item.component.DataComponen
 import org.geysermc.mcprotocollib.protocol.data.game.item.component.HolderSet;
 
 import java.util.Collections;
-import java.util.Locale;
 import java.util.UUID;
-import java.util.function.Function;
 
-public class WolfEntity extends TameableEntity {
+public class WolfEntity extends TameableEntity implements VariantHolder<WolfVariant> {
     private byte collarColor = 14; // Red - default
     private HolderSet repairableItems = null;
     private boolean isCurseOfBinding = false;
@@ -119,16 +114,19 @@ public class WolfEntity extends TameableEntity {
         dirtyMetadata.put(EntityDataTypes.COLOR, time != 0 ? (byte) 0 : collarColor);
     }
 
-    // 1.20.5+
-    public void setWolfVariant(ObjectEntityMetadata<Holder<WolfVariant>> entityMetadata) {
-        // TODO set to pale if custom holder?
-        entityMetadata.getValue().ifId(id -> {
-            BuiltInVariant wolfVariant = JavaRegistries.WOLF_VARIANT.fromNetworkId(session, id);
-            if (wolfVariant == null) {
-                wolfVariant = BuiltInVariant.PALE;
-            }
-            dirtyMetadata.put(EntityDataTypes.VARIANT, wolfVariant.ordinal());
-        });
+    @Override
+    public JavaRegistryKey<BuiltInVariant> variantRegistry() {
+        return JavaRegistries.WOLF_VARIANT;
+    }
+
+    @Override
+    public void setBedrockVariant(int bedrockId) {
+        dirtyMetadata.put(EntityDataTypes.VARIANT, bedrockId);
+    }
+
+    @Override
+    public BuiltIn<WolfVariant> defaultVariant() {
+        return BuiltInVariant.PALE;
     }
 
     @Override
@@ -201,7 +199,7 @@ public class WolfEntity extends TameableEntity {
     }
 
     // Ordered by bedrock id
-    public enum BuiltInVariant {
+    public enum BuiltInVariant implements BuiltIn<WolfVariant> {
         PALE,
         ASHEN,
         BLACK,
@@ -210,23 +208,6 @@ public class WolfEntity extends TameableEntity {
         SNOWY,
         SPOTTED,
         STRIPED,
-        WOODS;
-
-        public static final Function<RegistryEntryContext, BuiltInVariant> READER = context -> getByJavaIdentifier(context.id());
-
-        private final Key javaIdentifier;
-
-        BuiltInVariant() {
-            this.javaIdentifier = MinecraftKey.key(this.name().toLowerCase(Locale.ROOT));
-        }
-
-        public static @Nullable BuiltInVariant getByJavaIdentifier(Key identifier) {
-            for (BuiltInVariant wolfVariant : values()) {
-                if (wolfVariant.javaIdentifier.equals(identifier)) {
-                    return wolfVariant;
-                }
-            }
-            return null;
-        }
+        WOODS
     }
 }

core/src/main/java/org/geysermc/geyser/session/cache/RegistryCache.java

@@ -39,6 +39,7 @@ import org.cloudburstmc.protocol.bedrock.data.TrimMaterial;
 import org.cloudburstmc.protocol.bedrock.data.TrimPattern;
 import org.geysermc.geyser.GeyserImpl;
 import org.geysermc.geyser.entity.type.living.animal.FrogEntity;
+import org.geysermc.geyser.entity.type.living.animal.VariantHolder;
 import org.geysermc.geyser.entity.type.living.animal.farm.TemperatureVariantAnimal;
 import org.geysermc.geyser.entity.type.living.animal.tameable.CatEntity;
 import org.geysermc.geyser.entity.type.living.animal.tameable.WolfEntity;
@@ -96,9 +97,9 @@ public final class RegistryCache {
         register("worldgen/biome", (cache, array) -> cache.biomeTranslations = array, BiomeTranslator::loadServerBiome);
         register("banner_pattern", cache -> cache.bannerPatterns, context -> BannerPattern.getByJavaIdentifier(context.id()));
 
-        register(JavaRegistries.CAT_VARIANT, cache -> cache.catVariants, CatEntity.BuiltInVariant.READER);
-        register(JavaRegistries.FROG_VARIANT, cache -> cache.frogVariants, FrogEntity.BuiltInVariant.READER);
-        register(JavaRegistries.WOLF_VARIANT, cache -> cache.wolfVariants, WolfEntity.BuiltInVariant.READER);
+        register(JavaRegistries.CAT_VARIANT, cache -> cache.catVariants, VariantHolder.reader(CatEntity.BuiltInVariant.class));
+        register(JavaRegistries.FROG_VARIANT, cache -> cache.frogVariants, VariantHolder.reader(FrogEntity.BuiltInVariant.class));
+        register(JavaRegistries.WOLF_VARIANT, cache -> cache.wolfVariants, VariantHolder.reader(WolfEntity.BuiltInVariant.class));
 
         register(JavaRegistries.PIG_VARIANT, cache -> cache.pigVariants, TemperatureVariantAnimal.BuiltInVariant.READER);
         register(JavaRegistries.COW_VARIANT, cache -> cache.cowVariants, TemperatureVariantAnimal.BuiltInVariant.READER);