Fix the client entity being untracked on world changes

Fixes #2894 and probably other things

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

@@ -159,4 +159,12 @@ public interface EntityTracker {
     @Nullable DimensionData dimensionData(String dimension);
 
     void setDimensions(Map<String, DimensionData> dimensions);
+
+    /**
+     * Adds the client player entity to the tracker.
+     * If the client entity has not been set yet, this will return false.
+     *
+     * @return whether the client has been tracked
+     */
+    boolean trackClientEntity();
 }

common/src/main/java/com/viaversion/viaversion/connection/UserConnectionImpl.java

@@ -126,6 +126,7 @@ public class UserConnectionImpl implements UserConnection {
             storedObjects.values().removeIf(StorableObject::clearOnServerSwitch);
             for (EntityTracker tracker : entityTrackers.values()) {
                 tracker.clearEntities();
+                tracker.trackClientEntity();
             }
         } else {
             storedObjects.clear();

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

@@ -123,6 +123,15 @@ public class EntityTrackerBase implements EntityTracker, ClientEntityIdChangeLis
         this.clientEntityId = clientEntityId;
     }
 
+    @Override
+    public boolean trackClientEntity() {
+        if (clientEntityId != -1) {
+            entityTypes.put(clientEntityId, playerType);
+            return true;
+        }
+        return false;
+    }
+
     @Override
     public int currentWorldSectionHeight() {
         return currentWorldSectionHeight;

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

@@ -401,7 +401,11 @@ public abstract class EntityRewriter<T extends Protocol> extends RewriterBase<T>
     public PacketHandler playerTrackerHandler() {
         return wrapper -> {
             final EntityTracker tracker = tracker(wrapper.user());
-            tracker.addEntity(wrapper.get(Type.INT, 0), tracker.playerType());
+            final int entityId = wrapper.get(Type.INT, 0);
+            if (tracker.clientEntityId() == -1) {
+                tracker.setClientEntityId(entityId);
+            }
+            tracker.addEntity(entityId, tracker.playerType());
         };
     }
 
@@ -435,6 +439,7 @@ public abstract class EntityRewriter<T extends Protocol> extends RewriterBase<T>
             String world = wrapper.get(Type.STRING, 0);
             if (tracker.currentWorld() != null && !tracker.currentWorld().equals(world)) {
                 tracker.clearEntities();
+                tracker.trackClientEntity();
             }
             tracker.setCurrentWorld(world);
         };
@@ -457,6 +462,7 @@ public abstract class EntityRewriter<T extends Protocol> extends RewriterBase<T>
             String world = wrapper.get(Type.STRING, 0);
             if (tracker.currentWorld() != null && !tracker.currentWorld().equals(world)) {
                 tracker.clearEntities();
+                tracker.trackClientEntity();
             }
             tracker.setCurrentWorld(world);
         };

common/src/test/java/com/viaversion/viaversion/common/entities/EntityTypesTest.java

@@ -40,6 +40,7 @@ public class EntityTypesTest {
         testArrayOrder(Entity1_16Types.values(), Entity1_16Types::getTypeFromId);
         testArrayOrder(Entity1_16_2Types.values(), Entity1_16_2Types::getTypeFromId);
         testArrayOrder(Entity1_17Types.values(), Entity1_17Types::getTypeFromId);
+        // Newer type enums are automatically filled using mappings
     }
 
     private void testArrayOrder(EntityType[] types, Function<Integer, EntityType> returnFunction) {