Refactor dimension switch handling across all protocols

common/src/main/java/com/viaversion/viabackwards/api/rewriters/EntityRewriter.java

@@ -127,7 +127,6 @@ public PacketHandler worldTrackerHandlerByKey() {
             String world = wrapper.get(Types.STRING, 1);
             if (tracker.currentWorld() != null && !tracker.currentWorld().equals(world)) {
                 tracker.clearEntities();
-                tracker.trackClientEntity();
             }
             tracker.setCurrentWorld(world);
         };

common/src/main/java/com/viaversion/viabackwards/api/rewriters/EntityRewriterBase.java

@@ -23,6 +23,7 @@
 import com.viaversion.viabackwards.api.entities.storage.EntityReplacement;
 import com.viaversion.viabackwards.api.entities.storage.WrappedEntityData;
 import com.viaversion.viaversion.api.connection.UserConnection;
+import com.viaversion.viaversion.api.data.entity.EntityTracker;
 import com.viaversion.viaversion.api.data.entity.StoredEntityData;
 import com.viaversion.viaversion.api.data.entity.TrackedEntity;
 import com.viaversion.viaversion.api.minecraft.ClientWorld;
@@ -238,15 +239,27 @@ protected PacketHandler getTrackerHandler() {
         return getTrackerHandler(Types.VAR_INT, 1);
     }
 
-    protected PacketHandler getTrackerHandler(EntityType entityType, Type<? extends Number> intType) {
-        return wrapper -> tracker(wrapper.user()).addEntity((int) wrapper.get(intType, 0), entityType);
+    protected PacketHandler getTrackerHandler(EntityType entityType) {
+        return wrapper -> tracker(wrapper.user()).addEntity((int) wrapper.get(Types.VAR_INT, 0), entityType);
+    }
+
+    protected PacketHandler getPlayerTrackerHandler() {
+        return wrapper -> {
+            final int entityId = wrapper.get(Types.INT, 0);
+
+            final EntityTracker tracker = tracker(wrapper.user());
+            tracker(wrapper.user()).setClientEntityId(entityId);
+            tracker.addEntity(entityId, tracker.playerType());
+        };
     }
 
     protected PacketHandler getDimensionHandler(int index) {
         return wrapper -> {
-            ClientWorld clientWorld = wrapper.user().get(ClientWorld.class);
+            ClientWorld clientWorld = wrapper.user().getClientWorld(this.protocol.getClass());
             int dimensionId = wrapper.get(Types.INT, index);
-            clientWorld.setEnvironment(dimensionId);
+            if (clientWorld.setEnvironment(dimensionId)) {
+                tracker(wrapper.user()).clearEntities();
+            }
         };
     }
 }

common/src/main/java/com/viaversion/viabackwards/api/rewriters/LegacyEntityRewriter.java

@@ -66,8 +66,10 @@ protected void registerRespawn(C packetType) {
             public void register() {
                 map(Types.INT);
                 handler(wrapper -> {
-                    ClientWorld clientWorld = wrapper.user().get(ClientWorld.class);
-                    clientWorld.setEnvironment(wrapper.get(Types.INT, 0));
+                    ClientWorld clientWorld = wrapper.user().getClientWorld(protocol.getClass());
+                    if (clientWorld.setEnvironment(wrapper.get(Types.INT, 0))) {
+                        tracker(wrapper.user()).clearEntities();
+                    }
                 });
             }
         });
