From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: Samsuik <kfian294ma4@gmail.com>
Date: Thu, 14 Oct 2021 19:16:49 +0100
Subject: [PATCH] Copy EntityList methods to BasicEntityList


diff --git a/src/main/java/io/papermc/paper/world/ChunkEntitySlices.java b/src/main/java/io/papermc/paper/world/ChunkEntitySlices.java
index c78cbec447032de9fe69748591bef6be300160ed..01701cdd0d984e35ff3453e6253ba9bcaa0db7b9 100644
--- a/src/main/java/io/papermc/paper/world/ChunkEntitySlices.java
+++ b/src/main/java/io/papermc/paper/world/ChunkEntitySlices.java
@@ -26,6 +26,8 @@ import java.util.List;
 import java.util.function.Predicate;
 import org.bukkit.event.entity.EntityRemoveEvent;
 
+import it.unimi.dsi.fastutil.ints.Int2IntOpenHashMap; // Sakura
+
 public final class ChunkEntitySlices {
 
     protected final int minSection;
@@ -304,6 +306,13 @@ public final class ChunkEntitySlices {
 
     protected static final class BasicEntityList<E extends Entity> {
 
+        // Sakura start - copy entitylist methods across
+        protected final Int2IntOpenHashMap entityToIndex = new Int2IntOpenHashMap(2, 0.8f);
+        {
+            this.entityToIndex.defaultReturnValue(Integer.MIN_VALUE);
+        }
+        // Sakura end
+
         protected static final Entity[] EMPTY = new Entity[0];
         protected static final int DEFAULT_CAPACITY = 4;
 
@@ -326,55 +335,52 @@ public final class ChunkEntitySlices {
             return this.size;
         }
 
-        private void resize() {
-            if (this.storage == EMPTY) {
-                this.storage = (E[])new Entity[DEFAULT_CAPACITY];
-            } else {
-                this.storage = Arrays.copyOf(this.storage, this.storage.length * 2);
-            }
-        }
-
+        // Sakura start - copy entitylist methods across
         public void add(final E entity) {
-            final int idx = this.size++;
-            if (idx >= this.storage.length) {
-                this.resize();
-                this.storage[idx] = entity;
-            } else {
-                this.storage[idx] = entity;
+            final int count = this.size;
+            final int currIndex = this.entityToIndex.putIfAbsent(entity.getId(), count);
+
+            if (currIndex != Integer.MIN_VALUE) {
+                return; // already in this list
             }
-        }
 
-        public int indexOf(final E entity) {
-            final E[] storage = this.storage;
+            E[] list = this.storage;
 
-            for (int i = 0, len = Math.min(this.storage.length, this.size); i < len; ++i) {
-                if (storage[i] == entity) {
-                    return i;
-                }
+            if (list.length == count) {
+                // resize required
+                list = this.storage = Arrays.copyOf(list, (int)Math.max(4L, count * 2L)); // overflow results in negative
             }
 
-            return -1;
+            list[count] = entity;
+            this.size = count + 1;
+        }
+
+        public int indexOf(final E entity) {
+            return this.entityToIndex.getOrDefault(entity.getId(), -1);
         }
 
         public boolean remove(final E entity) {
-            final int idx = this.indexOf(entity);
-            if (idx == -1) {
+            final int index = this.entityToIndex.remove(entity.getId());
+            if (index == Integer.MIN_VALUE) {
                 return false;
             }
 
-            final int size = --this.size;
-            final E[] storage = this.storage;
-            if (idx != size) {
-                System.arraycopy(storage, idx + 1, storage, idx, size - idx);
+            // move the entity at the end to this index
+            final int endIndex = --this.size;
+            final E end = this.storage[endIndex];
+            if (index != endIndex) {
+                // not empty after this call
+                this.entityToIndex.put(end.getId(), index); // update index
             }
-
-            storage[size] = null;
+            this.storage[index] = end;
+            this.storage[endIndex] = null;
 
             return true;
         }
 
         public boolean has(final E entity) {
-            return this.indexOf(entity) != -1;
+            return this.entityToIndex.containsKey(entity.getId());
+            // Sakura end
         }
     }