Store tracked entity data a little differently

api/src/main/java/com/viaversion/viaversion/api/data/entity/EntityTracker.java

@@ -24,9 +24,8 @@ package com.viaversion.viaversion.api.data.entity;
 
 import com.viaversion.viaversion.api.connection.UserConnection;
 import com.viaversion.viaversion.api.minecraft.entities.EntityType;
-import org.checkerframework.checker.nullness.qual.Nullable;
-
 import java.util.Map;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 public interface EntityTracker {
 
@@ -53,6 +52,14 @@ public interface EntityTracker {
      */
     boolean hasEntity(int id);
 
+    /**
+     * Returns the tracked entity for the given entity id if present.
+     *
+     * @param entityId entity id
+     * @return tracked entity if tracked
+     */
+    @Nullable TrackedEntity entity(int entityId);
+
     /**
      * Entity type of the entity if tracked.
      * This returning null does not necessarily mean no entity by the id exists.

api/src/main/java/com/viaversion/viaversion/api/data/entity/TrackedEntity.java

@@ -0,0 +1,58 @@
+/*
+ * This file is part of ViaVersion - https://github.com/ViaVersion/ViaVersion
+ * Copyright (C) 2016-2023 ViaVersion and contributors
+ *
+ * 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.
+ */
+package com.viaversion.viaversion.api.data.entity;
+
+import com.viaversion.viaversion.api.minecraft.entities.EntityType;
+
+public interface TrackedEntity {
+
+    /**
+     * Returns the type of the stored entity.
+     *
+     * @return type of the entity
+     */
+    EntityType entityType();
+
+    /**
+     * Object to hold arbitrary additional data.
+     *
+     * @return entity data
+     */
+    StoredEntityData data();
+
+    /**
+     * Returns whether the stored entity currently has any additional data.
+     *
+     * @return whether the stored entity currently has additional data
+     */
+    boolean hasData();
+
+    /**
+     * Returns whether metadata has already been sent at least once for this entity.
+     *
+     * @return whether metadata has already been sent at least once for this entity
+     */
+    boolean hasSentMetadata();
+
+    void sentMetadata(boolean sentMetadata);
+}

common/src/main/java/com/viaversion/viaversion/data/entity/EntityTrackerBase.java

@@ -23,18 +23,17 @@ import com.viaversion.viaversion.api.connection.UserConnection;
 import com.viaversion.viaversion.api.data.entity.ClientEntityIdChangeListener;
 import com.viaversion.viaversion.api.data.entity.DimensionData;
 import com.viaversion.viaversion.api.data.entity.EntityTracker;
+import com.viaversion.viaversion.api.data.entity.TrackedEntity;
 import com.viaversion.viaversion.api.data.entity.StoredEntityData;
 import com.viaversion.viaversion.api.minecraft.entities.EntityType;
 import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
+import java.util.Collections;
+import java.util.Map;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import space.vectrix.flare.fastutil.Int2ObjectSyncMap;
 
-import java.util.Collections;
-import java.util.Map;
-
 public class EntityTrackerBase implements EntityTracker, ClientEntityIdChangeListener {
-    private final Int2ObjectMap<EntityType> entityTypes = Int2ObjectSyncMap.hashmap();
+    private final Int2ObjectMap<TrackedEntity> entities = Int2ObjectSyncMap.hashmap();
-    private final Int2ObjectMap<StoredEntityData> entityData;
     private final UserConnection connection;
     private final EntityType playerType;
     private int clientEntityId = -1;
@@ -45,13 +44,8 @@ public class EntityTrackerBase implements EntityTracker, ClientEntityIdChangeLis
     private Map<String, DimensionData> dimensions = Collections.emptyMap();
 
     public EntityTrackerBase(UserConnection connection, @Nullable EntityType playerType) {
-        this(connection, playerType, false);
-    }
-
-    public EntityTrackerBase(UserConnection connection, @Nullable EntityType playerType, boolean storesEntityData) {
         this.connection = connection;
         this.playerType = playerType;
-        this.entityData = storesEntityData ? Int2ObjectSyncMap.hashmap() : null;
     }
 
     @Override
@@ -61,47 +55,46 @@ public class EntityTrackerBase implements EntityTracker, ClientEntityIdChangeLis
 
     @Override
     public void addEntity(int id, EntityType type) {
-        entityTypes.put(id, type);
+        entities.put(id, new TrackedEntityImpl(type));
     }
 
     @Override
     public boolean hasEntity(int id) {
-        return entityTypes.containsKey(id);
+        return entities.containsKey(id);
+    }
+
+    @Override
+    public @Nullable TrackedEntity entity(final int entityId) {
+        return entities.get(entityId);
     }
 
     @Override
     public @Nullable EntityType entityType(int id) {
-        return entityTypes.get(id);
+        final TrackedEntity entity = entities.get(id);
+        return entity != null ? entity.entityType() : null;
     }
 
     @Override
     public @Nullable StoredEntityData entityData(int id) {
-        Preconditions.checkArgument(entityData != null, "Entity data storage has to be explicitly enabled via the constructor");
+        final TrackedEntity entity = entities.get(id);
-        EntityType type = entityType(id);
+        return entity != null ? entity.data() : null;
-        return type != null ? entityData.computeIfAbsent(id, s -> new StoredEntityImpl(type)) : null;
     }
 
     @Override
     public @Nullable StoredEntityData entityDataIfPresent(int id) {
-        Preconditions.checkArgument(entityData != null, "Entity data storage has to be explicitly enabled via the constructor");
+        final TrackedEntity entity = entities.get(id);
-        return entityData.get(id);
+        return entity != null && entity.hasData() ? entity.data() : null;
     }
 
     //TODO Soft memory leak: Remove entities on respawn in protocols prior to 1.18 (1.16+ only when the worldname is different)
     @Override
     public void removeEntity(int id) {
-        entityTypes.remove(id);
+        entities.remove(id);
-        if (entityData != null) {
-            entityData.remove(id);
-        }
     }
 
     @Override
     public void clearEntities() {
-        entityTypes.clear();
+        entities.clear();
-        if (entityData != null) {
-            entityData.clear();
-        }
     }
 
     @Override
@@ -112,12 +105,11 @@ public class EntityTrackerBase implements EntityTracker, ClientEntityIdChangeLis
     @Override
     public void setClientEntityId(int clientEntityId) {
         Preconditions.checkNotNull(playerType);
-        entityTypes.put(clientEntityId, playerType);
+        final TrackedEntity oldEntity;
-        if (this.clientEntityId != -1 && entityData != null) {
+        if (this.clientEntityId != -1 && (oldEntity = entities.remove(this.clientEntityId)) != null) {
-            StoredEntityData data = entityData.remove(this.clientEntityId);
+            entities.put(clientEntityId, oldEntity);
-            if (data != null) {
+        } else {
-                entityData.put(clientEntityId, data);
+            entities.put(clientEntityId, new TrackedEntityImpl(playerType));
-            }
         }
 
         this.clientEntityId = clientEntityId;
@@ -126,7 +118,7 @@ public class EntityTrackerBase implements EntityTracker, ClientEntityIdChangeLis
     @Override
     public boolean trackClientEntity() {
         if (clientEntityId != -1) {
-            entityTypes.put(clientEntityId, playerType);
+            entities.put(clientEntityId, new TrackedEntityImpl(playerType));
             return true;
         }
         return false;

common/src/main/java/com/viaversion/viaversion/data/entity/StoredEntityDataImpl.java

@@ -24,11 +24,11 @@ import org.checkerframework.checker.nullness.qual.Nullable;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
-public final class StoredEntityImpl implements StoredEntityData {
+public final class StoredEntityDataImpl implements StoredEntityData {
     private final Map<Class<?>, Object> storedObjects = new ConcurrentHashMap<>();
     private final EntityType type;
 
-    public StoredEntityImpl(EntityType type) {
+    public StoredEntityDataImpl(EntityType type) {
         this.type = type;
     }
 
@@ -39,11 +39,13 @@ public final class StoredEntityImpl implements StoredEntityData {
 
     @Override
     public @Nullable <T> T get(Class<T> objectClass) {
+        //noinspection unchecked
         return (T) storedObjects.get(objectClass);
     }
 
     @Override
     public <T> @Nullable T remove(Class<T> objectClass) {
+        //noinspection unchecked
         return (T) storedObjects.remove(objectClass);
     }
 

common/src/main/java/com/viaversion/viaversion/data/entity/TrackedEntityImpl.java

@@ -0,0 +1,60 @@
+/*
+ * This file is part of ViaVersion - https://github.com/ViaVersion/ViaVersion
+ * Copyright (C) 2023 ViaVersion and contributors
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+package com.viaversion.viaversion.data.entity;
+
+import com.viaversion.viaversion.api.data.entity.TrackedEntity;
+import com.viaversion.viaversion.api.data.entity.StoredEntityData;
+import com.viaversion.viaversion.api.minecraft.entities.EntityType;
+
+public final class TrackedEntityImpl implements TrackedEntity {
+    private final EntityType entityType;
+    private StoredEntityData data;
+    private boolean sentMetadata;
+
+    public TrackedEntityImpl(final EntityType entityType) {
+        this.entityType = entityType;
+    }
+
+    @Override
+    public EntityType entityType() {
+        return entityType;
+    }
+
+    @Override
+    public StoredEntityData data() {
+        if (data == null) {
+            data = new StoredEntityDataImpl(entityType);
+        }
+        return data;
+    }
+
+    @Override
+    public boolean hasData() {
+        return data != null;
+    }
+
+    @Override
+    public boolean hasSentMetadata() {
+        return sentMetadata;
+    }
+
+    @Override
+    public void sentMetadata(final boolean sentMetadata) {
+        this.sentMetadata = sentMetadata;
+    }
+}

common/src/main/java/com/viaversion/viaversion/protocols/protocol1_19_4to1_19_3/ClientboundPackets1_19_4.java

@@ -44,93 +44,93 @@ public enum ClientboundPackets1_19_4 implements ClientboundPacketType {
     COOLDOWN, // 0x14
     CUSTOM_CHAT_COMPLETIONS, // 0x15
     PLUGIN_MESSAGE, // 0x16
-    DAMAGE_EVENT,
+    DAMAGE_EVENT, // 0x17
-    DELETE_CHAT_MESSAGE, // 0x17
+    DELETE_CHAT_MESSAGE, // 0x18
-    DISCONNECT, // 0x18
+    DISCONNECT, // 0x19
-    DISGUISED_CHAT, // 0x19
+    DISGUISED_CHAT, // 0x1A
-    ENTITY_STATUS, // 0x1A
+    ENTITY_STATUS, // 0x1B
-    EXPLOSION, // 0x1B
+    EXPLOSION, // 0x1C
-    UNLOAD_CHUNK, // 0x1C
+    UNLOAD_CHUNK, // 0x1D
-    GAME_EVENT, // 0x1D
+    GAME_EVENT, // 0x1E
-    OPEN_HORSE_WINDOW, // 0x1E
+    OPEN_HORSE_WINDOW, // 0x1F
-    HIT_ANIMATION, // 0x1F
+    HIT_ANIMATION, // 0x20
-    WORLD_BORDER_INIT, // 0x20
+    WORLD_BORDER_INIT, // 0x21
-    KEEP_ALIVE, // 0x21
+    KEEP_ALIVE, // 0x22
-    CHUNK_DATA, // 0x22
+    CHUNK_DATA, // 0x23
-    EFFECT, // 0x23
+    EFFECT, // 0x24
-    SPAWN_PARTICLE, // 0x24
+    SPAWN_PARTICLE, // 0x25
-    UPDATE_LIGHT, // 0x25
+    UPDATE_LIGHT, // 0x26
-    JOIN_GAME, // 0x26
+    JOIN_GAME, // 0x27
-    MAP_DATA, // 0x27
+    MAP_DATA, // 0x28
-    TRADE_LIST, // 0x28
+    TRADE_LIST, // 0x29
-    ENTITY_POSITION, // 0x29
+    ENTITY_POSITION, // 0x2A
-    ENTITY_POSITION_AND_ROTATION, // 0x2A
+    ENTITY_POSITION_AND_ROTATION, // 0x2B
-    ENTITY_ROTATION, // 0x2B
+    ENTITY_ROTATION, // 0x2C
-    VEHICLE_MOVE, // 0x2C
+    VEHICLE_MOVE, // 0x2D
-    OPEN_BOOK, // 0x2D
+    OPEN_BOOK, // 0x2E
-    OPEN_WINDOW, // 0x2E
+    OPEN_WINDOW, // 0x2F
-    OPEN_SIGN_EDITOR, // 0x2F
+    OPEN_SIGN_EDITOR, // 0x30
-    PING, // 0x30
+    PING, // 0x31
-    CRAFT_RECIPE_RESPONSE, // 0x31
+    CRAFT_RECIPE_RESPONSE, // 0x32
-    PLAYER_ABILITIES, // 0x32
+    PLAYER_ABILITIES, // 0x33
-    PLAYER_CHAT, // 0x33
+    PLAYER_CHAT, // 0x34
-    COMBAT_END, // 0x34
+    COMBAT_END, // 0x35
-    COMBAT_ENTER, // 0x35
+    COMBAT_ENTER, // 0x36
-    COMBAT_KILL, // 0x36
+    COMBAT_KILL, // 0x37
-    PLAYER_INFO_REMOVE, // 0x37
+    PLAYER_INFO_REMOVE, // 0x38
-    PLAYER_INFO_UPDATE, // 0x38
+    PLAYER_INFO_UPDATE, // 0x39
-    FACE_PLAYER, // 0x39
+    FACE_PLAYER, // 0x3A
-    PLAYER_POSITION, // 0x3A
+    PLAYER_POSITION, // 0x3B
-    UNLOCK_RECIPES, // 0x3B
+    UNLOCK_RECIPES, // 0x3C
-    REMOVE_ENTITIES, // 0x3C
+    REMOVE_ENTITIES, // 0x3D
-    REMOVE_ENTITY_EFFECT, // 0x3D
+    REMOVE_ENTITY_EFFECT, // 0x3E
-    RESOURCE_PACK, // 0x3E
+    RESOURCE_PACK, // 0x3F
-    RESPAWN, // 0x3F
+    RESPAWN, // 0x40
-    ENTITY_HEAD_LOOK, // 0x40
+    ENTITY_HEAD_LOOK, // 0x41
-    MULTI_BLOCK_CHANGE, // 0x41
+    MULTI_BLOCK_CHANGE, // 0x42
-    SELECT_ADVANCEMENTS_TAB, // 0x42
+    SELECT_ADVANCEMENTS_TAB, // 0x43
-    SERVER_DATA, // 0x43
+    SERVER_DATA, // 0x44
-    ACTIONBAR, // 0x44
+    ACTIONBAR, // 0x45
-    WORLD_BORDER_CENTER, // 0x45
+    WORLD_BORDER_CENTER, // 0x46
-    WORLD_BORDER_LERP_SIZE, // 0x46
+    WORLD_BORDER_LERP_SIZE, // 0x47
-    WORLD_BORDER_SIZE, // 0x47
+    WORLD_BORDER_SIZE, // 0x48
-    WORLD_BORDER_WARNING_DELAY, // 0x48
+    WORLD_BORDER_WARNING_DELAY, // 0x49
-    WORLD_BORDER_WARNING_DISTANCE, // 0x49
+    WORLD_BORDER_WARNING_DISTANCE, // 0x4A
-    CAMERA, // 0x4A
+    CAMERA, // 0x4B
-    HELD_ITEM_CHANGE, // 0x4B
+    HELD_ITEM_CHANGE, // 0x4C
-    UPDATE_VIEW_POSITION, // 0x4C
+    UPDATE_VIEW_POSITION, // 0x4D
-    UPDATE_VIEW_DISTANCE, // 0x4D
+    UPDATE_VIEW_DISTANCE, // 0x4E
-    SPAWN_POSITION, // 0x4E
+    SPAWN_POSITION, // 0x4F
-    DISPLAY_SCOREBOARD, // 0x4F
+    DISPLAY_SCOREBOARD, // 0x50
-    ENTITY_METADATA, // 0x50
+    ENTITY_METADATA, // 0x51
-    ATTACH_ENTITY, // 0x51
+    ATTACH_ENTITY, // 0x52
-    ENTITY_VELOCITY, // 0x52
+    ENTITY_VELOCITY, // 0x53
-    ENTITY_EQUIPMENT, // 0x53
+    ENTITY_EQUIPMENT, // 0x54
-    SET_EXPERIENCE, // 0x54
+    SET_EXPERIENCE, // 0x55
-    UPDATE_HEALTH, // 0x55
+    UPDATE_HEALTH, // 0x56
-    SCOREBOARD_OBJECTIVE, // 0x56
+    SCOREBOARD_OBJECTIVE, // 0x57
-    SET_PASSENGERS, // 0x57
+    SET_PASSENGERS, // 0x58
-    TEAMS, // 0x58
+    TEAMS, // 0x59
-    UPDATE_SCORE, // 0x59
+    UPDATE_SCORE, // 0x5A
-    SET_SIMULATION_DISTANCE, // 0x5A
+    SET_SIMULATION_DISTANCE, // 0x5B
-    TITLE_SUBTITLE, // 0x5B
+    TITLE_SUBTITLE, // 0x5C
-    TIME_UPDATE, // 0x5C
+    TIME_UPDATE, // 0x5D
-    TITLE_TEXT, // 0x5D
+    TITLE_TEXT, // 0x5E
-    TITLE_TIMES, // 0x5E
+    TITLE_TIMES, // 0x5F
-    ENTITY_SOUND, // 0x5F
+    ENTITY_SOUND, // 0x60
-    SOUND, // 0x60
+    SOUND, // 0x61
-    STOP_SOUND, // 0x61
+    STOP_SOUND, // 0x62
-    SYSTEM_CHAT, // 0x62
+    SYSTEM_CHAT, // 0x63
-    TAB_LIST, // 0x63
+    TAB_LIST, // 0x64
-    NBT_QUERY, // 0x64
+    NBT_QUERY, // 0x65
-    COLLECT_ITEM, // 0x65
+    COLLECT_ITEM, // 0x66
-    ENTITY_TELEPORT, // 0x66
+    ENTITY_TELEPORT, // 0x67
-    ADVANCEMENTS, // 0x67
+    ADVANCEMENTS, // 0x68
-    ENTITY_PROPERTIES, // 0x68
+    ENTITY_PROPERTIES, // 0x69
-    UPDATE_ENABLED_FEATURES, // 0x69
+    UPDATE_ENABLED_FEATURES, // 0x6A
-    ENTITY_EFFECT, // 0x6A
+    ENTITY_EFFECT, // 0x6B
-    DECLARE_RECIPES, // 0x6B
+    DECLARE_RECIPES, // 0x6C
-    TAGS; // 0x6C
+    TAGS; // 0x6D
 
     @Override
     public int getId() {

common/src/main/java/com/viaversion/viaversion/rewriter/EntityRewriter.java

@@ -29,6 +29,7 @@ import com.viaversion.viaversion.api.data.Mappings;
 import com.viaversion.viaversion.api.data.ParticleMappings;
 import com.viaversion.viaversion.api.data.entity.DimensionData;
 import com.viaversion.viaversion.api.data.entity.EntityTracker;
+import com.viaversion.viaversion.api.data.entity.TrackedEntity;
 import com.viaversion.viaversion.api.minecraft.entities.EntityType;
 import com.viaversion.viaversion.api.minecraft.item.Item;
 import com.viaversion.viaversion.api.minecraft.metadata.MetaType;
@@ -102,7 +103,8 @@ public abstract class EntityRewriter<C extends ClientboundPacketType, T extends
 
     @Override
     public void handleMetadata(int entityId, List<Metadata> metadataList, UserConnection connection) {
-        EntityType type = tracker(connection).entityType(entityId);
+        final TrackedEntity entity = tracker(connection).entity(entityId);
+        final EntityType type = entity != null ? entity.entityType() : null;
         int i = 0; // Count index for fast removal
         for (Metadata metadata : metadataList.toArray(EMPTY_ARRAY)) { // Copy the list to allow mutation
             // Call handlers implementing the old handleMetadata
@@ -142,6 +144,10 @@ public abstract class EntityRewriter<C extends ClientboundPacketType, T extends
             }
             i++;
         }
+
+        if (entity != null) {
+            entity.sentMetadata(true);
+        }
     }
 
     @Deprecated