Refactor dimension switch handling across all protocols (#4146)

api/src/main/java/com/viaversion/viaversion/api/connection/UserConnection.java

@@ -24,6 +24,7 @@ package com.viaversion.viaversion.api.connection;
 
 import com.viaversion.viaversion.api.configuration.ViaVersionConfig;
 import com.viaversion.viaversion.api.data.entity.EntityTracker;
+import com.viaversion.viaversion.api.minecraft.ClientWorld;
 import com.viaversion.viaversion.api.protocol.Protocol;
 import com.viaversion.viaversion.api.protocol.packet.PacketTracker;
 import com.viaversion.viaversion.api.protocol.packet.PacketWrapper;
@@ -99,14 +100,33 @@ public interface UserConnection {
     void addEntityTracker(Class<? extends Protocol> protocolClass, EntityTracker tracker);
 
     /**
-     * Clear stored objects and entity trackers.
+     * Returns the client world by the given protocol class if present.
+     *
+     * @param protocolClass protocol class
+     * @param <T>           client world type
+     * @return client world if present
+     */
+    @Nullable
+    <T extends ClientWorld> T getClientWorld(Class<? extends Protocol> protocolClass);
+
+    /**
+     * Adds a client world to the user connection.
+     * Does not override existing client worlds.
+     *
+     * @param protocolClass protocol class
+     * @param clientWorld   client world
+     */
+    void addClientWorld(Class<? extends Protocol> protocolClass, ClientWorld clientWorld);
+
+    /**
+     * Clear stored objects, entity trackers and client worlds.
      */
     default void clearStoredObjects() {
         clearStoredObjects(false);
     }
 
     /**
-     * Clear stored objects and entity trackers.
+     * Clear stored objects, entity trackers and client worlds.
      * If cleared for a proxy server switch, some stored objects and tracker data will be retained.
      *
      * @param isServerSwitch whether the clear is due to a server switch

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

@@ -78,7 +78,7 @@ public interface EntityTracker {
     void removeEntity(int id);
 
     /**
-     * Clears stored entity types and data.
+     * Clears stored entity types and data, only leaving behind the client entity.
      */
     void clearEntities();
 
@@ -177,12 +177,4 @@ public interface EntityTracker {
     @Nullable DimensionData dimensionData(int dimensionId);
 
     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();
 }

api/src/main/java/com/viaversion/viaversion/api/minecraft/ClientWorld.java

@@ -23,13 +23,12 @@
 package com.viaversion.viaversion.api.minecraft;
 
 import com.viaversion.viaversion.api.connection.StorableObject;
-import org.checkerframework.checker.nullness.qual.Nullable;
 
 /**
  * Stored up until 1.14 to be used in chunk sending.
  */
 public class ClientWorld implements StorableObject {
-    private Environment environment;
+    private Environment environment = Environment.NORMAL;
 
     public ClientWorld() {
     }
@@ -38,11 +37,20 @@ public class ClientWorld implements StorableObject {
         this.environment = environment;
     }
 
-    public @Nullable Environment getEnvironment() {
+    public Environment getEnvironment() {
         return environment;
     }
 
-    public void setEnvironment(final int environmentId) {
+    /**
+     * Sets the environment of the world and returns whether the environment was changed.
+     *
+     * @param environmentId the id of the environment to set
+     * @return whether the environment was changed
+     */
+    public boolean setEnvironment(final int environmentId) {
+        final int previousEnvironmentId = environment.id();
         this.environment = Environment.getEnvironmentById(environmentId);
+
+        return previousEnvironmentId != environmentId;
     }
 }

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

@@ -23,6 +23,7 @@ import com.viaversion.viaversion.api.connection.ProtocolInfo;
 import com.viaversion.viaversion.api.connection.StorableObject;
 import com.viaversion.viaversion.api.connection.UserConnection;
 import com.viaversion.viaversion.api.data.entity.EntityTracker;
+import com.viaversion.viaversion.api.minecraft.ClientWorld;
 import com.viaversion.viaversion.api.platform.ViaInjector;
 import com.viaversion.viaversion.api.protocol.Protocol;
 import com.viaversion.viaversion.api.protocol.packet.Direction;
@@ -58,6 +59,7 @@ public class UserConnectionImpl implements UserConnection {
     private final long id = IDS.incrementAndGet();
     private final Map<Class<?>, StorableObject> storedObjects = new ConcurrentHashMap<>();
     private final Map<Class<? extends Protocol>, EntityTracker> entityTrackers = new HashMap<>();
+    private final Map<Class<? extends Protocol>, ClientWorld> clientWorlds = new HashMap<>();
     private final PacketTracker packetTracker = new PacketTracker(this);
     private final Set<UUID> passthroughTokens = Collections.newSetFromMap(CacheBuilder.newBuilder()
         .expireAfterWrite(10, TimeUnit.SECONDS)
@@ -125,9 +127,17 @@ public class UserConnectionImpl implements UserConnection {
 
     @Override
     public void addEntityTracker(Class<? extends Protocol> protocolClass, EntityTracker tracker) {
-        if (!entityTrackers.containsKey(protocolClass)) {
+        entityTrackers.putIfAbsent(protocolClass, tracker);
-            entityTrackers.put(protocolClass, tracker);
     }
+
+    @Override
+    public @Nullable <T extends ClientWorld> T getClientWorld(final Class<? extends Protocol> protocolClass) {
+        return (T) clientWorlds.get(protocolClass);
+    }
+
+    @Override
+    public void addClientWorld(final Class<? extends Protocol> protocolClass, final ClientWorld clientWorld) {
+        clientWorlds.putIfAbsent(protocolClass, clientWorld);
     }
 
     @Override
@@ -142,7 +152,6 @@ public class UserConnectionImpl implements UserConnection {
             });
             for (EntityTracker tracker : entityTrackers.values()) {
                 tracker.clearEntities();
-                tracker.trackClientEntity();
             }
         } else {
             for (StorableObject object : storedObjects.values()) {
@@ -150,6 +159,7 @@ public class UserConnectionImpl implements UserConnection {
             }
             storedObjects.clear();
             entityTrackers.clear();
+            clientWorlds.clear();
         }
     }
 

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

@@ -87,7 +87,6 @@ public class EntityTrackerBase implements EntityTracker, ClientEntityIdChangeLis
         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) {
         entities.remove(id);
@@ -99,6 +98,11 @@ public class EntityTrackerBase implements EntityTracker, ClientEntityIdChangeLis
         for (final int id : entities.keySet().toIntArray()) {
             removeEntity(id);
         }
+
+        // Re-add the client entity. Keep the call above to untrack attached data if necessary
+        if (clientEntityId != null) {
+            entities.put(clientEntityId.intValue(), new TrackedEntityImpl(playerType));
+        }
     }
 
     @Override
@@ -127,15 +131,6 @@ public class EntityTrackerBase implements EntityTracker, ClientEntityIdChangeLis
         this.clientEntityId = clientEntityId;
     }
 
-    @Override
-    public boolean trackClientEntity() {
-        if (clientEntityId != null) {
-            entities.put(clientEntityId.intValue(), new TrackedEntityImpl(playerType));
-            return true;
-        }
-        return false;
-    }
-
     @Override
     public int currentWorldSectionHeight() {
         return currentWorldSectionHeight;

common/src/main/java/com/viaversion/viaversion/protocols/v1_10to1_11/Protocol1_10To1_11.java

@@ -99,7 +99,7 @@ public class Protocol1_10To1_11 extends AbstractProtocol<ClientboundPackets1_9_3
         });
 
         registerClientbound(ClientboundPackets1_9_3.LEVEL_CHUNK, wrapper -> {
-            ClientWorld clientWorld = wrapper.user().get(ClientWorld.class);
+            ClientWorld clientWorld = wrapper.user().getClientWorld(Protocol1_10To1_11.class);
 
             Chunk chunk = wrapper.passthrough(ChunkType1_9_3.forEnvironment(clientWorld.getEnvironment()));
 
@@ -120,31 +120,6 @@ public class Protocol1_10To1_11 extends AbstractProtocol<ClientboundPackets1_9_3
             }
         });
 
-        registerClientbound(ClientboundPackets1_9_3.LOGIN, new PacketHandlers() {
-            @Override
-            public void register() {
-                map(Types.INT); // 0 - Entity ID
-                map(Types.UNSIGNED_BYTE); // 1 - Gamemode
-                map(Types.INT); // 2 - Dimension
-                handler(wrapper -> {
-                    ClientWorld clientChunks = wrapper.user().get(ClientWorld.class);
-                    int dimensionId = wrapper.get(Types.INT, 1);
-                    clientChunks.setEnvironment(dimensionId);
-                });
-            }
-        });
-        registerClientbound(ClientboundPackets1_9_3.RESPAWN, new PacketHandlers() {
-            @Override
-            public void register() {
-                map(Types.INT);
-                handler(wrapper -> {
-                    ClientWorld clientWorld = wrapper.user().get(ClientWorld.class);
-                    int dimensionId = wrapper.get(Types.INT, 0);
-                    clientWorld.setEnvironment(dimensionId);
-                });
-            }
-        });
-
         this.registerClientbound(ClientboundPackets1_9_3.LEVEL_EVENT, new PacketHandlers() {
             @Override
             public void register() {
@@ -238,11 +213,8 @@ public class Protocol1_10To1_11 extends AbstractProtocol<ClientboundPackets1_9_3
 
     @Override
     public void init(UserConnection userConnection) {
-        // Entity tracker
         userConnection.addEntityTracker(this.getClass(), new EntityTracker1_11(userConnection));
-
+        userConnection.addClientWorld(this.getClass(), new ClientWorld());
-        if (!userConnection.has(ClientWorld.class))
-            userConnection.put(new ClientWorld());
     }
 
     @Override

common/src/main/java/com/viaversion/viaversion/protocols/v1_10to1_11/rewriter/EntityPacketRewriter1_11.java

@@ -20,6 +20,7 @@ package com.viaversion.viaversion.protocols.v1_10to1_11.rewriter;
 import com.viaversion.nbt.tag.CompoundTag;
 import com.viaversion.nbt.tag.StringTag;
 import com.viaversion.viaversion.api.Via;
+import com.viaversion.viaversion.api.minecraft.ClientWorld;
 import com.viaversion.viaversion.api.minecraft.entities.EntityTypes1_10;
 import com.viaversion.viaversion.api.minecraft.entities.EntityTypes1_11;
 import com.viaversion.viaversion.api.minecraft.entities.EntityTypes1_11.EntityType;
@@ -48,6 +49,34 @@ public class EntityPacketRewriter1_11 extends EntityRewriter<ClientboundPackets1
 
     @Override
     protected void registerPackets() {
+        protocol.registerClientbound(ClientboundPackets1_9_3.LOGIN, new PacketHandlers() {
+            @Override
+            public void register() {
+                map(Types.INT); // 0 - Entity ID
+                map(Types.UNSIGNED_BYTE); // 1 - Gamemode
+                map(Types.INT); // 2 - Dimension
+                handler(wrapper -> {
+                    ClientWorld clientWorld = wrapper.user().getClientWorld(Protocol1_10To1_11.class);
+                    int dimensionId = wrapper.get(Types.INT, 1);
+                    clientWorld.setEnvironment(dimensionId);
+                });
+            }
+        });
+
+        protocol.registerClientbound(ClientboundPackets1_9_3.RESPAWN, new PacketHandlers() {
+            @Override
+            public void register() {
+                map(Types.INT);
+                handler(wrapper -> {
+                    ClientWorld clientWorld = wrapper.user().getClientWorld(Protocol1_10To1_11.class);
+                    int dimensionId = wrapper.get(Types.INT, 0);
+                    if (clientWorld.setEnvironment(dimensionId)) {
+                        tracker(wrapper.user()).clearEntities();
+                    }
+                });
+            }
+        });
+
         protocol.registerClientbound(ClientboundPackets1_9_3.ADD_ENTITY, new PacketHandlers() {
             @Override
             public void register() {

common/src/main/java/com/viaversion/viaversion/protocols/v1_11_1to1_12/Protocol1_11_1To1_12.java

@@ -72,7 +72,7 @@ public class Protocol1_11_1To1_12 extends AbstractProtocol<ClientboundPackets1_9
         });
 
         registerClientbound(ClientboundPackets1_9_3.LEVEL_CHUNK, wrapper -> {
-            ClientWorld clientWorld = wrapper.user().get(ClientWorld.class);
+            ClientWorld clientWorld = wrapper.user().getClientWorld(Protocol1_11_1To1_12.class);
 
             ChunkType1_9_3 type = ChunkType1_9_3.forEnvironment(clientWorld.getEnvironment());
             Chunk chunk = wrapper.passthrough(type);
@@ -101,38 +101,6 @@ public class Protocol1_11_1To1_12 extends AbstractProtocol<ClientboundPackets1_9
             }
         });
 
-        registerClientbound(ClientboundPackets1_9_3.LOGIN, new PacketHandlers() {
-            @Override
-            public void register() {
-                map(Types.INT);
-                map(Types.UNSIGNED_BYTE);
-                map(Types.INT);
-                handler(wrapper -> {
-                    UserConnection user = wrapper.user();
-                    ClientWorld clientChunks = user.get(ClientWorld.class);
-                    int dimensionId = wrapper.get(Types.INT, 1);
-                    clientChunks.setEnvironment(dimensionId);
-
-                    // Reset recipes
-                    if (user.getProtocolInfo().protocolVersion().newerThanOrEqualTo(ProtocolVersion.v1_13)) {
-                        wrapper.create(ClientboundPackets1_13.UPDATE_RECIPES, packetWrapper -> packetWrapper.write(Types.VAR_INT, 0))
-                            .scheduleSend(Protocol1_12_2To1_13.class);
-                    }
-                });
-            }
-        });
-        registerClientbound(ClientboundPackets1_9_3.RESPAWN, new PacketHandlers() {
-            @Override
-            public void register() {
-                map(Types.INT);
-                handler(wrapper -> {
-                    ClientWorld clientWorld = wrapper.user().get(ClientWorld.class);
-                    int dimensionId = wrapper.get(Types.INT, 0);
-                    clientWorld.setEnvironment(dimensionId);
-                });
-            }
-        });
-
         new SoundRewriter<>(this, this::getNewSoundId).registerSound(ClientboundPackets1_9_3.SOUND);
 
 
@@ -205,9 +173,7 @@ public class Protocol1_11_1To1_12 extends AbstractProtocol<ClientboundPackets1_9
     @Override
     public void init(UserConnection userConnection) {
         userConnection.addEntityTracker(this.getClass(), new EntityTrackerBase(userConnection, EntityTypes1_12.EntityType.PLAYER));
-        if (!userConnection.has(ClientWorld.class)) {
+        userConnection.addClientWorld(this.getClass(), new ClientWorld());
-            userConnection.put(new ClientWorld());
-        }
     }
 
     @Override

common/src/main/java/com/viaversion/viaversion/protocols/v1_11_1to1_12/rewriter/EntityPacketRewriter1_12.java

@@ -17,14 +17,18 @@
  */
 package com.viaversion.viaversion.protocols.v1_11_1to1_12.rewriter;
 
+import com.viaversion.viaversion.api.minecraft.ClientWorld;
 import com.viaversion.viaversion.api.minecraft.entities.EntityType;
 import com.viaversion.viaversion.api.minecraft.entities.EntityTypes1_12;
 import com.viaversion.viaversion.api.minecraft.item.Item;
 import com.viaversion.viaversion.api.protocol.remapper.PacketHandlers;
+import com.viaversion.viaversion.api.protocol.version.ProtocolVersion;
 import com.viaversion.viaversion.api.type.Types;
 import com.viaversion.viaversion.api.type.types.version.Types1_12;
 import com.viaversion.viaversion.api.type.types.version.Types1_9;
 import com.viaversion.viaversion.protocols.v1_11_1to1_12.Protocol1_11_1To1_12;
+import com.viaversion.viaversion.protocols.v1_12_2to1_13.Protocol1_12_2To1_13;
+import com.viaversion.viaversion.protocols.v1_12_2to1_13.packet.ClientboundPackets1_13;
 import com.viaversion.viaversion.protocols.v1_9_1to1_9_3.packet.ClientboundPackets1_9_3;
 import com.viaversion.viaversion.rewriter.EntityRewriter;
 
@@ -36,6 +40,40 @@ public class EntityPacketRewriter1_12 extends EntityRewriter<ClientboundPackets1
 
     @Override
     protected void registerPackets() {
+        protocol.registerClientbound(ClientboundPackets1_9_3.LOGIN, new PacketHandlers() {
+            @Override
+            public void register() {
+                map(Types.INT);
+                map(Types.UNSIGNED_BYTE);
+                map(Types.INT);
+                handler(wrapper -> {
+                    ClientWorld clientWorld = wrapper.user().getClientWorld(Protocol1_11_1To1_12.class);
+                    int dimensionId = wrapper.get(Types.INT, 1);
+                    clientWorld.setEnvironment(dimensionId);
+
+                    // Reset recipes
+                    if (wrapper.user().getProtocolInfo().protocolVersion().newerThanOrEqualTo(ProtocolVersion.v1_13)) {
+                        wrapper.create(ClientboundPackets1_13.UPDATE_RECIPES, packetWrapper -> packetWrapper.write(Types.VAR_INT, 0))
+                            .scheduleSend(Protocol1_12_2To1_13.class);
+                    }
+                });
+            }
+        });
+
+        protocol.registerClientbound(ClientboundPackets1_9_3.RESPAWN, new PacketHandlers() {
+            @Override
+            public void register() {
+                map(Types.INT);
+                handler(wrapper -> {
+                    ClientWorld clientWorld = wrapper.user().getClientWorld(Protocol1_11_1To1_12.class);
+                    int dimensionId = wrapper.get(Types.INT, 0);
+                    if (clientWorld.setEnvironment(dimensionId)) {
+                        tracker(wrapper.user()).clearEntities();
+                    }
+                });
+            }
+        });
+
         protocol.registerClientbound(ClientboundPackets1_9_3.ADD_ENTITY, new PacketHandlers() {
             @Override
             public void register() {

common/src/main/java/com/viaversion/viaversion/protocols/v1_12_2to1_13/Protocol1_12_2To1_13.java

@@ -401,23 +401,6 @@ public class Protocol1_12_2To1_13 extends AbstractProtocol<ClientboundPackets1_1
             }
         });
 
-        registerClientbound(ClientboundPackets1_12_1.RESPAWN, new PacketHandlers() {
-            @Override
-            public void register() {
-                map(Types.INT); // 0 - Dimension ID
-                handler(wrapper -> {
-                    ClientWorld clientWorld = wrapper.user().get(ClientWorld.class);
-                    int dimensionId = wrapper.get(Types.INT, 0);
-                    clientWorld.setEnvironment(dimensionId);
-
-                    if (Via.getConfig().isServersideBlockConnections()) {
-                        ConnectionData.clearBlockStorage(wrapper.user());
-                    }
-                });
-                handler(SEND_DECLARE_COMMANDS_AND_TAGS);
-            }
-        });
-
         registerClientbound(ClientboundPackets1_12_1.SET_OBJECTIVE, new PacketHandlers() {
             @Override
             public void register() {
@@ -844,9 +827,9 @@ public class Protocol1_12_2To1_13 extends AbstractProtocol<ClientboundPackets1_1
     @Override
     public void init(UserConnection userConnection) {
         userConnection.addEntityTracker(this.getClass(), new EntityTrackerBase(userConnection, EntityTypes1_13.EntityType.PLAYER));
+        userConnection.addClientWorld(this.getClass(), new ClientWorld());
+
         userConnection.put(new TabCompleteTracker());
-        if (!userConnection.has(ClientWorld.class))
-            userConnection.put(new ClientWorld());
         userConnection.put(new BlockStorage());
         if (Via.getConfig().isServersideBlockConnections()) {
             if (Via.getManager().getProviders().get(BlockConnectionProvider.class) instanceof PacketBlockConnectionProvider) {

common/src/main/java/com/viaversion/viaversion/protocols/v1_12_2to1_13/rewriter/EntityPacketRewriter1_13.java

@@ -28,6 +28,7 @@ import com.viaversion.viaversion.api.type.Types;
 import com.viaversion.viaversion.api.type.types.version.Types1_12;
 import com.viaversion.viaversion.api.type.types.version.Types1_13;
 import com.viaversion.viaversion.protocols.v1_12_2to1_13.Protocol1_12_2To1_13;
+import com.viaversion.viaversion.protocols.v1_12_2to1_13.blockconnections.ConnectionData;
 import com.viaversion.viaversion.protocols.v1_12_2to1_13.data.EntityIdMappings1_13;
 import com.viaversion.viaversion.protocols.v1_12_2to1_13.data.ParticleIdMappings1_13;
 import com.viaversion.viaversion.protocols.v1_12to1_12_1.packet.ClientboundPackets1_12_1;
@@ -132,7 +133,7 @@ public class EntityPacketRewriter1_13 extends EntityRewriter<ClientboundPackets1
                 map(Types.INT); // 2 - Dimension
 
                 handler(wrapper -> {
-                    ClientWorld clientChunks = wrapper.user().get(ClientWorld.class);
+                    ClientWorld clientChunks = wrapper.user().getClientWorld(Protocol1_12_2To1_13.class);
                     int dimensionId = wrapper.get(Types.INT, 1);
                     clientChunks.setEnvironment(dimensionId);
                 });
@@ -141,6 +142,24 @@ public class EntityPacketRewriter1_13 extends EntityRewriter<ClientboundPackets1
             }
         });
 
+        protocol.registerClientbound(ClientboundPackets1_12_1.RESPAWN, new PacketHandlers() {
+            @Override
+            public void register() {
+                map(Types.INT); // 0 - Dimension ID
+                handler(wrapper -> {
+                    ClientWorld clientWorld = wrapper.user().getClientWorld(Protocol1_12_2To1_13.class);
+                    int dimensionId = wrapper.get(Types.INT, 0);
+                    if (clientWorld.setEnvironment(dimensionId)) {
+                        if (Via.getConfig().isServersideBlockConnections()) {
+                            ConnectionData.clearBlockStorage(wrapper.user());
+                        }
+                        tracker(wrapper.user()).clearEntities();
+                    }
+                });
+                handler(Protocol1_12_2To1_13.SEND_DECLARE_COMMANDS_AND_TAGS);
+            }
+        });
+
         protocol.registerClientbound(ClientboundPackets1_12_1.UPDATE_MOB_EFFECT, new PacketHandlers() {
             @Override
             public void register() {

common/src/main/java/com/viaversion/viaversion/protocols/v1_12_2to1_13/rewriter/WorldPacketRewriter1_13.java

@@ -325,7 +325,7 @@ public class WorldPacketRewriter1_13 {
         });
 
         protocol.registerClientbound(ClientboundPackets1_12_1.LEVEL_CHUNK, wrapper -> {
-            ClientWorld clientWorld = wrapper.user().get(ClientWorld.class);
+            ClientWorld clientWorld = wrapper.user().getClientWorld(Protocol1_12_2To1_13.class);
             BlockStorage storage = wrapper.user().get(BlockStorage.class);
 
             ChunkType1_9_3 type = ChunkType1_9_3.forEnvironment(clientWorld.getEnvironment());

common/src/main/java/com/viaversion/viaversion/protocols/v1_13_2to1_14/Protocol1_13_2To1_14.java

@@ -121,9 +121,7 @@ public class Protocol1_13_2To1_14 extends AbstractProtocol<ClientboundPackets1_1
     @Override
     public void init(UserConnection userConnection) {
         userConnection.addEntityTracker(this.getClass(), new EntityTracker1_14(userConnection));
-        if (!userConnection.has(ClientWorld.class)) {
+        userConnection.addClientWorld(this.getClass(), new ClientWorld());
-            userConnection.put(new ClientWorld());
-        }
     }
 
     @Override

common/src/main/java/com/viaversion/viaversion/protocols/v1_13_2to1_14/rewriter/EntityPacketRewriter1_14.java

@@ -204,7 +204,7 @@ public class EntityPacketRewriter1_14 extends EntityRewriter<ClientboundPackets1
                 map(Types.INT); // 2 - Dimension
                 handler(wrapper -> {
                     // Store the player
-                    ClientWorld clientChunks = wrapper.user().get(ClientWorld.class);
+                    ClientWorld clientChunks = wrapper.user().getClientWorld(Protocol1_13_2To1_14.class);
                     int dimensionId = wrapper.get(Types.INT, 1);
                     clientChunks.setEnvironment(dimensionId);
                 });

common/src/main/java/com/viaversion/viaversion/protocols/v1_13_2to1_14/rewriter/WorldPacketRewriter1_14.java

@@ -125,7 +125,7 @@ public class WorldPacketRewriter1_14 {
         });
 
         protocol.registerClientbound(ClientboundPackets1_13.LEVEL_CHUNK, wrapper -> {
-            ClientWorld clientWorld = wrapper.user().get(ClientWorld.class);
+            ClientWorld clientWorld = wrapper.user().getClientWorld(Protocol1_13_2To1_14.class);
             Chunk chunk = wrapper.read(ChunkType1_13.forEnvironment(clientWorld.getEnvironment()));
             wrapper.write(ChunkType1_14.TYPE, chunk);
 
@@ -282,21 +282,26 @@ public class WorldPacketRewriter1_14 {
             public void register() {
                 map(Types.INT); // 0 - Dimension ID
                 handler(wrapper -> {
-                    ClientWorld clientWorld = wrapper.user().get(ClientWorld.class);
+                    short difficulty = wrapper.read(Types.UNSIGNED_BYTE); // 19w11a removed difficulty from respawn
+
+                    ClientWorld clientWorld = wrapper.user().getClientWorld(Protocol1_13_2To1_14.class);
                     int dimensionId = wrapper.get(Types.INT, 0);
-                    clientWorld.setEnvironment(dimensionId);
+
+                    if (!clientWorld.setEnvironment(dimensionId)) {
+                        return;
+                    }
+
                     EntityTracker1_14 entityTracker = wrapper.user().getEntityTracker(Protocol1_13_2To1_14.class);
+                    entityTracker.clearEntities();
+
                     // The client may reset the center chunk if dimension is changed
                     entityTracker.setForceSendCenterChunk(true);
-                });
+
-                handler(wrapper -> {
-                    short difficulty = wrapper.read(Types.UNSIGNED_BYTE); // 19w11a removed difficulty from respawn
                     PacketWrapper difficultyPacket = wrapper.create(ClientboundPackets1_14.CHANGE_DIFFICULTY);
                     difficultyPacket.write(Types.UNSIGNED_BYTE, difficulty);
                     difficultyPacket.write(Types.BOOLEAN, false); // Unknown value added in 19w11a
                     difficultyPacket.scheduleSend(protocol.getClass());
-                });
+
-                handler(wrapper -> {
                     // Manually send the packet and update the viewdistance after
                     wrapper.send(Protocol1_13_2To1_14.class);
                     wrapper.cancel();

common/src/main/java/com/viaversion/viaversion/protocols/v1_13to1_13_1/Protocol1_13To1_13_1.java

@@ -134,9 +134,7 @@ public class Protocol1_13To1_13_1 extends AbstractProtocol<ClientboundPackets1_1
     @Override
     public void init(UserConnection userConnection) {
         userConnection.addEntityTracker(this.getClass(), new EntityTrackerBase(userConnection, EntityTypes1_13.EntityType.PLAYER));
-        if (!userConnection.has(ClientWorld.class)) {
+        userConnection.addClientWorld(this.getClass(), new ClientWorld());
-            userConnection.put(new ClientWorld());
-        }
     }
 
     @Override

common/src/main/java/com/viaversion/viaversion/protocols/v1_13to1_13_1/rewriter/EntityPacketRewriter1_13_1.java

@@ -17,6 +17,7 @@
  */
 package com.viaversion.viaversion.protocols.v1_13to1_13_1.rewriter;
 
+import com.viaversion.viaversion.api.minecraft.ClientWorld;
 import com.viaversion.viaversion.api.minecraft.entities.EntityType;
 import com.viaversion.viaversion.api.minecraft.entities.EntityTypes1_13;
 import com.viaversion.viaversion.api.protocol.remapper.PacketHandlers;
@@ -34,6 +35,36 @@ public class EntityPacketRewriter1_13_1 extends EntityRewriter<ClientboundPacket
 
     @Override
     protected void registerPackets() {
+        protocol.registerClientbound(ClientboundPackets1_13.LOGIN, new PacketHandlers() {
+            @Override
+            public void register() {
+                map(Types.INT); // 0 - Entity ID
+                map(Types.UNSIGNED_BYTE); // 1 - Gamemode
+                map(Types.INT); // 2 - Dimension
+
+                handler(wrapper -> {
+                    // Store the player
+                    ClientWorld clientWorld = wrapper.user().getClientWorld(Protocol1_13To1_13_1.class);
+                    int dimensionId = wrapper.get(Types.INT, 1);
+                    clientWorld.setEnvironment(dimensionId);
+                });
+            }
+        });
+
+        protocol.registerClientbound(ClientboundPackets1_13.RESPAWN, new PacketHandlers() {
+            @Override
+            public void register() {
+                map(Types.INT); // 0 - Dimension ID
+                handler(wrapper -> {
+                    ClientWorld clientWorld = wrapper.user().getClientWorld(Protocol1_13To1_13_1.class);
+                    int dimensionId = wrapper.get(Types.INT, 0);
+                    if (clientWorld.setEnvironment(dimensionId)) {
+                        tracker(wrapper.user()).clearEntities();
+                    }
+                });
+            }
+        });
+
         protocol.registerClientbound(ClientboundPackets1_13.ADD_ENTITY, new PacketHandlers() {
             @Override
             public void register() {

common/src/main/java/com/viaversion/viaversion/protocols/v1_13to1_13_1/rewriter/WorldPacketRewriter1_13_1.java

@@ -19,9 +19,6 @@ package com.viaversion.viaversion.protocols.v1_13to1_13_1.rewriter;
 
 import com.viaversion.viaversion.api.minecraft.ClientWorld;
 import com.viaversion.viaversion.api.minecraft.chunks.Chunk;
-import com.viaversion.viaversion.api.minecraft.chunks.ChunkSection;
-import com.viaversion.viaversion.api.minecraft.chunks.DataPalette;
-import com.viaversion.viaversion.api.minecraft.chunks.PaletteType;
 import com.viaversion.viaversion.api.protocol.remapper.PacketHandlers;
 import com.viaversion.viaversion.api.type.Types;
 import com.viaversion.viaversion.api.type.types.chunk.ChunkType1_13;
@@ -35,7 +32,7 @@ public class WorldPacketRewriter1_13_1 {
         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_13To1_13_1.class);
             Chunk chunk = wrapper.passthrough(ChunkType1_13.forEnvironment(clientWorld.getEnvironment()));
 
             blockRewriter.handleChunk(chunk);
@@ -84,33 +81,5 @@ public class WorldPacketRewriter1_13_1 {
                 });
             }
         });
-
-        protocol.registerClientbound(ClientboundPackets1_13.LOGIN, new PacketHandlers() {
-            @Override
-            public void register() {
-                map(Types.INT); // 0 - Entity ID
-                map(Types.UNSIGNED_BYTE); // 1 - Gamemode
-                map(Types.INT); // 2 - Dimension
-
-                handler(wrapper -> {
-                    // Store the player
-                    ClientWorld clientChunks = wrapper.user().get(ClientWorld.class);
-                    int dimensionId = wrapper.get(Types.INT, 1);
-                    clientChunks.setEnvironment(dimensionId);
-                });
-            }
-        });
-
-        protocol.registerClientbound(ClientboundPackets1_13.RESPAWN, new PacketHandlers() {
-            @Override
-            public void register() {
-                map(Types.INT); // 0 - Dimension ID
-                handler(wrapper -> {
-                    ClientWorld clientWorld = wrapper.user().get(ClientWorld.class);
-                    int dimensionId = wrapper.get(Types.INT, 0);
-                    clientWorld.setEnvironment(dimensionId);
-                });
-            }
-        });
     }
 }

common/src/main/java/com/viaversion/viaversion/protocols/v1_14_4to1_15/rewriter/EntityPacketRewriter1_15.java

@@ -86,7 +86,10 @@ public class EntityPacketRewriter1_15 extends EntityRewriter<ClientboundPackets1
             @Override
             public void register() {
                 map(Types.INT);
-                handler(wrapper -> wrapper.write(Types.LONG, 0L)); // Level Seed
+                handler(wrapper -> {
+                    tracker(wrapper.user()).clearEntities();
+                    wrapper.write(Types.LONG, 0L); // Level Seed
+                });
             }
         });
 

common/src/main/java/com/viaversion/viaversion/protocols/v1_15_2to1_16/rewriter/EntityPacketRewriter1_16.java

@@ -119,6 +119,8 @@ public class EntityPacketRewriter1_16 extends EntityRewriter<ClientboundPackets1
                 map(Types.LONG); // Seed
                 map(Types.UNSIGNED_BYTE); // Gamemode
                 handler(wrapper -> {
+                    tracker(wrapper.user()).clearEntities();
+
                     wrapper.write(Types.BYTE, (byte) -1); // Previous gamemode, set to none
 
                     // <= 1.14.4 didn't keep attributes on respawn and 1.15.x always kept them
@@ -145,9 +147,8 @@ public class EntityPacketRewriter1_16 extends EntityRewriter<ClientboundPackets1
                 handler(DIMENSION_HANDLER); // Dimension
                 map(Types.LONG); // Seed
                 map(Types.UNSIGNED_BYTE); // Max players
+                handler(playerTrackerHandler());
                 handler(wrapper -> {
-                    wrapper.user().getEntityTracker(Protocol1_15_2To1_16.class).addEntity(wrapper.get(Types.INT, 0), EntityTypes1_16.PLAYER);
-
                     final String type = wrapper.read(Types.STRING);// level type
                     wrapper.passthrough(Types.VAR_INT); // View distance
                     wrapper.passthrough(Types.BOOLEAN); // Reduced debug info

common/src/main/java/com/viaversion/viaversion/protocols/v1_16_1to1_16_2/rewriter/EntityPacketRewriter1_16_2.java

@@ -72,6 +72,8 @@ public class EntityPacketRewriter1_16_2 extends EntityRewriter<ClientboundPacket
         protocol.registerClientbound(ClientboundPackets1_16.RESPAWN, wrapper -> {
             String dimensionType = wrapper.read(Types.STRING);
             wrapper.write(Types.NAMED_COMPOUND_TAG, getDimensionData(dimensionType));
+
+            tracker(wrapper.user()).clearEntities();
         });
     }
 

common/src/main/java/com/viaversion/viaversion/protocols/v1_16_4to1_17/rewriter/EntityPacketRewriter1_17.java

@@ -116,6 +116,8 @@ public final class EntityPacketRewriter1_17 extends EntityRewriter<ClientboundPa
         protocol.registerClientbound(ClientboundPackets1_16_2.RESPAWN, wrapper -> {
             CompoundTag dimensionData = wrapper.passthrough(Types.NAMED_COMPOUND_TAG);
             addNewDimensionData(dimensionData);
+
+            tracker(wrapper.user()).clearEntities();
         });
 
         protocol.registerClientbound(ClientboundPackets1_16_2.UPDATE_ATTRIBUTES, new PacketHandlers() {

common/src/main/java/com/viaversion/viaversion/protocols/v1_8to1_9/Protocol1_8To1_9.java

@@ -43,7 +43,7 @@ import com.viaversion.viaversion.protocols.v1_8to1_9.rewriter.ItemPacketRewriter
 import com.viaversion.viaversion.protocols.v1_8to1_9.rewriter.PlayerPacketRewriter1_9;
 import com.viaversion.viaversion.protocols.v1_8to1_9.rewriter.SpawnPacketRewriter1_9;
 import com.viaversion.viaversion.protocols.v1_8to1_9.rewriter.WorldPacketRewriter1_9;
-import com.viaversion.viaversion.protocols.v1_8to1_9.storage.ClientChunks;
+import com.viaversion.viaversion.protocols.v1_8to1_9.storage.ClientWorld1_9;
 import com.viaversion.viaversion.protocols.v1_8to1_9.storage.CommandBlockStorage;
 import com.viaversion.viaversion.protocols.v1_8to1_9.storage.EntityTracker1_9;
 import com.viaversion.viaversion.protocols.v1_8to1_9.storage.InventoryTracker;
@@ -104,20 +104,12 @@ public class Protocol1_8To1_9 extends AbstractProtocol<ClientboundPackets1_8, Cl
 
     @Override
     public void init(UserConnection userConnection) {
-        // Entity tracker
         userConnection.addEntityTracker(this.getClass(), new EntityTracker1_9(userConnection));
-        // Chunk tracker
+        userConnection.addClientWorld(this.getClass(), new ClientWorld1_9());
-        userConnection.put(new ClientChunks());
-        // Movement tracker
-        userConnection.put(new MovementTracker());
-        // Inventory tracker
-        userConnection.put(new InventoryTracker());
-        // CommandBlock storage
-        userConnection.put(new CommandBlockStorage());
 
-        if (!userConnection.has(ClientWorld.class)) {
+        userConnection.put(new MovementTracker());
-            userConnection.put(new ClientWorld());
+        userConnection.put(new InventoryTracker());
-        }
+        userConnection.put(new CommandBlockStorage());
     }
 
     @Override

common/src/main/java/com/viaversion/viaversion/protocols/v1_8to1_9/rewriter/PlayerPacketRewriter1_9.java

@@ -36,7 +36,7 @@ import com.viaversion.viaversion.protocols.v1_8to1_9.packet.ServerboundPackets1_
 import com.viaversion.viaversion.protocols.v1_8to1_9.provider.CommandBlockProvider;
 import com.viaversion.viaversion.protocols.v1_8to1_9.provider.CompressionProvider;
 import com.viaversion.viaversion.protocols.v1_8to1_9.provider.MainHandProvider;
-import com.viaversion.viaversion.protocols.v1_8to1_9.storage.ClientChunks;
+import com.viaversion.viaversion.protocols.v1_8to1_9.storage.ClientWorld1_9;
 import com.viaversion.viaversion.protocols.v1_8to1_9.storage.EntityTracker1_9;
 import com.viaversion.viaversion.protocols.v1_8to1_9.storage.MovementTracker;
 import com.viaversion.viaversion.util.ComponentUtil;
@@ -203,7 +203,7 @@ public class PlayerPacketRewriter1_9 {
 
                 // Track player's dimension
                 handler(wrapper -> {
-                    ClientWorld clientWorld = wrapper.user().get(ClientWorld.class);
+                    ClientWorld clientWorld = wrapper.user().getClientWorld(Protocol1_8To1_9.class);
                     int dimensionId = wrapper.get(Types.BYTE, 0);
                     clientWorld.setEnvironment(dimensionId);
                 });
@@ -295,27 +295,25 @@ public class PlayerPacketRewriter1_9 {
                 map(Types.UNSIGNED_BYTE); // 2 - GameMode
                 map(Types.STRING); // 3 - Level Type
 
-                // Track player's dimension
-                handler(wrapper -> {
-                    ClientWorld clientWorld = wrapper.user().get(ClientWorld.class);
-                    int dimensionId = wrapper.get(Types.INT, 0);
-                    clientWorld.setEnvironment(dimensionId);
-                });
-
-                handler(wrapper -> {
-                    // Client unloads chunks on respawn
-                    wrapper.user().get(ClientChunks.class).getLoadedChunks().clear();
-
-                    int gamemode = wrapper.get(Types.UNSIGNED_BYTE, 0);
-                    EntityTracker1_9 tracker = wrapper.user().getEntityTracker(Protocol1_8To1_9.class);
-                    tracker.setGameMode(GameMode.getById(gamemode));
-                });
-
-                // Fake permissions to get Commandblocks working
                 handler(wrapper -> {
                     CommandBlockProvider provider = Via.getManager().getProviders().get(CommandBlockProvider.class);
+                    // Fake permissions to get Commandblocks working
                     provider.sendPermission(wrapper.user());
+
+                    EntityTracker1_9 tracker = wrapper.user().getEntityTracker(Protocol1_8To1_9.class);
+                    int gamemode = wrapper.get(Types.UNSIGNED_BYTE, 0);
+                    tracker.setGameMode(GameMode.getById(gamemode));
+
+                    ClientWorld1_9 clientWorld = wrapper.user().getClientWorld(Protocol1_8To1_9.class);
+                    int dimensionId = wrapper.get(Types.INT, 0);
+
+                    // Track player's dimension
+                    if (clientWorld.setEnvironment(dimensionId)) {
+                        tracker.clearEntities();
+
+                        clientWorld.getLoadedChunks().clear();
                         provider.unloadChunks(wrapper.user());
+                    }
                 });
             }
         });

common/src/main/java/com/viaversion/viaversion/protocols/v1_8to1_9/rewriter/WorldPacketRewriter1_9.java

@@ -45,7 +45,7 @@ import com.viaversion.viaversion.protocols.v1_8to1_9.packet.ServerboundPackets1_
 import com.viaversion.viaversion.protocols.v1_8to1_9.packet.ServerboundPackets1_9;
 import com.viaversion.viaversion.protocols.v1_8to1_9.provider.CommandBlockProvider;
 import com.viaversion.viaversion.protocols.v1_8to1_9.provider.HandItemProvider;
-import com.viaversion.viaversion.protocols.v1_8to1_9.storage.ClientChunks;
+import com.viaversion.viaversion.protocols.v1_8to1_9.storage.ClientWorld1_9;
 import com.viaversion.viaversion.protocols.v1_8to1_9.storage.EntityTracker1_9;
 import com.viaversion.viaversion.util.ComponentUtil;
 import com.viaversion.viaversion.util.Key;
@@ -130,11 +130,10 @@ public class WorldPacketRewriter1_9 {
         });
 
         protocol.registerClientbound(ClientboundPackets1_8.LEVEL_CHUNK, wrapper -> {
-            ClientWorld clientWorld = wrapper.user().get(ClientWorld.class);
+            ClientWorld1_9 clientWorld = wrapper.user().getClientWorld(Protocol1_8To1_9.class);
-            ClientChunks clientChunks = wrapper.user().get(ClientChunks.class);
             Chunk chunk = wrapper.read(ChunkType1_8.forEnvironment(clientWorld.getEnvironment()));
 
-            long chunkHash = ClientChunks.toLong(chunk.getX(), chunk.getZ());
+            long chunkHash = ClientWorld1_9.toLong(chunk.getX(), chunk.getZ());
 
             // Check if the chunk should be handled as an unload packet
             if (chunk.isFullChunk() && chunk.getBitmask() == 0) {
@@ -146,14 +145,14 @@ public class WorldPacketRewriter1_9 {
                 CommandBlockProvider provider = Via.getManager().getProviders().get(CommandBlockProvider.class);
                 provider.unloadChunk(wrapper.user(), chunk.getX(), chunk.getZ());
 
-                clientChunks.getLoadedChunks().remove(chunkHash);
+                clientWorld.getLoadedChunks().remove(chunkHash);
 
                 // Unload the empty chunks
                 if (Via.getConfig().isChunkBorderFix()) {
                     for (BlockFace face : BlockFace.HORIZONTAL) {
                         int chunkX = chunk.getX() + face.modX();
                         int chunkZ = chunk.getZ() + face.modZ();
-                        if (!clientChunks.getLoadedChunks().contains(ClientChunks.toLong(chunkX, chunkZ))) {
+                        if (!clientWorld.getLoadedChunks().contains(ClientWorld1_9.toLong(chunkX, chunkZ))) {
                             PacketWrapper unloadChunk = wrapper.create(ClientboundPackets1_9.FORGET_LEVEL_CHUNK);
                             unloadChunk.write(Types.INT, chunkX);
                             unloadChunk.write(Types.INT, chunkZ);
@@ -165,14 +164,14 @@ public class WorldPacketRewriter1_9 {
                 Type<Chunk> chunkType = ChunkType1_9_1.forEnvironment(clientWorld.getEnvironment());
                 wrapper.write(chunkType, chunk);
 
-                clientChunks.getLoadedChunks().add(chunkHash);
+                clientWorld.getLoadedChunks().add(chunkHash);
 
                 // Send empty chunks surrounding the loaded chunk to force 1.9+ clients to render the new chunk
                 if (Via.getConfig().isChunkBorderFix()) {
                     for (BlockFace face : BlockFace.HORIZONTAL) {
                         int chunkX = chunk.getX() + face.modX();
                         int chunkZ = chunk.getZ() + face.modZ();
-                        if (!clientChunks.getLoadedChunks().contains(ClientChunks.toLong(chunkX, chunkZ))) {
+                        if (!clientWorld.getLoadedChunks().contains(ClientWorld1_9.toLong(chunkX, chunkZ))) {
                             PacketWrapper emptyChunk = wrapper.create(ClientboundPackets1_9.LEVEL_CHUNK);
                             Chunk c = new BaseChunk(chunkX, chunkZ, true, false, 0, new ChunkSection[16], new int[256], new ArrayList<>());
                             emptyChunk.write(chunkType, c);
@@ -185,8 +184,7 @@ public class WorldPacketRewriter1_9 {
 
         protocol.registerClientbound(ClientboundPackets1_8.MAP_BULK_CHUNK, null, wrapper -> {
             wrapper.cancel(); // Cancel the packet from being sent
-            ClientWorld clientWorld = wrapper.user().get(ClientWorld.class);
+            ClientWorld1_9 clientWorld = wrapper.user().getClientWorld(Protocol1_8To1_9.class);
-            ClientChunks clientChunks = wrapper.user().get(ClientChunks.class);
             Chunk[] chunks = wrapper.read(BulkChunkType1_8.TYPE);
 
             Type<Chunk> chunkType = ChunkType1_9_1.forEnvironment(clientWorld.getEnvironment());
@@ -196,14 +194,14 @@ public class WorldPacketRewriter1_9 {
                 chunkData.write(chunkType, chunk);
                 chunkData.send(Protocol1_8To1_9.class);
 
-                clientChunks.getLoadedChunks().add(ClientChunks.toLong(chunk.getX(), chunk.getZ()));
+                clientWorld.getLoadedChunks().add(ClientWorld1_9.toLong(chunk.getX(), chunk.getZ()));
 
                 // Send empty chunks surrounding the loaded chunk to force 1.9+ clients to render the new chunk
                 if (Via.getConfig().isChunkBorderFix()) {
                     for (BlockFace face : BlockFace.HORIZONTAL) {
                         int chunkX = chunk.getX() + face.modX();
                         int chunkZ = chunk.getZ() + face.modZ();
-                        if (!clientChunks.getLoadedChunks().contains(ClientChunks.toLong(chunkX, chunkZ))) {
+                        if (!clientWorld.getLoadedChunks().contains(ClientWorld1_9.toLong(chunkX, chunkZ))) {
                             PacketWrapper emptyChunk = wrapper.create(ClientboundPackets1_9.LEVEL_CHUNK);
                             Chunk c = new BaseChunk(chunkX, chunkZ, true, false, 0, new ChunkSection[16], new int[256], new ArrayList<>());
                             emptyChunk.write(chunkType, c);

common/src/main/java/com/viaversion/viaversion/protocols/v1_8to1_9/storage/ClientWorld1_9.java

@@ -18,10 +18,10 @@
 package com.viaversion.viaversion.protocols.v1_8to1_9.storage;
 
 import com.google.common.collect.Sets;
-import com.viaversion.viaversion.api.connection.StorableObject;
+import com.viaversion.viaversion.api.minecraft.ClientWorld;
 import java.util.Set;
 
-public class ClientChunks implements StorableObject {
+public class ClientWorld1_9 extends ClientWorld {
     private final Set<Long> loadedChunks = Sets.newConcurrentHashSet();
 
     public static long toLong(int msw, int lsw) {

common/src/main/java/com/viaversion/viaversion/protocols/v1_9_1to1_9_3/Protocol1_9_1To1_9_3.java

@@ -87,7 +87,7 @@ public class Protocol1_9_1To1_9_3 extends AbstractProtocol<ClientboundPackets1_9
         });
 
         registerClientbound(ClientboundPackets1_9.LEVEL_CHUNK, wrapper -> {
-            ClientWorld clientWorld = wrapper.user().get(ClientWorld.class);
+            ClientWorld clientWorld = wrapper.user().getClientWorld(Protocol1_9_1To1_9_3.class);
 
             Chunk chunk = wrapper.read(ChunkType1_9_1.forEnvironment(clientWorld.getEnvironment()));
             wrapper.write(ChunkType1_9_3.forEnvironment(clientWorld.getEnvironment()), chunk);
@@ -120,7 +120,7 @@ public class Protocol1_9_1To1_9_3 extends AbstractProtocol<ClientboundPackets1_9
                 map(Types.INT); // 2 - Dimension
 
                 handler(wrapper -> {
-                    ClientWorld clientWorld = wrapper.user().get(ClientWorld.class);
+                    ClientWorld clientWorld = wrapper.user().getClientWorld(Protocol1_9_1To1_9_3.class);
                     int dimensionId = wrapper.get(Types.INT, 1);
                     clientWorld.setEnvironment(dimensionId);
                 });
@@ -132,7 +132,7 @@ public class Protocol1_9_1To1_9_3 extends AbstractProtocol<ClientboundPackets1_9
             public void register() {
                 map(Types.INT); // 0 - Dimension ID
                 handler(wrapper -> {
-                    ClientWorld clientWorld = wrapper.user().get(ClientWorld.class);
+                    ClientWorld clientWorld = wrapper.user().getClientWorld(Protocol1_9_1To1_9_3.class);
                     int dimensionId = wrapper.get(Types.INT, 0);
                     clientWorld.setEnvironment(dimensionId);
                 });
@@ -156,8 +156,6 @@ public class Protocol1_9_1To1_9_3 extends AbstractProtocol<ClientboundPackets1_9
 
     @Override
     public void init(UserConnection user) {
-        if (!user.has(ClientWorld.class)) {
+        user.addClientWorld(this.getClass(), new ClientWorld());
-            user.put(new ClientWorld());
-        }
     }
 }

common/src/main/java/com/viaversion/viaversion/protocols/v1_9_3to1_10/Protocol1_9_3To1_10.java

@@ -153,7 +153,7 @@ public class Protocol1_9_3To1_10 extends AbstractProtocol<ClientboundPackets1_9_
                 map(Types.INT); // 2 - Dimension
 
                 handler(wrapper -> {
-                    ClientWorld clientWorld = wrapper.user().get(ClientWorld.class);
+                    ClientWorld clientWorld = wrapper.user().getClientWorld(Protocol1_9_3To1_10.class);
 
                     int dimensionId = wrapper.get(Types.INT, 1);
                     clientWorld.setEnvironment(dimensionId);
@@ -168,7 +168,7 @@ public class Protocol1_9_3To1_10 extends AbstractProtocol<ClientboundPackets1_9_
                 map(Types.INT); // 0 - Dimension ID
 
                 handler(wrapper -> {
-                    ClientWorld clientWorld = wrapper.user().get(ClientWorld.class);
+                    ClientWorld clientWorld = wrapper.user().getClientWorld(Protocol1_9_3To1_10.class);
 
                     int dimensionId = wrapper.get(Types.INT, 0);
                     clientWorld.setEnvironment(dimensionId);
@@ -178,7 +178,7 @@ public class Protocol1_9_3To1_10 extends AbstractProtocol<ClientboundPackets1_9_
 
         // Chunk Data
         registerClientbound(ClientboundPackets1_9_3.LEVEL_CHUNK, wrapper -> {
-            ClientWorld clientWorld = wrapper.user().get(ClientWorld.class);
+            ClientWorld clientWorld = wrapper.user().getClientWorld(Protocol1_9_3To1_10.class);
             Chunk chunk = wrapper.passthrough(ChunkType1_9_3.forEnvironment(clientWorld.getEnvironment()));
 
             if (Via.getConfig().isReplacePistons()) {
@@ -234,11 +234,9 @@ public class Protocol1_9_3To1_10 extends AbstractProtocol<ClientboundPackets1_9_
 
     @Override
     public void init(UserConnection userConnection) {
-        userConnection.put(new ResourcePackTracker());
+        userConnection.addClientWorld(this.getClass(), new ClientWorld());
 
-        if (!userConnection.has(ClientWorld.class)) {
+        userConnection.put(new ResourcePackTracker());
-            userConnection.put(new ClientWorld());
-        }
     }
 
     @Override

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

@@ -367,6 +367,11 @@ public abstract class EntityRewriter<C extends ClientboundPacketType, T extends
         registerSetEntityData(packetType, null, dataType);
     }
 
+    public void clearEntities(final UserConnection connection) {
+        final EntityTracker tracker = tracker(connection);
+        tracker.clearEntities();
+    }
+
     public PacketHandler trackerHandler() {
         return trackerAndRewriterHandler(null);
     }
@@ -410,7 +415,6 @@ public abstract class EntityRewriter<C extends ClientboundPacketType, T extends
             String world = wrapper.get(Types.STRING, 0);
             if (tracker.currentWorld() != null && !tracker.currentWorld().equals(world)) {
                 tracker.clearEntities();
-                tracker.trackClientEntity();
             }
             tracker.setCurrentWorld(world);
         };
@@ -433,7 +437,6 @@ public abstract class EntityRewriter<C extends ClientboundPacketType, T extends
             String world = wrapper.get(Types.STRING, 1);
             if (tracker.currentWorld() != null && !tracker.currentWorld().equals(world)) {
                 tracker.clearEntities();
-                tracker.trackClientEntity();
             }
             tracker.setCurrentWorld(world);
         };
@@ -462,7 +465,6 @@ public abstract class EntityRewriter<C extends ClientboundPacketType, T extends
         // Clear entities if the world changes
         if (tracker.currentWorld() != null && !tracker.currentWorld().equals(world)) {
             tracker.clearEntities();
-            tracker.trackClientEntity();
         }
         tracker.setCurrentWorld(world);
     }