From b94303d2be2c2afc3ec161865dd76854f8f1f94a Mon Sep 17 00:00:00 2001
From: Andrew Steinborn <git@steinborn.me>
Date: Tue, 16 Feb 2021 01:14:22 -0500
Subject: [PATCH] Map command nodes being serialized by their identities

Fixes #429
---
 .../protocol/packet/AvailableCommands.java    |  5 +++-
 .../util/collect/IdentityHashStrategy.java    | 24 +++++++++++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)
 create mode 100644 proxy/src/main/java/com/velocitypowered/proxy/util/collect/IdentityHashStrategy.java

diff --git a/proxy/src/main/java/com/velocitypowered/proxy/protocol/packet/AvailableCommands.java b/proxy/src/main/java/com/velocitypowered/proxy/protocol/packet/AvailableCommands.java
index 891115511..90a9d3bc2 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/protocol/packet/AvailableCommands.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/protocol/packet/AvailableCommands.java
@@ -23,7 +23,9 @@ import com.velocitypowered.proxy.protocol.MinecraftPacket;
 import com.velocitypowered.proxy.protocol.ProtocolUtils;
 import com.velocitypowered.proxy.protocol.ProtocolUtils.Direction;
 import com.velocitypowered.proxy.protocol.packet.brigadier.ArgumentPropertyRegistry;
+import com.velocitypowered.proxy.util.collect.IdentityHashStrategy;
 import io.netty.buffer.ByteBuf;
+import it.unimi.dsi.fastutil.objects.Object2IntLinkedOpenCustomHashMap;
 import it.unimi.dsi.fastutil.objects.Object2IntLinkedOpenHashMap;
 import it.unimi.dsi.fastutil.objects.Object2IntMap;
 import java.util.ArrayDeque;
@@ -96,7 +98,8 @@ public class AvailableCommands implements MinecraftPacket {
   public void encode(ByteBuf buf, Direction direction, ProtocolVersion protocolVersion) {
     // Assign all the children an index.
     Deque<CommandNode<CommandSource>> childrenQueue = new ArrayDeque<>(ImmutableList.of(rootNode));
-    Object2IntMap<CommandNode<CommandSource>> idMappings = new Object2IntLinkedOpenHashMap<>();
+    Object2IntMap<CommandNode<CommandSource>> idMappings = new Object2IntLinkedOpenCustomHashMap<>(
+        IdentityHashStrategy.instance());
     while (!childrenQueue.isEmpty()) {
       CommandNode<CommandSource> child = childrenQueue.poll();
       if (!idMappings.containsKey(child)) {
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/util/collect/IdentityHashStrategy.java b/proxy/src/main/java/com/velocitypowered/proxy/util/collect/IdentityHashStrategy.java
new file mode 100644
index 000000000..b741d3ac5
--- /dev/null
+++ b/proxy/src/main/java/com/velocitypowered/proxy/util/collect/IdentityHashStrategy.java
@@ -0,0 +1,24 @@
+package com.velocitypowered.proxy.util.collect;
+
+import it.unimi.dsi.fastutil.Hash.Strategy;
+
+public final class IdentityHashStrategy<T> implements Strategy<T> {
+
+  @SuppressWarnings("rawtypes")
+  private static final IdentityHashStrategy INSTANCE = new IdentityHashStrategy();
+
+  public static <T> Strategy<T> instance() {
+    //noinspection unchecked
+    return INSTANCE;
+  }
+
+  @Override
+  public int hashCode(T o) {
+    return System.identityHashCode(o);
+  }
+
+  @Override
+  public boolean equals(T a, T b) {
+    return a == b;
+  }
+}