Make entity data type id mapping less manual

2026-04-06 13:34:19 +02:00
1 parent 31f257fee1
commit 4dc57e3d53
@@ -19,6 +19,8 @@ classes/
 *.ipr
 *.iws

+benchmark/
+
 # VSCode

 .settings/
@@ -43,6 +43,10 @@ public abstract class AbstractEntityDataTypes implements EntityDataTypes {
        return values;
    }

+    public Class<? extends EntityDataType> getDataTypeClass() {
+        return byId(0).getClass();
+    }
+
    protected EntityDataType add(final int typeId, final Type<?> type) {
        if (values[typeId] != null) {
            throw new IllegalArgumentException("Entity data type ID " + typeId + " is already registered as " + values[typeId]);
@@ -41,15 +41,13 @@ final class EntityPacketRewriter99_1 extends EntityRewriter<ClientboundPacket26_
    @Override
    protected void registerRewrites() {
        final EntityDataTypes26_1 entityDataTypes = protocol.mappedTypes().entityDataTypes();
-        filter().mapDataType(entityDataTypes::byId);
-        /* ... or something this if entity data classes changed
-        filter().mapDataType(typeId -> {
-            int id = typeId;
-            if (id >= SomeAddedIndex) {
-                id++;
-            }
-            return entityDataTypes.byId(id);
-        });*/
+        dataTypeMapper().register();
+        /* ... or like this for additions and removals that are not at the very end
+        dataTypeMapper()
+            .added(entityDataTypes.catSoundVariant)
+            .removed(entityDataTypes.cowSoundVariant)
+            .skip(entityDataTypes.pigSoundVariant) // if neither removed nor added, but the value type has to be changed separately
+            .register();*/

        // Registers registry type id changes
        registerEntityDataTypeHandler(
@@ -26,7 +26,6 @@ import com.viaversion.viaversion.api.minecraft.PaintingVariant;
 import com.viaversion.viaversion.api.minecraft.RegistryEntry;
 import com.viaversion.viaversion.api.minecraft.entities.EntityType;
 import com.viaversion.viaversion.api.minecraft.entities.EntityTypes1_20_5;
-import com.viaversion.viaversion.api.minecraft.entitydata.EntityDataType;
 import com.viaversion.viaversion.api.protocol.packet.PacketWrapper;
 import com.viaversion.viaversion.api.type.Types;
 import com.viaversion.viaversion.api.type.types.version.VersionedTypes;
@@ -177,17 +176,17 @@ public final class EntityPacketRewriter1_21 extends EntityRewriter<ClientboundPa

    @Override
    protected void registerRewrites() {
-        filter().handler((event, data) -> {
-            final EntityDataType type = data.dataType();
-            if (type == VersionedTypes.V1_20_5.entityDataTypes.wolfVariantType) {
-                final int variant = data.value();
-                data.setTypeAndValue(VersionedTypes.V1_21.entityDataTypes.wolfVariantType, Holder.of(variant));
-            } else if (type == VersionedTypes.V1_20_5.entityDataTypes.paintingVariantType) {
-                final int variant = data.value();
-                data.setTypeAndValue(VersionedTypes.V1_21.entityDataTypes.paintingVariantType, Holder.of(variant));
-            } else {
-                data.setDataType(VersionedTypes.V1_21.entityDataTypes.byId(type.typeId()));
-            }
+        dataTypeMapper()
+            .skip(VersionedTypes.V1_20_5.entityDataTypes.wolfVariantType)
+            .skip(VersionedTypes.V1_20_5.entityDataTypes.paintingVariantType)
+            .register();
+        filter().dataType(VersionedTypes.V1_20_5.entityDataTypes.paintingVariantType).handler((event, data) -> {
+            final int variant = data.value();
+            data.setTypeAndValue(VersionedTypes.V1_21.entityDataTypes.paintingVariantType, Holder.of(variant));
+        });
+        filter().dataType(VersionedTypes.V1_20_5.entityDataTypes.wolfVariantType).handler((event, data) -> {
+            final int variant = data.value();
+            data.setTypeAndValue(VersionedTypes.V1_21.entityDataTypes.wolfVariantType, Holder.of(variant));
        });
        registerEntityDataTypeHandler(
            VersionedTypes.V1_21.entityDataTypes.itemType,
@@ -80,22 +80,12 @@ public final class EntityPacketRewriter26_1 extends EntityRewriter<ClientboundPa
    @Override
    protected void registerRewrites() {
        final EntityDataTypes26_1 entityDataTypes = protocol.mappedTypes().entityDataTypes();
-        filter().mapDataType(id -> {
-            int mappedId = id;
-            if (mappedId >= entityDataTypes.catSoundVariant.typeId()) {
-                mappedId++;
-            }
-            if (mappedId >= entityDataTypes.cowSoundVariant.typeId()) {
-                mappedId++;
-            }
-            if (mappedId >= entityDataTypes.pigSoundVariant.typeId()) {
-                mappedId++;
-            }
-            if (mappedId >= entityDataTypes.chickenSoundVariant.typeId()) {
-                mappedId++;
-            }
-            return entityDataTypes.byId(mappedId);
-        });
+        dataTypeMapper()
+            .added(entityDataTypes.catSoundVariant)
+            .added(entityDataTypes.cowSoundVariant)
+            .added(entityDataTypes.pigSoundVariant)
+            .added(entityDataTypes.chickenSoundVariant)
+            .register();
        registerEntityDataTypeHandler(
            entityDataTypes.itemType,
            entityDataTypes.blockStateType,
@@ -23,12 +23,10 @@ import com.viaversion.viaversion.api.minecraft.entities.EntityType;
 import com.viaversion.viaversion.api.minecraft.entities.EntityTypes1_21_4;
 import com.viaversion.viaversion.api.type.Types;
 import com.viaversion.viaversion.api.type.types.version.VersionedTypes;
-import com.viaversion.viaversion.protocols.v1_20_5to1_21.packet.ClientboundConfigurationPackets1_21;
 import com.viaversion.viaversion.protocols.v1_21_2to1_21_4.Protocol1_21_2To1_21_4;
 import com.viaversion.viaversion.protocols.v1_21_2to1_21_4.packet.ServerboundPackets1_21_4;
 import com.viaversion.viaversion.protocols.v1_21to1_21_2.packet.ClientboundPacket1_21_2;
 import com.viaversion.viaversion.rewriter.EntityRewriter;
-import com.viaversion.viaversion.rewriter.RegistryDataRewriter;

 public final class EntityPacketRewriter1_21_4 extends EntityRewriter<ClientboundPacket1_21_2, Protocol1_21_2To1_21_4> {

@@ -67,7 +65,7 @@ public final class EntityPacketRewriter1_21_4 extends EntityRewriter<Clientbound

    @Override
    protected void registerRewrites() {
-        filter().mapDataType(VersionedTypes.V1_21_4.entityDataTypes::byId);
+        dataTypeMapper().register();

        registerEntityDataTypeHandler(
            VersionedTypes.V1_21_4.entityDataTypes.itemType,
@@ -186,28 +186,16 @@ public final class EntityPacketRewriter1_21_5 extends EntityRewriter<Clientbound

    @Override
    protected void registerRewrites() {
-        filter().handler((event, data) -> {
-            final int id = data.dataType().typeId();
-            if (id == VersionedTypes.V1_21_4.entityDataTypes.wolfVariantType.typeId()) {
-                final Holder<WolfVariant> wolfVariant = data.value();
-                data.setTypeAndValue(VersionedTypes.V1_21_5.entityDataTypes.wolfVariantType, wolfVariant.hasId() ? wolfVariant.id() : 0);
-                return;
-            }
-
-            int mappedId = id;
-            if (mappedId >= VersionedTypes.V1_21_5.entityDataTypes.cowVariantType.typeId()) {
-                mappedId++;
-            }
-            if (mappedId >= VersionedTypes.V1_21_5.entityDataTypes.wolfSoundVariantType.typeId()) {
-                mappedId++;
-            }
-            if (mappedId >= VersionedTypes.V1_21_5.entityDataTypes.pigVariantType.typeId()) {
-                mappedId++;
-            }
-            if (mappedId >= VersionedTypes.V1_21_5.entityDataTypes.chickenVariantType.typeId()) {
-                mappedId++;
-            }
-            data.setDataType(VersionedTypes.V1_21_5.entityDataTypes.byId(mappedId));
+        dataTypeMapper()
+            .added(VersionedTypes.V1_21_5.entityDataTypes.cowVariantType)
+            .added(VersionedTypes.V1_21_5.entityDataTypes.wolfSoundVariantType)
+            .added(VersionedTypes.V1_21_5.entityDataTypes.pigVariantType)
+            .added(VersionedTypes.V1_21_5.entityDataTypes.chickenVariantType)
+            .skip(VersionedTypes.V1_21_4.entityDataTypes.wolfVariantType)
+            .register();
+        filter().dataType(VersionedTypes.V1_21_4.entityDataTypes.wolfVariantType).handler((event, data) -> {
+            final Holder<WolfVariant> wolfVariant = data.value();
+            data.setTypeAndValue(VersionedTypes.V1_21_5.entityDataTypes.wolfVariantType, wolfVariant.hasId() ? wolfVariant.id() : 0);
        });

        registerEntityDataTypeHandler(
@@ -101,7 +101,7 @@ public final class EntityPacketRewriter1_21_6 extends EntityRewriter<Clientbound
    @Override
    protected void registerRewrites() {
        final EntityDataTypes1_21_5 entityDataTypes = VersionedTypes.V1_21_6.entityDataTypes;
-        filter().mapDataType(entityDataTypes::byId);
+        dataTypeMapper().register();
        registerEntityDataTypeHandler(
            entityDataTypes.itemType,
            entityDataTypes.blockStateType,
@@ -103,24 +103,11 @@ public final class EntityPacketRewriter1_21_9 extends EntityRewriter<Clientbound
    protected void registerRewrites() {
        final EntityDataTypes1_21_5 unmappedEntityDataTypes = protocol.types().entityDataTypes();
        final EntityDataTypes1_21_9 entityDataTypes = protocol.mappedTypes().entityDataTypes();
-        filter().handler((event, data) -> {
-            int id = data.dataType().typeId();
-            if (id == unmappedEntityDataTypes.compoundTagType.typeId()) {
-                if (event.entityType() == null) {
-                    // Remove unhandled data from bad packets here
-                    event.cancel();
-                }
-                return; // Handled below
-            }
-            if (id > unmappedEntityDataTypes.compoundTagType.typeId()) {
-                id--;
-            }
-            if (id > entityDataTypes.armadilloState.typeId()) {
-                id += 2; // copper golem and weathering copper state
-            }
-            data.setDataType(entityDataTypes.byId(id));
-        });
-
+        dataTypeMapper()
+            .added(entityDataTypes.copperGolemState)
+            .added(entityDataTypes.weatheringCopperState)
+            .removed(unmappedEntityDataTypes.compoundTagType)
+            .register();
        registerEntityDataTypeHandler(
            entityDataTypes.itemType,
            entityDataTypes.blockStateType,
@@ -53,12 +53,7 @@ public final class EntityPacketRewriter1_21_11 extends EntityRewriter<Clientboun
    @Override
    protected void registerRewrites() {
        final EntityDataTypes1_21_11 entityDataTypes = protocol.mappedTypes().entityDataTypes();
-        filter().mapDataType(id -> {
-            if (id >= entityDataTypes.zombieNautilusVariantType.typeId()) {
-                id++;
-            }
-            return entityDataTypes.byId(id);
-        });
+        dataTypeMapper().added(entityDataTypes.zombieNautilusVariantType).register();
        registerEntityDataTypeHandler(
            entityDataTypes.itemType,
            entityDataTypes.blockStateType,
@@ -475,7 +475,7 @@ public final class EntityPacketRewriter1_21_2 extends EntityRewriter<Clientbound

    @Override
    protected void registerRewrites() {
-        filter().mapDataType(VersionedTypes.V1_21_2.entityDataTypes::byId);
+        dataTypeMapper().register();

        registerEntityDataTypeHandler(
            VersionedTypes.V1_21_2.entityDataTypes.itemType,
@@ -91,6 +91,19 @@ public abstract class EntityRewriter<C extends ClientboundPacketType, T extends
        return new EntityDataFilter.Builder(this);
    }

+    /**
+     * Registers a first-pass handler to remap entity data types via {@code rewriter.protocol().mappedTypes().entityDataTypes().byId(typeId)}
+     * <p>
+     * Removed source types are intentionally left untouched during the initial handler, so entity-specific handlers can
+     * still manually remap them before the final cancellation runs (unless something sent invalid packets for untracked entities).
+     *
+     * @return configurable type remapper
+     */
+    public EntityDataFilter.DataTypeMapper dataTypeMapper() {
+        Preconditions.checkNotNull(protocol.mappedTypes(), "Protocol does not override mappedTypes, use filter().mapDataType instead");
+        return new EntityDataFilter.DataTypeMapper(this);
+    }
+
    /**
     * Registers an entity data filter.
     * Note that {@link EntityDataFilter.Builder#register()} already calls this method.
@@ -99,7 +112,6 @@ public abstract class EntityRewriter<C extends ClientboundPacketType, T extends
     * @throws IllegalArgumentException if the filter is already registered
     */
    public void registerFilter(EntityDataFilter filter) {
-        Preconditions.checkArgument(!entityDataFilters.contains(filter));
        entityDataFilters.add(filter);
    }

@@ -22,8 +22,14 @@ import com.google.common.base.Preconditions;
 import com.viaversion.viaversion.api.minecraft.entities.EntityType;
 import com.viaversion.viaversion.api.minecraft.entitydata.EntityData;
 import com.viaversion.viaversion.api.minecraft.entitydata.EntityDataType;
+import com.viaversion.viaversion.api.minecraft.entitydata.types.AbstractEntityDataTypes;
 import com.viaversion.viaversion.rewriter.EntityRewriter;
 import it.unimi.dsi.fastutil.ints.Int2ObjectFunction;
+import it.unimi.dsi.fastutil.ints.IntArrayList;
+import it.unimi.dsi.fastutil.ints.IntArraySet;
+import it.unimi.dsi.fastutil.ints.IntComparators;
+import it.unimi.dsi.fastutil.ints.IntList;
+import it.unimi.dsi.fastutil.ints.IntSet;
 import java.util.Objects;
 import org.checkerframework.checker.nullness.qual.Nullable;

@@ -34,6 +40,10 @@ public record EntityDataFilter(@Nullable EntityType type, boolean filterFamily,
        Preconditions.checkNotNull(handler, "EntityDataHandler cannot be null");
    }

+    public EntityDataFilter(final EntityDataHandler handler) {
+        this(null, true, null, -1, handler);
+    }
+
    /**
     * Returns the entity data index to filter, or -1.
     *
@@ -82,7 +92,7 @@ public record EntityDataFilter(@Nullable EntityType type, boolean filterFamily,
    /**
     * Returns whether if the entity data should be handled by this filter.
     *
-     * @param type     entity type
+     * @param type       entity type
     * @param entityData entityData
     * @return whether the data should be filtered
     */
@@ -211,6 +221,7 @@ public record EntityDataFilter(@Nullable EntityType type, boolean filterFamily,
            register();
        }

+        // Mainly for pre-1.20.5 use
        public void mapDataType(Int2ObjectFunction<EntityDataType> updateFunction) {
            handler((event, data) -> {
                EntityDataType mappedType = updateFunction.apply(data.dataType().typeId());
@@ -293,7 +304,112 @@ public record EntityDataFilter(@Nullable EntityType type, boolean filterFamily,
         * @return created data filter
         */
        public EntityDataFilter build() {
+            Preconditions.checkNotNull(handler, "EntityDataHandler cannot be null");
            return new EntityDataFilter(type, filterFamily, dataType, index, handler);
        }
    }
+
+    public static final class DataTypeMapper {
+        private final IntList addedTypeIds = new IntArrayList();
+        private final IntList removedTypeIds = new IntArrayList();
+        private final IntSet skippedTypeIds = new IntArraySet();
+        private final EntityRewriter<?, ?> rewriter;
+        private final AbstractEntityDataTypes mappedDataTypes;
+
+        public DataTypeMapper(final EntityRewriter<?, ?> rewriter) {
+            this.rewriter = rewriter;
+            this.mappedDataTypes = rewriter.protocol().mappedTypes().entityDataTypes();
+        }
+
+        /**
+         * Marks a mapped entity data type as added.
+         *
+         * @param dataType added mapped entity data type
+         * @return self
+         */
+        public DataTypeMapper added(final EntityDataType dataType) {
+            Preconditions.checkArgument(dataType.getClass() == mappedDataTypes.getDataTypeClass());
+            this.addedTypeIds.add(dataType.typeId());
+            return this;
+        }
+
+        /**
+         * Marks an unmapped entity data type as removed.
+         * <p>
+         * Matching unmapped types are left untouched for the regular handlers to process and must be canceled or updated manually.
+         *
+         * @param dataType removed unmapped entity data type
+         * @return self
+         */
+        public DataTypeMapper removed(final EntityDataType dataType) {
+            Preconditions.checkArgument(dataType.getClass() == rewriter.protocol().types().entityDataTypes().getDataTypeClass());
+            this.removedTypeIds.add(dataType.typeId());
+            return this;
+        }
+
+        /**
+         * Skips automatic remapping for an unmapped entity data type so another handler can process it later.
+         *
+         * @param dataType unmapped entity data type to leave untouched
+         * @return self
+         */
+        public DataTypeMapper skip(final EntityDataType dataType) {
+            Preconditions.checkArgument(dataType.getClass() == rewriter.protocol().types().entityDataTypes().getDataTypeClass());
+            this.skippedTypeIds.add(dataType.typeId());
+            return this;
+        }
+
+        public void register() {
+            if (addedTypeIds.isEmpty() && removedTypeIds.isEmpty() && skippedTypeIds.isEmpty()) {
+                // Simple handler
+                rewriter.filter().handler((event, data) -> {
+                    final EntityDataType dataType = mappedDataTypes.byId(data.dataType().typeId());
+                    data.setDataType(dataType);
+                });
+                return;
+            }
+
+            this.addedTypeIds.sort(IntComparators.NATURAL_COMPARATOR);
+            this.removedTypeIds.sort(IntComparators.OPPOSITE_COMPARATOR);
+
+            rewriter.registerFilter(new EntityDataFilter((event, data) -> {
+                // Do not cancel removed elements yet so that filters can handle them (unless something sent an invalid packet for untracked entities)
+                final EntityDataType mappedType = mappedType(data.dataType());
+                if (mappedType != null) {
+                    data.setDataType(mappedType);
+                } else if (event.entityType() == null) {
+                    event.cancel();
+                }
+            }));
+        }
+
+        private @Nullable EntityDataType mappedType(final EntityDataType dataType) {
+            if (skippedTypeIds.contains(dataType.typeId()) || shouldCancel(dataType)) {
+                return null;
+            }
+
+            final int mappedId = mappedTypeId(dataType.typeId());
+            return mappedDataTypes.byId(mappedId);
+        }
+
+        private boolean shouldCancel(final EntityDataType dataType) {
+            return !removedTypeIds.isEmpty() && removedTypeIds.contains(dataType.typeId());
+        }
+
+        private int mappedTypeId(int typeId) {
+            // Starts with the highest removed
+            for (final int removedTypeId : removedTypeIds) {
+                if (typeId > removedTypeId) {
+                    typeId--;
+                }
+            }
+            // Starts with the first added
+            for (final int addedTypeId : addedTypeIds) {
+                if (typeId >= addedTypeId) {
+                    typeId++;
+                }
+            }
+            return typeId;
+        }
+    }
 }