Leaf/leaf-server/minecraft-patches/features/0090-Smart-sort-entities-in-NearestLivingEntitySensor.patch

From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: HaHaWTH <102713261+HaHaWTH@users.noreply.github.com>
Date: Tue, 9 Nov 2077 00:00:00 +0800
Subject: [PATCH] Smart sort entities in NearestLivingEntitySensor

Co-authored-by: Taiyou06 <kaandindar21@gmail.com>

FastBitRadix Sort will be used. (see https://ieeexplore.ieee.org/document/7822019 for more)
When entity count reached the threshold, Bucket Sort will be used.

In non-strict test, this can give ~60-110% improvement (524ms on Paper, 204ms on Leaf),
under 625 villagers situation.

diff --git a/net/minecraft/world/entity/ai/sensing/NearestLivingEntitySensor.java b/net/minecraft/world/entity/ai/sensing/NearestLivingEntitySensor.java
index b0c5e41fefc7c9adf1a61bd5b52861736657d37e..9ef531551bb135d4b6a1ca9c3320e35fa62d7c41 100644
--- a/net/minecraft/world/entity/ai/sensing/NearestLivingEntitySensor.java
+++ b/net/minecraft/world/entity/ai/sensing/NearestLivingEntitySensor.java
@@ -13,18 +13,92 @@ import net.minecraft.world.entity.ai.memory.NearestVisibleLivingEntities;
 import net.minecraft.world.phys.AABB;
 
 public class NearestLivingEntitySensor<T extends LivingEntity> extends Sensor<T> {
+    // Leaf start - Optimized entity sorting
+    private static final int SMALL_ARRAY_THRESHOLD = 2;
+
     @Override
     protected void doTick(ServerLevel level, T entity) {
-        double attributeValue = entity.getAttributeValue(Attributes.FOLLOW_RANGE);
-        AABB aabb = entity.getBoundingBox().inflate(attributeValue, attributeValue, attributeValue);
-        List<LivingEntity> entitiesOfClass = level.getEntitiesOfClass(
-            LivingEntity.class, aabb, matchableEntity -> matchableEntity != entity && matchableEntity.isAlive()
+        double range = entity.getAttributeValue(Attributes.FOLLOW_RANGE);
+        AABB aabb = entity.getBoundingBox().inflate(range, range, range);
+
+        List<LivingEntity> entities = level.getEntitiesOfClass(
+            LivingEntity.class, aabb, e -> e != entity && e.isAlive()
         );
-        entitiesOfClass.sort(Comparator.comparingDouble(entity::distanceToSqr));
+
+        LivingEntity[] sorted = smartSort(entities.toArray(new LivingEntity[0]), entity);
+        List<LivingEntity> sortedList = List.of(sorted);
+
         Brain<?> brain = entity.getBrain();
-        brain.setMemory(MemoryModuleType.NEAREST_LIVING_ENTITIES, entitiesOfClass);
-        brain.setMemory(MemoryModuleType.NEAREST_VISIBLE_LIVING_ENTITIES, new NearestVisibleLivingEntities(level, entity, entitiesOfClass));
+        brain.setMemory(MemoryModuleType.NEAREST_LIVING_ENTITIES, sortedList);
+        brain.setMemory(MemoryModuleType.NEAREST_VISIBLE_LIVING_ENTITIES,
+            new NearestVisibleLivingEntities(level, entity, sortedList));
+    }
+
+    private LivingEntity[] smartSort(LivingEntity[] entities, T reference) {
+        if (entities.length <= 1) return entities;
+
+        long[] bits = new long[entities.length];
+        for (int i = 0; i < entities.length; i++) {
+            bits[i] = Double.doubleToRawLongBits(reference.distanceToSqr(entities[i]));
+        }
+
+        fastRadixSort(entities, bits, 0, entities.length - 1, 62);
+        return entities;
+    }
+
+    private void fastRadixSort(LivingEntity[] ents, long[] bits, int low, int high, int bit) {
+        if (bit < 0 || low >= high) return;
+
+        if (high - low <= SMALL_ARRAY_THRESHOLD) {
+            insertionSort(ents, bits, low, high);
+            return;
+        }
+
+        int i = low, j = high;
+        final long mask = 1L << bit;
+
+        while (i <= j) {
+            while (i <= j && (bits[i] & mask) == 0) i++;
+            while (i <= j && (bits[j] & mask) != 0) j--;
+
+            if (i < j) {
+                swap(ents, bits, i++, j--);
+            }
+        }
+
+        if (low < j) fastRadixSort(ents, bits, low, j, bit - 1);
+        if (i < high) fastRadixSort(ents, bits, i, high, bit - 1);
+    }
+
+    private void insertionSort(LivingEntity[] ents, long[] bits, int low, int high) {
+        for (int i = low + 1; i <= high; i++) {
+            int j = i;
+            LivingEntity e = ents[j];
+            long b = bits[j];
+
+            while (j > low && bits[j - 1] > b) {
+                ents[j] = ents[j - 1];
+                bits[j] = bits[j - 1];
+                j--;
+            }
+
+            ents[j] = e;
+            bits[j] = b;
+        }
+    }
+
+    private void swap(LivingEntity[] ents, long[] bits, int a, int b) {
+        // Swap entities
+        LivingEntity te = ents[a];
+        ents[a] = ents[b];
+        ents[b] = te;
+
+        // Swap bits
+        long tb = bits[a];
+        bits[a] = bits[b];
+        bits[b] = tb;
     }
+    // Leaf end - Optimized entity sorting
 
     @Override
     public Set<MemoryModuleType<?>> requires() {