Implement getNewEntityTracker for Folia (#2464)

dmulloy2 · Jul 3, 2023 · 12d8141 · 12d8141
1 parent 2c48b1c
commit 12d8141
Showing 1 changed file with 21 additions and 1 deletion.
diff --git a/src/main/java/com/comphenix/protocol/injector/EntityUtilities.java b/src/main/java/com/comphenix/protocol/injector/EntityUtilities.java
@@ -27,14 +27,14 @@
 import com.comphenix.protocol.utility.MinecraftFields;
 import com.comphenix.protocol.utility.MinecraftReflection;
 import com.comphenix.protocol.utility.MinecraftVersion;
+import com.comphenix.protocol.utility.Util;
 import com.comphenix.protocol.wrappers.WrappedIntHashMap;
 
 import java.lang.reflect.Field;
 import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
@@ -61,6 +61,7 @@ class EntityUtilities {
     private FieldAccessor entityTrackerField;
     private FieldAccessor trackedPlayersField;
     private FieldAccessor trackedEntitiesField;
+    private FieldAccessor foliaTrackerField;
 
     private MethodAccessor getEntity;
     private MethodAccessor getChunkProvider;
@@ -203,6 +204,25 @@ private Collection<?> getTrackedPlayers(Entity entity) {
 
     @SuppressWarnings("unchecked")
     private Object getNewEntityTracker(Object worldServer, int entityId) {
+        if (Util.isUsingFolia()) {
+            if (this.getEntity == null) {
+                Method entityGetter = FuzzyReflection.fromObject(worldServer).getMethodByReturnTypeAndParameters(
+                        "getEntity",
+                        MinecraftReflection.getEntityClass(),
+                        int.class);
+                this.getEntity = Accessors.getMethodAccessor(entityGetter);
+            }
+
+            Object entity = this.getEntity.invoke(worldServer, entityId);
+
+            if (this.foliaTrackerField == null) {
+                this.foliaTrackerField = Accessors.getFieldAccessor(FuzzyReflection.fromClass(entity.getClass(), false)
+                        .getField(FuzzyFieldContract.newBuilder().typeExact(MinecraftReflection.getEntityTrackerClass()).build()));
+            }
+
+            return this.foliaTrackerField.get(entity);
+        }
+
         if (this.getChunkProvider == null) {
             Class<?> chunkProviderClass = MinecraftReflection.getChunkProviderServer();
             this.getChunkProvider = Accessors.getMethodAccessor(FuzzyReflection.fromClass(worldServer.getClass(), false)