From ce8fab7c4483be624a3667c35f2366d73f758bee Mon Sep 17 00:00:00 2001
From: Nassim Jahnke <nassim@njahnke.dev>
Date: Fri, 26 Jan 2024 14:35:19 +0100
Subject: [PATCH] Cleanup config packet registering in AbstractProtocol

---
 .../api/protocol/AbstractProtocol.java        | 56 +++++++++++--------
 .../packet/mapping/PacketIdMapping.java       |  5 ++
 .../packet/mapping/PacketMapping.java         |  7 +++
 .../packet/mapping/PacketTypeMapping.java     |  5 ++
 .../api/protocol/remapper/PacketHandler.java  |  7 +++
 5 files changed, 56 insertions(+), 24 deletions(-)

diff --git a/api/src/main/java/com/viaversion/viaversion/api/protocol/AbstractProtocol.java b/api/src/main/java/com/viaversion/viaversion/api/protocol/AbstractProtocol.java
index b1edff5d6..d27270ce3 100644
--- a/api/src/main/java/com/viaversion/viaversion/api/protocol/AbstractProtocol.java
+++ b/api/src/main/java/com/viaversion/viaversion/api/protocol/AbstractProtocol.java
@@ -125,45 +125,53 @@ public abstract class AbstractProtocol<CU extends ClientboundPacketType, CM exte
         // even if there will be multiple of these handlers
         final SU configurationAcknowledgedPacket = configurationAcknowledgedPacket();
         if (configurationAcknowledgedPacket != null) {
-            registerServerbound(configurationAcknowledgedPacket, setClientStateHandler(State.CONFIGURATION));
+            appendServerbound(configurationAcknowledgedPacket, setClientStateHandler(State.CONFIGURATION));
         }
 
         final CU startConfigurationPacket = startConfigurationPacket();
         if (startConfigurationPacket != null) {
-            registerClientbound(startConfigurationPacket, setServerStateHandler(State.CONFIGURATION));
+            appendClientbound(startConfigurationPacket, setServerStateHandler(State.CONFIGURATION));
         }
 
         final ServerboundPacketType finishConfigurationPacket = serverboundFinishConfigurationPacket();
         if (finishConfigurationPacket != null) {
-            final int id = finishConfigurationPacket.getId();
-            final PacketMapping mapping = serverboundMappings.mappedPacket(State.CONFIGURATION, id); // Use existing handler if present
-            registerServerbound(State.CONFIGURATION, id, id, wrapper -> {
-                if (mapping != null) {
-                    mapping.applyType(wrapper);
-                    if (mapping.handler() != null) {
-                        mapping.handler().handle(wrapper);
-                    }
-                }
-                setClientStateHandler(State.PLAY).handle(wrapper);
-            }, true);
+            appendServerbound(finishConfigurationPacket, setClientStateHandler(State.PLAY));
         }
 
         final ClientboundPacketType clientboundFinishConfigurationPacket = clientboundFinishConfigurationPacket();
         if (clientboundFinishConfigurationPacket != null) {
-            final int id = clientboundFinishConfigurationPacket.getId();
-            final PacketMapping mapping = clientboundMappings.mappedPacket(State.CONFIGURATION, id); // Use existing handler if present
-            registerClientbound(State.CONFIGURATION, id, id, wrapper -> {
-                if (mapping != null) {
-                    mapping.applyType(wrapper);
-                    if (mapping.handler() != null) {
-                        mapping.handler().handle(wrapper);
-                    }
-                }
-                setServerStateHandler(State.PLAY).handle(wrapper);
-            }, true);
+            appendClientbound(clientboundFinishConfigurationPacket, setServerStateHandler(State.PLAY));
         }
     }
 