@@ -81,8 +83,8 @@ public void register() {
                 map(Types.UNSIGNED_BYTE); // 1 - Gamemode
                 map(Types.INT); // 2 - Dimension
                 handler(wrapper -> {
-                    ClientWorld clientChunks = wrapper.user().get(ClientWorld.class);
-                    clientChunks.setEnvironment(wrapper.get(Types.INT, 1));
+                    ClientWorld clientWorld = wrapper.user().getClientWorld(protocol.getClass());
+                    clientWorld.setEnvironment(wrapper.get(Types.INT, 1));
 
                     final int entityId = wrapper.get(Types.INT, 0);
                     addTrackedEntity(wrapper, entityId, playerType);

common/src/main/java/com/viaversion/viabackwards/protocol/v1_10to1_9_3/Protocol1_10To1_9_3.java

@@ -98,11 +98,8 @@ public void register() {
 
     @Override
     public void init(UserConnection user) {
-        if (!user.has(ClientWorld.class)) {
-            user.put(new ClientWorld());
-        }
-
         user.addEntityTracker(this.getClass(), new EntityTrackerBase(user, EntityTypes1_10.EntityType.PLAYER));
+        user.addClientWorld(this.getClass(), new ClientWorld());
     }
 
     @Override

common/src/main/java/com/viaversion/viabackwards/protocol/v1_10to1_9_3/rewriter/BlockItemPacketRewriter1_10.java

@@ -50,7 +50,7 @@ protected void registerPackets() {
         registerSetCreativeModeSlot(ServerboundPackets1_9_3.SET_CREATIVE_MODE_SLOT);
 
         protocol.registerClientbound(ClientboundPackets1_9_3.LEVEL_CHUNK, wrapper -> {
-            ClientWorld clientWorld = wrapper.user().get(ClientWorld.class);
+            ClientWorld clientWorld = wrapper.user().getClientWorld(Protocol1_10To1_9_3.class);
 
             ChunkType1_9_3 type = ChunkType1_9_3.forEnvironment(clientWorld.getEnvironment());
             Chunk chunk = wrapper.passthrough(type);

common/src/main/java/com/viaversion/viabackwards/protocol/v1_11_1to1_11/Protocol1_11_1To1_11.java

@@ -39,11 +39,8 @@ public Protocol1_11_1To1_11() {
 
     @Override
     public void init(UserConnection user) {
-        if (!user.has(ClientWorld.class)) {
-            user.put(new ClientWorld());
-        }
-
         user.addEntityTracker(this.getClass(), new EntityTrackerBase(user, EntityTypes1_11.EntityType.PLAYER));
+        user.addClientWorld(this.getClass(), new ClientWorld());
     }
 
     @Override

common/src/main/java/com/viaversion/viabackwards/protocol/v1_11to1_10/Protocol1_11To1_10.java

@@ -62,11 +62,8 @@ protected void registerPackets() {
 
     @Override
     public void init(UserConnection user) {
-        if (!user.has(ClientWorld.class)) {
-            user.put(new ClientWorld());
-        }
-
         user.addEntityTracker(this.getClass(), new EntityTrackerBase(user, EntityTypes1_11.EntityType.PLAYER));
+        user.addClientWorld(this.getClass(), new ClientWorld());
 
         if (!user.has(WindowTracker.class)) {
             user.put(new WindowTracker());

common/src/main/java/com/viaversion/viabackwards/protocol/v1_11to1_10/rewriter/BlockItemPacketRewriter1_11.java

@@ -147,7 +147,7 @@ public void register() {
         registerSetCreativeModeSlot(ServerboundPackets1_9_3.SET_CREATIVE_MODE_SLOT);
 
         protocol.registerClientbound(ClientboundPackets1_9_3.LEVEL_CHUNK, wrapper -> {
-            ClientWorld clientWorld = wrapper.user().get(ClientWorld.class);
+            ClientWorld clientWorld = wrapper.user().getClientWorld(Protocol1_11To1_10.class);
 
             ChunkType1_9_3 type = ChunkType1_9_3.forEnvironment(clientWorld.getEnvironment()); // Use the 1.10 Chunk type since nothing changed.
             Chunk chunk = wrapper.passthrough(type);

common/src/main/java/com/viaversion/viabackwards/protocol/v1_12to1_11_1/Protocol1_12To1_11_1.java

@@ -71,11 +71,8 @@ protected void registerPackets() {
 
     @Override
     public void init(UserConnection user) {
-        if (!user.has(ClientWorld.class)) {
-            user.put(new ClientWorld());
-        }
-
         user.addEntityTracker(this.getClass(), new EntityTrackerBase(user, EntityTypes1_12.EntityType.PLAYER));
+        user.addClientWorld(this.getClass(), new ClientWorld());
 
         user.put(new ShoulderTracker(user));
     }

common/src/main/java/com/viaversion/viabackwards/protocol/v1_12to1_11_1/rewriter/BlockItemPacketRewriter1_12.java

@@ -131,7 +131,7 @@ public void register() {
         registerSetCreativeModeSlot(ServerboundPackets1_9_3.SET_CREATIVE_MODE_SLOT);
 
         protocol.registerClientbound(ClientboundPackets1_12.LEVEL_CHUNK, wrapper -> {
-            ClientWorld clientWorld = wrapper.user().get(ClientWorld.class);
+            ClientWorld clientWorld = wrapper.user().getClientWorld(Protocol1_12To1_11_1.class);
 
             ChunkType1_9_3 type = ChunkType1_9_3.forEnvironment(clientWorld.getEnvironment()); // Use the 1.9.4 Chunk type since nothing changed.
             Chunk chunk = wrapper.passthrough(type);

common/src/main/java/com/viaversion/viabackwards/protocol/v1_12to1_11_1/rewriter/EntityPacketRewriter1_12.java

@@ -119,9 +119,8 @@ public void register() {
                 map(Types.UNSIGNED_BYTE); // 1 - Gamemode
                 map(Types.INT); // 2 - Dimension
 
-                handler(getTrackerHandler(EntityTypes1_12.EntityType.PLAYER, Types.INT));
-
                 handler(getDimensionHandler(1));
+                handler(getPlayerTrackerHandler());
 
                 handler(wrapper -> {
                     ShoulderTracker tracker = wrapper.user().get(ShoulderTracker.class);

common/src/main/java/com/viaversion/viabackwards/protocol/v1_13_1to1_13/Protocol1_13_1To1_13.java

@@ -203,10 +203,7 @@ public void register() {
     @Override
     public void init(UserConnection user) {
         user.addEntityTracker(getClass(), new EntityTrackerBase(user, EntityTypes1_13.EntityType.PLAYER));
-
-        if (!user.has(ClientWorld.class)) {
-            user.put(new ClientWorld());
-        }
+        user.addClientWorld(getClass(), new ClientWorld());
     }
 
     @Override

common/src/main/java/com/viaversion/viabackwards/protocol/v1_13_1to1_13/rewriter/WorldPacketRewriter1_13_1.java

@@ -34,7 +34,7 @@ public static void register(Protocol1_13_1To1_13 protocol) {
         BlockRewriter<ClientboundPackets1_13> blockRewriter = BlockRewriter.legacy(protocol);
 
         protocol.registerClientbound(ClientboundPackets1_13.LEVEL_CHUNK, wrapper -> {
-            ClientWorld clientWorld = wrapper.user().get(ClientWorld.class);
+            ClientWorld clientWorld = wrapper.user().getClientWorld(Protocol1_13_1To1_13.class);
             Chunk chunk = wrapper.passthrough(ChunkType1_13.forEnvironment(clientWorld.getEnvironment()));
 
             blockRewriter.handleChunk(chunk);

common/src/main/java/com/viaversion/viabackwards/protocol/v1_13to1_12_2/Protocol1_13To1_12_2.java

@@ -114,11 +114,8 @@ protected void registerPackets() {
 
     @Override
     public void init(UserConnection user) {
-        if (!user.has(ClientWorld.class)) {
-            user.put(new ClientWorld());
-        }
-
         user.addEntityTracker(this.getClass(), new EntityTrackerBase(user, EntityTypes1_13.EntityType.PLAYER));
+        user.addClientWorld(this.getClass(), new ClientWorld());
 
         user.put(new BackwardsBlockStorage());
         user.put(new TabCompleteStorage());

common/src/main/java/com/viaversion/viabackwards/protocol/v1_13to1_12_2/rewriter/BlockItemPacketRewriter1_13.java

@@ -279,7 +279,7 @@ public void register() {
         });
 
         protocol.registerClientbound(ClientboundPackets1_13.LEVEL_CHUNK, wrapper -> {
-            ClientWorld clientWorld = wrapper.user().get(ClientWorld.class);
+            ClientWorld clientWorld = wrapper.user().getClientWorld(Protocol1_13To1_12_2.class);
 
             ChunkType1_9_3 type_old = ChunkType1_9_3.forEnvironment(clientWorld.getEnvironment());
             ChunkType1_13 type = ChunkType1_13.forEnvironment(clientWorld.getEnvironment());

common/src/main/java/com/viaversion/viabackwards/protocol/v1_13to1_12_2/rewriter/EntityPacketRewriter1_13.java

@@ -27,6 +27,7 @@
 import com.viaversion.viabackwards.protocol.v1_13to1_12_2.storage.BackwardsBlockStorage;
 import com.viaversion.viabackwards.protocol.v1_13to1_12_2.storage.NoteBlockStorage;
 import com.viaversion.viabackwards.protocol.v1_13to1_12_2.storage.PlayerPositionStorage1_13;
+import com.viaversion.viaversion.api.minecraft.ClientWorld;
 import com.viaversion.viaversion.api.minecraft.Particle;
 import com.viaversion.viaversion.api.minecraft.entities.EntityType;
 import com.viaversion.viaversion.api.minecraft.entities.EntityTypes1_13;
@@ -172,7 +173,7 @@ public void register() {
                 map(Types.VAR_INT);
                 map(Types.UUID);
 
-                handler(getTrackerHandler(EntityTypes1_13.EntityType.PAINTING, Types.VAR_INT));
+                handler(getTrackerHandler(EntityTypes1_13.EntityType.PAINTING));
                 handler(wrapper -> {
                     int motive = wrapper.read(Types.VAR_INT);
                     String title = PaintingNames1_13.getStringId(motive);
@@ -188,10 +189,15 @@ public void register() {
             public void register() {
                 map(Types.INT); // 0 - Dimension ID
 
-                handler(getDimensionHandler(0));
                 handler(wrapper -> {
-                    wrapper.user().get(BackwardsBlockStorage.class).clear();
-                    wrapper.user().get(NoteBlockStorage.class).clear();
+                    ClientWorld clientWorld = wrapper.user().getClientWorld(Protocol1_13To1_12_2.class);
+                    int dimensionId = wrapper.get(Types.INT, 0);
+
+                    if (clientWorld.setEnvironment(dimensionId)) {
+                        tracker(wrapper.user()).clearEntities();
+                        wrapper.user().get(BackwardsBlockStorage.class).clear();
+                        wrapper.user().get(NoteBlockStorage.class).clear();
+                    }
                 });
             }
         });

common/src/main/java/com/viaversion/viabackwards/protocol/v1_14to1_13_2/Protocol1_14To1_13_2.java

@@ -162,11 +162,8 @@ private static boolean isSet(int mask, int i) {
 
     @Override
     public void init(UserConnection user) {
-        if (!user.has(ClientWorld.class)) {
-            user.put(new ClientWorld());
-        }
-
         user.addEntityTracker(this.getClass(), new EntityTrackerBase(user, EntityTypes1_14.PLAYER));
+        user.addClientWorld(this.getClass(), new ClientWorld());
 
         if (!user.has(ChunkLightStorage.class)) {
             user.put(new ChunkLightStorage());

common/src/main/java/com/viaversion/viabackwards/protocol/v1_14to1_13_2/rewriter/BlockItemPacketRewriter1_14.java

@@ -360,7 +360,7 @@ public void register() {
         });
 
         protocol.registerClientbound(ClientboundPackets1_14.LEVEL_CHUNK, wrapper -> {
-            ClientWorld clientWorld = wrapper.user().get(ClientWorld.class);
+            ClientWorld clientWorld = wrapper.user().getClientWorld(Protocol1_14To1_13_2.class);
             Chunk chunk = wrapper.read(ChunkType1_14.TYPE);
             wrapper.write(ChunkType1_13.forEnvironment(clientWorld.getEnvironment()), chunk);
 

common/src/main/java/com/viaversion/viabackwards/protocol/v1_14to1_13_2/rewriter/EntityPacketRewriter1_14.java

@@ -304,8 +304,8 @@ public void register() {
                 map(Types.UNSIGNED_BYTE); // 1 - Gamemode
                 map(Types.INT); // 2 - Dimension
 
-                handler(getTrackerHandler(EntityTypes1_14.PLAYER, Types.INT));
                 handler(getDimensionHandler(1));
+                handler(getPlayerTrackerHandler());
                 handler(wrapper -> {
                     short difficulty = wrapper.user().get(DifficultyStorage.class).getDifficulty();
                     wrapper.write(Types.UNSIGNED_BYTE, difficulty);
@@ -329,14 +329,16 @@ public void register() {
                 map(Types.INT); // 0 - Dimension ID
 
                 handler(wrapper -> {
-                    ClientWorld clientWorld = wrapper.user().get(ClientWorld.class);
+                    ClientWorld clientWorld = wrapper.user().getClientWorld(Protocol1_14To1_13_2.class);
                     int dimensionId = wrapper.get(Types.INT, 0);
-                    clientWorld.setEnvironment(dimensionId);
+
+                    if (clientWorld.setEnvironment(dimensionId)) {
+                        tracker(wrapper.user()).clearEntities();
+                        wrapper.user().get(ChunkLightStorage.class).clear();
+                    }
 
                     short difficulty = wrapper.user().get(DifficultyStorage.class).getDifficulty();
                     wrapper.write(Types.UNSIGNED_BYTE, difficulty);
-
-                    wrapper.user().get(ChunkLightStorage.class).clear();
                 });
             }
         });

common/src/main/java/com/viaversion/viabackwards/protocol/v1_15to1_14_4/Protocol1_15To1_14_4.java

@@ -25,6 +25,7 @@
 import com.viaversion.viabackwards.protocol.v1_15to1_14_4.rewriter.BlockItemPacketRewriter1_15;
 import com.viaversion.viabackwards.protocol.v1_15to1_14_4.rewriter.EntityPacketRewriter1_15;
 import com.viaversion.viaversion.api.connection.UserConnection;
+import com.viaversion.viaversion.api.minecraft.ClientWorld;
 import com.viaversion.viaversion.api.minecraft.RegistryType;
 import com.viaversion.viaversion.api.minecraft.entities.EntityTypes1_15;
 import com.viaversion.viaversion.api.protocol.packet.PacketWrapper;
@@ -102,8 +103,10 @@ private int toEffectCoordinate(float coordinate) {
 
     @Override
     public void init(UserConnection user) {
-        user.put(new ImmediateRespawnStorage());
         user.addEntityTracker(getClass(), new EntityTrackerBase(user, EntityTypes1_15.PLAYER));
+        user.addClientWorld(getClass(), new ClientWorld());
+
+        user.put(new ImmediateRespawnStorage());
     }
 
     @Override

common/src/main/java/com/viaversion/viabackwards/protocol/v1_15to1_14_4/rewriter/EntityPacketRewriter1_15.java

@@ -96,6 +96,7 @@ public void register() {
             public void register() {
                 map(Types.INT);
                 read(Types.LONG); // Seed
+                handler(getDimensionHandler(0));
             }
         });
 
@@ -113,7 +114,7 @@ public void register() {
                 map(Types.VAR_INT); // 5 - View Distance
                 map(Types.BOOLEAN); // 6 - Reduce Debug Info
 
-                handler(getTrackerHandler(EntityTypes1_15.PLAYER, Types.INT));
+                handler(getPlayerTrackerHandler());
 
                 handler(wrapper -> {
                     boolean immediateRespawn = !wrapper.read(Types.BOOLEAN); // Inverted
@@ -138,7 +139,7 @@ public void register() {
                 map(Types.BYTE); // 6 - Pitch
                 handler(wrapper -> wrapper.write(Types1_14.ENTITY_DATA_LIST, new ArrayList<>())); // Entity data is no longer sent in 1.15, so we have to send an empty one
 
-                handler(getTrackerHandler(EntityTypes1_15.PLAYER, Types.VAR_INT));
+                handler(getTrackerHandler(EntityTypes1_15.PLAYER));
             }
         });
 

common/src/main/java/com/viaversion/viabackwards/protocol/v1_16_2to1_16_1/rewriter/EntityPacketRewriter1_16_2.java

@@ -101,13 +101,15 @@ public void register() {
                     wrapper.write(Types.UNSIGNED_BYTE, (short) Math.min(maxPlayers, 255));
                 });
                 // ...
-                handler(getTrackerHandler(EntityTypes1_16_2.PLAYER, Types.INT));
+                handler(getPlayerTrackerHandler());
             }
         });
 
         protocol.registerClientbound(ClientboundPackets1_16_2.RESPAWN, wrapper -> {
             CompoundTag dimensionData = wrapper.read(Types.NAMED_COMPOUND_TAG);
             wrapper.write(Types.STRING, getDimensionFromData(dimensionData));
+
+            tracker(wrapper.user()).clearEntities();
         });
     }
 

common/src/main/java/com/viaversion/viabackwards/protocol/v1_16to1_15_2/Protocol1_16To1_15_2.java

@@ -168,14 +168,12 @@ public void register() {
 
     @Override
     public void init(UserConnection user) {
-        if (!user.has(ClientWorld.class)) {
-            user.put(new ClientWorld());
-        }
+        user.addEntityTracker(this.getClass(), new EntityTrackerBase(user, EntityTypes1_16.PLAYER));
+        user.addClientWorld(this.getClass(), new ClientWorld());
 
         user.put(new PlayerSneakStorage());
         user.put(new WorldNameTracker());
         user.put(new PlayerAttributesStorage());
-        user.addEntityTracker(this.getClass(), new EntityTrackerBase(user, EntityTypes1_16.PLAYER));
     }
 
     @Override