From 5fca206b8dcdace56fe87506b44165237ae8368a Mon Sep 17 00:00:00 2001
From: KennyTV <kennytv@t-online.de>
Date: Sun, 7 Jun 2020 10:14:18 +0200
Subject: [PATCH] Add method to get protocol instance by class

---
 .../api/protocol/ProtocolRegistry.java        | 20 ++++++++++++++++---
 1 file changed, 17 insertions(+), 3 deletions(-)

diff --git a/common/src/main/java/us/myles/ViaVersion/api/protocol/ProtocolRegistry.java b/common/src/main/java/us/myles/ViaVersion/api/protocol/ProtocolRegistry.java
index 8ef9a3276..50331e45c 100644
--- a/common/src/main/java/us/myles/ViaVersion/api/protocol/ProtocolRegistry.java
+++ b/common/src/main/java/us/myles/ViaVersion/api/protocol/ProtocolRegistry.java
@@ -53,10 +53,11 @@ public class ProtocolRegistry {
     public static int SERVER_PROTOCOL = -1;
     // Input Version -> Output Version & Protocol (Allows fast lookup)
     private static final Map<Integer, Map<Integer, Protocol>> registryMap = new ConcurrentHashMap<>();
+    private static final Map<Class<? extends Protocol>, Protocol> protocols = new HashMap<>();
     private static final Map<Pair<Integer, Integer>, List<Pair<Integer, Protocol>>> pathCache = new ConcurrentHashMap<>();
-    private static final List<Protocol> registerList = new ArrayList<>();
     private static final Set<Integer> supportedVersions = new HashSet<>();
     private static final List<Pair<Range<Integer>, Protocol>> baseProtocols = Lists.newCopyOnWriteArrayList();
+    private static final List<Protocol> registerList = new ArrayList<>();
 
     private static final Object MAPPING_LOADER_LOCK = new Object();
     private static Map<Class<? extends Protocol>, CompletableFuture<Void>> mappingLoaderFutures = new HashMap<>();
@@ -131,7 +132,9 @@ public class ProtocolRegistry {
             pathCache.clear();
         }
 
-        for (Integer version : supported) {
+        protocols.put(protocol.getClass(), protocol);
+
+        for (int version : supported) {
             Map<Integer, Protocol> protocolMap = registryMap.computeIfAbsent(version, k -> new HashMap<>());
             protocolMap.put(output, protocol);
         }
@@ -269,7 +272,7 @@ public class ProtocolRegistry {
     }
 
     /**
-     * Calculate a path from a client version to server version
+     * Calculate a path from a client version to server version.
      *
      * @param clientVersion The input client version
      * @param serverVersion The desired output server version
@@ -292,6 +295,17 @@ public class ProtocolRegistry {
         return outputPath;
     }
 
+    /**
+     * Returns a protocol instance by its class.
+     *
+     * @param protocolClass class of the protocol
+     * @return protocol if present
+     */
+    @Nullable
+    public static Protocol getProtocol(Class<? extends Protocol> protocolClass) {
+        return protocols.get(protocolClass);
+    }
+
     public static Protocol getBaseProtocol(int serverVersion) {
         for (Pair<Range<Integer>, Protocol> rangeProtocol : Lists.reverse(baseProtocols)) {
             if (rangeProtocol.getKey().contains(serverVersion)) {