+    private void appendClientbound(final ClientboundPacketType type, final PacketHandler handler) {
+        final PacketMapping mapping = clientboundMappings.mappedPacket(type.state(), type.getId()); // Use existing handler if present
+        final PacketHandler newHandler;
+        final int mappedPacketId;
+        if (mapping != null) {
+            newHandler = mapping.handler().append(handler);
+            mappedPacketId = mapping.mappedPacketId() != null ? mapping.mappedPacketId() : type.getId();
+        } else {
+            newHandler = handler;
+            mappedPacketId = type.getId();
+        }
+        registerClientbound(type.state(), type.getId(), mappedPacketId, newHandler, true);
+    }
+
+    private void appendServerbound(final ServerboundPacketType type, final PacketHandler handler) {
+        final PacketMapping mapping = serverboundMappings.mappedPacket(type.state(), type.getId()); // Use existing handler if present
+        final PacketHandler newHandler;
+        final int mappedPacketId;
+        if (mapping != null) {
+            newHandler = mapping.handler().append(handler);
+            mappedPacketId = mapping.mappedPacketId() != null ? mapping.mappedPacketId() : type.getId();
+        } else {
+            newHandler = handler;
+            mappedPacketId = type.getId();
+        }
+        registerServerbound(type.state(), type.getId(), mappedPacketId, newHandler, true);
+    }
+
     private <U extends PacketType, M extends PacketType> void registerPacketIdChanges(
             Map<State, PacketTypeMap<U>> unmappedPacketTypes,
             Map<State, PacketTypeMap<M>> mappedPacketTypes,
diff --git a/api/src/main/java/com/viaversion/viaversion/api/protocol/packet/mapping/PacketIdMapping.java b/api/src/main/java/com/viaversion/viaversion/api/protocol/packet/mapping/PacketIdMapping.java
index fa80099d4..4b2875318 100644
--- a/api/src/main/java/com/viaversion/viaversion/api/protocol/packet/mapping/PacketIdMapping.java
+++ b/api/src/main/java/com/viaversion/viaversion/api/protocol/packet/mapping/PacketIdMapping.java
@@ -41,6 +41,11 @@ final class PacketIdMapping implements PacketMapping {
         wrapper.setId(mappedPacketId);
     }
 
+    @Override
+    public Integer mappedPacketId() {
+        return mappedPacketId;
+    }
+
     @Override
     public @Nullable PacketHandler handler() {
         return handler;
diff --git a/api/src/main/java/com/viaversion/viaversion/api/protocol/packet/mapping/PacketMapping.java b/api/src/main/java/com/viaversion/viaversion/api/protocol/packet/mapping/PacketMapping.java
index 590f8d4c6..6d4880213 100644
--- a/api/src/main/java/com/viaversion/viaversion/api/protocol/packet/mapping/PacketMapping.java
+++ b/api/src/main/java/com/viaversion/viaversion/api/protocol/packet/mapping/PacketMapping.java
@@ -39,6 +39,13 @@ public interface PacketMapping {
      */
     void applyType(PacketWrapper wrapper);
 
+    /**
+     * Returns the mapped packet id if present.
+     *
+     * @return mapped packet type, or null if no action has to be taken
+     */
+    @Nullable Integer mappedPacketId();
+
     /**
      * Returns a packet transformer to transform a packet from one protocol version to another.
      *
diff --git a/api/src/main/java/com/viaversion/viaversion/api/protocol/packet/mapping/PacketTypeMapping.java b/api/src/main/java/com/viaversion/viaversion/api/protocol/packet/mapping/PacketTypeMapping.java
index e467b1965..b4f33c8d3 100644
--- a/api/src/main/java/com/viaversion/viaversion/api/protocol/packet/mapping/PacketTypeMapping.java
+++ b/api/src/main/java/com/viaversion/viaversion/api/protocol/packet/mapping/PacketTypeMapping.java
@@ -43,6 +43,11 @@ final class PacketTypeMapping implements PacketMapping {
         }
     }
 
+    @Override
+    public @Nullable Integer mappedPacketId() {
+        return mappedPacketType != null ? mappedPacketType.getId() : null;
+    }
+
     @Override
     public @Nullable PacketHandler handler() {
         return handler;
diff --git a/api/src/main/java/com/viaversion/viaversion/api/protocol/remapper/PacketHandler.java b/api/src/main/java/com/viaversion/viaversion/api/protocol/remapper/PacketHandler.java
index fa0acec1a..4322c7df8 100644
--- a/api/src/main/java/com/viaversion/viaversion/api/protocol/remapper/PacketHandler.java
+++ b/api/src/main/java/com/viaversion/viaversion/api/protocol/remapper/PacketHandler.java
@@ -34,4 +34,11 @@ public interface PacketHandler {
      * @throws Exception if an error occurs during the packet handling
      */
     void handle(PacketWrapper wrapper) throws Exception;
+
+    default PacketHandler append(final PacketHandler handler) {
+        return wrapper -> {
+            this.handle(wrapper);
+            handler.handle(wrapper);
+        };
+    }
 }