From a111271113e2abb2e6f4c98f118fdc830d1739f6 Mon Sep 17 00:00:00 2001
From: themode <themode@outlook.fr>
Date: Wed, 26 Jan 2022 18:37:21 +0100
Subject: [PATCH] Fix tests

---
 .../net/minestom/server/instance/EntityTracker.java  | 10 +++-------
 .../minestom/server/instance/EntityTrackerImpl.java  | 12 ++++++------
 .../java/net/minestom/server/utils/ViewEngine.java   |  6 ++++--
 3 files changed, 13 insertions(+), 15 deletions(-)

diff --git a/src/main/java/net/minestom/server/instance/EntityTracker.java b/src/main/java/net/minestom/server/instance/EntityTracker.java
index 51b45e522..48a93a01f 100644
--- a/src/main/java/net/minestom/server/instance/EntityTracker.java
+++ b/src/main/java/net/minestom/server/instance/EntityTracker.java
@@ -1,6 +1,5 @@
 package net.minestom.server.instance;
 
-import net.minestom.server.MinecraftServer;
 import net.minestom.server.coordinate.Point;
 import net.minestom.server.entity.Entity;
 import net.minestom.server.entity.ExperienceOrb;
@@ -81,13 +80,10 @@ public sealed interface EntityTracker permits EntityTrackerImpl {
         visibleEntities(point.chunkX(), point.chunkZ(), target, query);
     }
 
-    /**
-     * Gets a list containing references to all the entity from {@link MinecraftServer#getChunkViewDistance()} range.
-     */
-    <T extends Entity> @NotNull List<List<T>> references(int chunkX, int chunkZ, @NotNull Target<T> target);
+    <T extends Entity> @NotNull List<List<T>> references(int chunkX, int chunkZ, int range, @NotNull Target<T> target);
 
-    default <T extends Entity> @NotNull List<List<T>> references(@NotNull Point point, @NotNull Target<T> target) {
-        return references(point.chunkX(), point.chunkZ(), target);
+    default <T extends Entity> @NotNull List<List<T>> references(@NotNull Point point, int range, @NotNull Target<T> target) {
+        return references(point.chunkX(), point.chunkZ(), range, target);
     }
 
     /**
diff --git a/src/main/java/net/minestom/server/instance/EntityTrackerImpl.java b/src/main/java/net/minestom/server/instance/EntityTrackerImpl.java
index 916298499..2e293ec05 100644
--- a/src/main/java/net/minestom/server/instance/EntityTrackerImpl.java
+++ b/src/main/java/net/minestom/server/instance/EntityTrackerImpl.java
@@ -128,22 +128,22 @@ final class EntityTrackerImpl implements EntityTracker {
 
     @Override
     public <T extends Entity> void visibleEntities(int chunkX, int chunkZ, @NotNull Target<T> target, @NotNull Query<T> query) {
-        for (List<T> entities : references(chunkX, chunkZ, target)) {
+        for (List<T> entities : references(chunkX, chunkZ, MinecraftServer.getEntityViewDistance(), target)) {
             if (entities.isEmpty()) continue;
             for (T entity : entities) query.consume(entity);
         }
     }
 
     @Override
-    public @NotNull <T extends Entity> List<List<T>> references(int chunkX, int chunkZ, @NotNull Target<T> target) {
+    public @NotNull <T extends Entity> List<List<T>> references(int chunkX, int chunkZ, int range, @NotNull Target<T> target) {
         // Gets reference to all chunk entities lists within the range
         // This is used to avoid a map lookup per chunk
         final TargetEntry<T> entry = (TargetEntry<T>) entries[target.ordinal()];
-        return entry.chunkRangeEntities.computeIfAbsent(ChunkUtils.getChunkIndex(chunkX, chunkZ),
+        var rangeRef = entry.references.computeIfAbsent(range, integer -> Long2ObjectSyncMap.hashmap());
+        return rangeRef.computeIfAbsent(ChunkUtils.getChunkIndex(chunkX, chunkZ),
                 chunkIndex -> {
                     List<List<T>> entities = new ArrayList<>();
-                    ChunkUtils.forChunksInRange(ChunkUtils.getChunkCoordX(chunkIndex), ChunkUtils.getChunkCoordZ(chunkIndex),
-                            MinecraftServer.getChunkViewDistance(),
+                    ChunkUtils.forChunksInRange(ChunkUtils.getChunkCoordX(chunkIndex), ChunkUtils.getChunkCoordZ(chunkIndex), range,
                             (x, z) -> entities.add(entry.chunkEntities.computeIfAbsent(getChunkIndex(x, z), i -> (List<T>) LIST_SUPPLIER.get())));
                     return List.copyOf(entities);
                 });
@@ -177,7 +177,7 @@ final class EntityTrackerImpl implements EntityTracker {
         // Chunk index -> entities inside it
         private final Long2ObjectSyncMap<List<T>> chunkEntities = Long2ObjectSyncMap.hashmap();
         // Chunk index -> lists of visible entities (references to chunkEntities entries)
-        private final Long2ObjectSyncMap<List<List<T>>> chunkRangeEntities = Long2ObjectSyncMap.hashmap();
+        private final Int2ObjectSyncMap<Long2ObjectSyncMap<List<List<T>>>> references = Int2ObjectSyncMap.hashmap();
 
         TargetEntry(Target<T> target) {
             this.target = target;
diff --git a/src/main/java/net/minestom/server/utils/ViewEngine.java b/src/main/java/net/minestom/server/utils/ViewEngine.java
index 267a7bbc9..2d6a166d1 100644
--- a/src/main/java/net/minestom/server/utils/ViewEngine.java
+++ b/src/main/java/net/minestom/server/utils/ViewEngine.java
@@ -2,6 +2,7 @@ package net.minestom.server.utils;
 
 import it.unimi.dsi.fastutil.ints.IntOpenHashSet;
 import it.unimi.dsi.fastutil.ints.IntSet;
+import net.minestom.server.MinecraftServer;
 import net.minestom.server.coordinate.Point;
 import net.minestom.server.entity.Entity;
 import net.minestom.server.entity.Player;
@@ -55,8 +56,9 @@ public final class ViewEngine {
         synchronized (mutex) {
             this.tracker = tracker;
             if (tracker != null) {
-                this.viewableOption.references = tracker.references(point, EntityTracker.Target.PLAYERS);
-                this.viewerOption.references = tracker.references(point, EntityTracker.Target.ENTITIES);
+                final int range = entity != null ? MinecraftServer.getEntityViewDistance() : MinecraftServer.getChunkViewDistance();
+                this.viewableOption.references = tracker.references(point, range, EntityTracker.Target.PLAYERS);
+                this.viewerOption.references = tracker.references(point, range, EntityTracker.Target.ENTITIES);
             } else {
                 this.viewableOption.references = null;
                 this.viewerOption.references = null;