SakuraMC/patches/server/0020-Limited-Get-Entities.patch

From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: Samsuik <kfian294ma4@gmail.com>
Date: Thu, 23 Sep 2021 15:19:31 +0100
Subject: [PATCH] Limited Get Entities


diff --git a/src/main/java/io/papermc/paper/chunk/system/entity/EntityLookup.java b/src/main/java/io/papermc/paper/chunk/system/entity/EntityLookup.java
index 15ee41452992714108efe53b708b5a4e1da7c1ff..b62a1ae53a891802db17046271b9981ab3af1df0 100644
--- a/src/main/java/io/papermc/paper/chunk/system/entity/EntityLookup.java
+++ b/src/main/java/io/papermc/paper/chunk/system/entity/EntityLookup.java
@@ -522,6 +522,128 @@ public final class EntityLookup implements LevelEntityGetter<Entity> {
         return slices;
     }
 
+    // Sakura start - limited get entities
+    public void getLimitedEntities(final Entity except, final AABB box, final List<Entity> into,
+                               final Predicate<? super Entity> predicate, final int limit, final int search) {
+        final int minChunkX = (Mth.floor(box.minX) - 2) >> 4;
+        final int minChunkZ = (Mth.floor(box.minZ) - 2) >> 4;
+        final int maxChunkX = (Mth.floor(box.maxX) + 2) >> 4;
+        final int maxChunkZ = (Mth.floor(box.maxZ) + 2) >> 4;
+
+        final int minRegionX = minChunkX >> REGION_SHIFT;
+        final int minRegionZ = minChunkZ >> REGION_SHIFT;
+        final int maxRegionX = maxChunkX >> REGION_SHIFT;
+        final int maxRegionZ = maxChunkZ >> REGION_SHIFT;
+
+        for (int currRegionZ = minRegionZ; currRegionZ <= maxRegionZ; ++currRegionZ) {
+            final int minZ = currRegionZ == minRegionZ ? minChunkZ & REGION_MASK : 0;
+            final int maxZ = currRegionZ == maxRegionZ ? maxChunkZ & REGION_MASK : REGION_MASK;
+
+            for (int currRegionX = minRegionX; currRegionX <= maxRegionX; ++currRegionX) {
+                final ChunkSlicesRegion region = this.getRegion(currRegionX, currRegionZ);
+
+                if (region == null) {
+                    continue;
+                }
+
+                final int minX = currRegionX == minRegionX ? minChunkX & REGION_MASK : 0;
+                final int maxX = currRegionX == maxRegionX ? maxChunkX & REGION_MASK : REGION_MASK;
+
+                for (int currZ = minZ; currZ <= maxZ; ++currZ) {
+                    for (int currX = minX; currX <= maxX; ++currX) {
+                        final ChunkEntitySlices chunk = region.get(currX | (currZ << REGION_SHIFT));
+                        if (chunk == null || !chunk.status.isOrAfter(FullChunkStatus.FULL)) {
+                            continue;
+                        }
+
+                        chunk.getLimitedEntities(except, box, into, predicate, limit, search);
+                    }
+                }
+            }
+        }
+    }
+
+    public <T extends Entity> void getLimitedEntities(final EntityType<?> type, final AABB box, final List<? super T> into,
+                                               final Predicate<? super T> predicate, final int limit, final int search) {
+        final int minChunkX = (Mth.floor(box.minX) - 2) >> 4;
+        final int minChunkZ = (Mth.floor(box.minZ) - 2) >> 4;
+        final int maxChunkX = (Mth.floor(box.maxX) + 2) >> 4;
+        final int maxChunkZ = (Mth.floor(box.maxZ) + 2) >> 4;
+
+        final int minRegionX = minChunkX >> REGION_SHIFT;
+        final int minRegionZ = minChunkZ >> REGION_SHIFT;
+        final int maxRegionX = maxChunkX >> REGION_SHIFT;
+        final int maxRegionZ = maxChunkZ >> REGION_SHIFT;
+
+        for (int currRegionZ = minRegionZ; currRegionZ <= maxRegionZ; ++currRegionZ) {
+            final int minZ = currRegionZ == minRegionZ ? minChunkZ & REGION_MASK : 0;
+            final int maxZ = currRegionZ == maxRegionZ ? maxChunkZ & REGION_MASK : REGION_MASK;
+
+            for (int currRegionX = minRegionX; currRegionX <= maxRegionX; ++currRegionX) {
+                final ChunkSlicesRegion region = this.getRegion(currRegionX, currRegionZ);
+
+                if (region == null) {
+                    continue;
+                }
+
+                final int minX = currRegionX == minRegionX ? minChunkX & REGION_MASK : 0;
+                final int maxX = currRegionX == maxRegionX ? maxChunkX & REGION_MASK : REGION_MASK;
+
+                for (int currZ = minZ; currZ <= maxZ; ++currZ) {
+                    for (int currX = minX; currX <= maxX; ++currX) {
+                        final ChunkEntitySlices chunk = region.get(currX | (currZ << REGION_SHIFT));
+                        if (chunk == null || !chunk.status.isOrAfter(FullChunkStatus.FULL)) {
+                            continue;
+                        }
+
+                        chunk.getLimitedEntities(type, box, (List)into, (Predicate)predicate, limit, search);
+                    }
+                }
+            }
+        }
+    }
+
+    public <T extends Entity> void getLimitedEntities(final Class<? extends T> clazz, final Entity except, final AABB box, final List<? super T> into,
+                                               final Predicate<? super T> predicate, final int limit, final int search) {
+        final int minChunkX = (Mth.floor(box.minX) - 2) >> 4;
+        final int minChunkZ = (Mth.floor(box.minZ) - 2) >> 4;
+        final int maxChunkX = (Mth.floor(box.maxX) + 2) >> 4;
+        final int maxChunkZ = (Mth.floor(box.maxZ) + 2) >> 4;
+
+        final int minRegionX = minChunkX >> REGION_SHIFT;
+        final int minRegionZ = minChunkZ >> REGION_SHIFT;
+        final int maxRegionX = maxChunkX >> REGION_SHIFT;
+        final int maxRegionZ = maxChunkZ >> REGION_SHIFT;
+
+        for (int currRegionZ = minRegionZ; currRegionZ <= maxRegionZ; ++currRegionZ) {
+            final int minZ = currRegionZ == minRegionZ ? minChunkZ & REGION_MASK : 0;
+            final int maxZ = currRegionZ == maxRegionZ ? maxChunkZ & REGION_MASK : REGION_MASK;
+
+            for (int currRegionX = minRegionX; currRegionX <= maxRegionX; ++currRegionX) {
+                final ChunkSlicesRegion region = this.getRegion(currRegionX, currRegionZ);
+
+                if (region == null) {
+                    continue;
+                }
+
+                final int minX = currRegionX == minRegionX ? minChunkX & REGION_MASK : 0;
+                final int maxX = currRegionX == maxRegionX ? maxChunkX & REGION_MASK : REGION_MASK;
+
+                for (int currZ = minZ; currZ <= maxZ; ++currZ) {
+                    for (int currX = minX; currX <= maxX; ++currX) {
+                        final ChunkEntitySlices chunk = region.get(currX | (currZ << REGION_SHIFT));
+                        if (chunk == null || !chunk.status.isOrAfter(FullChunkStatus.FULL)) {
+                            continue;
+                        }
+
+                        chunk.getLimitedEntities(clazz, except, box, into, predicate, limit, search);
+                    }
+                }
+            }
+        }
+    }
+    // Sakura end
+
     public void getEntitiesWithoutDragonParts(final Entity except, final AABB box, final List<Entity> into, final Predicate<? super Entity> predicate) {
         final int minChunkX = (Mth.floor(box.minX) - 2) >> 4;
         final int minChunkZ = (Mth.floor(box.minZ) - 2) >> 4;
diff --git a/src/main/java/io/papermc/paper/world/ChunkEntitySlices.java b/src/main/java/io/papermc/paper/world/ChunkEntitySlices.java
index 8fcaa00e461c7f4413bf655ddd8165a2b908f900..404b99def4562942e036089085a667979b1cac81 100644
--- a/src/main/java/io/papermc/paper/world/ChunkEntitySlices.java
+++ b/src/main/java/io/papermc/paper/world/ChunkEntitySlices.java
@@ -259,6 +259,30 @@ public final class ChunkEntitySlices {
     }
     // Sakura end
 
+    // Sakura start - limited get entities
+    public void getLimitedEntities(final Entity except, final AABB box, final List<Entity> into, final Predicate<? super Entity> predicate, final int limit, final int search) {
+        this.allEntities.getLimitedEntities(except, box, into, predicate, limit, search);
+    }
+
+    public <T extends Entity> void getLimitedEntities(final EntityType<?> type, final AABB box, final List<? super T> into,
+                                                      final Predicate<? super T> predicate, final int limit, final int search) {
+        this.allEntities.getLimitedEntities(type, box, (List)into, (Predicate)predicate, limit, search);
+    }
+
+    public <T extends Entity> void getLimitedEntities(final Class<? extends T> clazz, final Entity except, final AABB box, final List<? super T> into,
+                                               final Predicate<? super T> predicate, final int limit, final int search) {
+        EntityCollectionBySection collection = this.entitiesByClass.get(clazz);
+        if (collection != null) {
+            collection.getLimitedEntities(except, clazz, box, (List)into, (Predicate)predicate, limit, search);
+        } else {
+            synchronized (this) {
+                this.entitiesByClass.putIfAbsent(clazz, collection = this.initClass(clazz));
+            }
+            collection.getLimitedEntities(except, clazz, box, (List)into, (Predicate)predicate, limit, search);
+        }
+    }
+    // Sakura end
+
     public void getHardCollidingEntities(final Entity except, final AABB box, final List<Entity> into, final Predicate<? super Entity> predicate) {
         this.hardCollidingEntities.getEntities(except, box, into, predicate);
     }
@@ -449,6 +473,155 @@ public final class ChunkEntitySlices {
         }
         // Sakura end
 
+        // Sakura start - limited get entities
+        public void getLimitedEntities(final Entity except, final Class<?> clazz, final AABB box, final List<Entity> into,
+                                                    final Predicate<? super Entity> predicate, final int limit, int search) {
+            if (this.count == 0) {
+                return;
+            }
+
+            final int minSection = this.manager.minSection;
+            final int maxSection = this.manager.maxSection;
+
+            final int min = Mth.clamp(Mth.floor(box.minY - 2.0) >> 4, minSection, maxSection);
+            final int max = Mth.clamp(Mth.floor(box.maxY + 2.0) >> 4, minSection, maxSection);
+
+            // TODO use the bitset
+
+            final BasicEntityList<Entity>[] entitiesBySection = this.entitiesBySection;
+
+            for (int section = min; section <= max; ++section) {
+                final BasicEntityList<Entity> list = entitiesBySection[section - minSection];
+
+                if (list == null) {
+                    continue;
+                }
+
+                final Entity[] storage = list.storage;
+
+                int len = Math.min(storage.length, list.size());
+                int index = manager.world.random.nextInt(len); // starting position
+
+                for (int i = 0; i < len; ++i) {
+                    index = (index + 1) % len; // Update index
+                    final Entity entity = storage[index];
+
+                    if (into.size() >= limit || search-- <= 0) {
+                        return;
+                    }
+
+                    if (entity == null || entity == except || !clazz.isInstance(entity) || !entity.getBoundingBox().intersects(box)) {
+                        continue;
+                    }
+
+                    if (predicate != null && !predicate.test(entity)) {
+                        continue;
+                    }
+
+                    into.add(entity);
+                }
+            }
+        }
+
+        public <T extends Entity> void getLimitedEntities(final EntityType<?> type, final AABB box, final List<? super T> into,
+                                                   final Predicate<? super T> predicate, final int limit, int search) {
+            if (this.count == 0) {
+                return;
+            }
+
+            final int minSection = this.manager.minSection;
+            final int maxSection = this.manager.maxSection;
+
+            final int min = Mth.clamp(Mth.floor(box.minY - 2.0) >> 4, minSection, maxSection);
+            final int max = Mth.clamp(Mth.floor(box.maxY + 2.0) >> 4, minSection, maxSection);
+
+            // TODO use the bitset
+
+            final BasicEntityList<Entity>[] entitiesBySection = this.entitiesBySection;
+
+            for (int section = min; section <= max; ++section) {
+                final BasicEntityList<Entity> list = entitiesBySection[section - minSection];
+
+                if (list == null) {
+                    continue;
+                }
+
+                final Entity[] storage = list.storage;
+
+                int len = Math.min(storage.length, list.size());
+                int index = manager.world.random.nextInt(len); // starting position
+
+                for (int i = 0; i < len; ++i) {
+                    index = (index + 1) % len; // Update index
+                    final Entity entity = storage[index];
+
+                    if (into.size() >= limit || search-- <= 0) {
+                        return;
+                    }
+
+                    if (entity == null || (type != null && entity.getType() != type) || !entity.getBoundingBox().intersects(box)) {
+                        continue;
+                    }
+
+                    if (predicate != null && !predicate.test((T)entity)) {
+                        continue;
+                    }
+
+                    into.add((T)entity);
+                }
+            }
+        }
+
+        public void getLimitedEntities(final Entity except, final AABB box, final List<Entity> into,
+                                       final Predicate<? super Entity> predicate, final int limit, int search) {
+            if (this.count == 0) {
+                return;
+            }
+
+            final int minSection = this.manager.minSection;
+            final int maxSection = this.manager.maxSection;
+
+            final int min = Mth.clamp(Mth.floor(box.minY - 2.0) >> 4, minSection, maxSection);
+            final int max = Mth.clamp(Mth.floor(box.maxY + 2.0) >> 4, minSection, maxSection);
+
+            // TODO use the bitset
+
+            final BasicEntityList<Entity>[] entitiesBySection = this.entitiesBySection;
+
+            for (int section = min; section <= max; ++section) {
+                final BasicEntityList<Entity> list = entitiesBySection[section - minSection];
+
+                if (list == null) {
+                    continue;
+                }
+
+                final Entity[] storage = list.storage;
+
+                int len = Math.min(storage.length, list.size());
+                int index = manager.world.random.nextInt(len); // starting position
+
+                for (int i = 0; i < len; ++i) {
+                    index = (index + 1) % len; // Update index
+                    final Entity entity = storage[index];
+
+                    if (into.size() >= limit || search-- <= 0) {
+                        return;
+                    }
+
+                    if (entity == null || entity == except || !entity.getBoundingBox().intersects(box)) {
+                        continue;
+                    }
+
+                    if (predicate != null && !predicate.test(entity)) {
+                        continue;
+                    }
+
+                    into.add(entity);
+                }
+            }
+        }
+        // Sakura end
+
         public void getEntities(final Entity except, final AABB box, final List<Entity> into, final Predicate<? super Entity> predicate) {
             if (this.count == 0) {
                 return;
diff --git a/src/main/java/net/minecraft/world/level/Level.java b/src/main/java/net/minecraft/world/level/Level.java
index 5b6b9de92749a5c1c5f4f52b78d0332212ef351c..2ab6bfd7887cd7a8b4bb6a9b09938d9c30f1cb41 100644
--- a/src/main/java/net/minecraft/world/level/Level.java
+++ b/src/main/java/net/minecraft/world/level/Level.java
@@ -232,6 +232,39 @@ public abstract class Level implements LevelAccessor, AutoCloseable {
     public final me.samsuik.sakura.entity.merge.MergeHistory mergeHistory = new me.samsuik.sakura.entity.merge.MergeHistory(); // Sakura - cannon entity merging
     public final me.samsuik.sakura.explosion.DensityCache densityCache = new me.samsuik.sakura.explosion.DensityCache(); // Sakura - specialised density cache for swinging explosions
 
+    // Sakura start - add entity retrival methods with search limits
+    public void getLimitedEntities(Entity except, AABB box, Predicate<? super Entity> predicate, List<Entity> into, int limit, int search) {
+        ((ServerLevel)this).getEntityLookup().getLimitedEntities(except, box, into, predicate, limit, search);
+    }
+
+    public <T extends Entity> List<T> getLimitedEntities(net.minecraft.world.entity.EntityType<?> filter, AABB box, Predicate<? super T> predicate, int limit, int search) {
+        List<T> ret = new java.util.ArrayList<>();
+        ((ServerLevel)this).getEntityLookup().getLimitedEntities(filter, box, ret, predicate, limit, search);
+        return ret;
+    }
+
+    public <T> void getLimitedEntitiesByClass(Class<? extends T> clazz, Entity except, final AABB box, List<? super T> into,
+                                              Predicate<? super T> predicate, int limit, int search) {
+        ((ServerLevel)this).getEntityLookup().getLimitedEntities((Class)clazz, except, box, (List)into, (Predicate)predicate, limit, search);
+    }
+
+    public <T extends Entity> List<T> getLimitedEntitiesOfClass(Class<T> entityClass, AABB box, Predicate<? super T> predicate, int limit, int search) {
+        List<T> ret = new java.util.ArrayList<>();
+        ((ServerLevel)this).getEntityLookup().getLimitedEntities(entityClass, null, box, ret, predicate, limit, search);
+        return ret;
+    }
+
+    public List<Entity> getLimitedEntities(@Nullable Entity except, AABB box, Predicate<? super Entity> predicate, int limit, int search) {
+        List<Entity> list = Lists.newArrayList();
+        getLimitedEntities(except, box, predicate, list, limit, search);
+        return list;
+    }
+
+    public List<Entity> getLimitedEntities(@Nullable Entity except, AABB box, int limit, int search) {
+        return this.getLimitedEntities(except, box, net.minecraft.world.entity.EntitySelector.NO_SPECTATORS, limit, search);
+    }
+    // Sakura end - add entity retrival methods with search limits
+
     protected Level(WritableLevelData worlddatamutable, ResourceKey<Level> resourcekey, RegistryAccess iregistrycustom, Holder<DimensionType> holder, Supplier<ProfilerFiller> supplier, boolean flag, boolean flag1, long i, int j, org.bukkit.generator.ChunkGenerator gen, org.bukkit.generator.BiomeProvider biomeProvider, org.bukkit.World.Environment env, java.util.function.Function<org.spigotmc.SpigotWorldConfig, io.papermc.paper.configuration.WorldConfiguration> paperWorldConfigCreator, Supplier<me.samsuik.sakura.configuration.WorldConfiguration> sakuraWorldConfigCreator, java.util.concurrent.Executor executor) { // Sakura // Paper - create paper world config; Async-Anti-Xray: Pass executor
         this.spigotConfig = new org.spigotmc.SpigotWorldConfig(((net.minecraft.world.level.storage.PrimaryLevelData) worlddatamutable).getLevelName()); // Spigot
         this.paperConfig = paperWorldConfigCreator.apply(this.spigotConfig); // Paper - create paper world config