Fix entity counts & chunk info with Moonrise on Fabric and NeoForge (#445)

2 files changed, 19 insertions, 37 deletions

diff --git a/spark-forge/src/main/java/me/lucko/spark/forge/ForgeWorldInfoProvider.java b/spark-forge/src/main/java/me/lucko/spark/forge/ForgeWorldInfoProvider.java
index 286605a..0d1724e 100644
--- a/spark-forge/src/main/java/me/lucko/spark/forge/ForgeWorldInfoProvider.java
+++ b/spark-forge/src/main/java/me/lucko/spark/forge/ForgeWorldInfoProvider.java
@@ -20,8 +20,7 @@
 
 package me.lucko.spark.forge;
 
-import it.unimi.dsi.fastutil.longs.LongIterator;
-import it.unimi.dsi.fastutil.longs.LongSet;
+import it.unimi.dsi.fastutil.longs.Long2ObjectOpenHashMap;
 import me.lucko.spark.common.platform.world.AbstractChunkInfo;
 import me.lucko.spark.common.platform.world.CountMap;
 import me.lucko.spark.common.platform.world.WorldInfoProvider;
@@ -34,32 +33,14 @@ import net.minecraft.world.entity.EntityType;
 import net.minecraft.world.level.ChunkPos;
 import net.minecraft.world.level.GameRules;
 import net.minecraft.world.level.entity.EntityLookup;
-import net.minecraft.world.level.entity.EntitySection;
-import net.minecraft.world.level.entity.EntitySectionStorage;
 import net.minecraft.world.level.entity.PersistentEntitySectionManager;
 import net.minecraft.world.level.entity.TransientEntitySectionManager;
 
-import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
-import java.util.stream.Stream;
 
 public abstract class ForgeWorldInfoProvider implements WorldInfoProvider {
 
-    protected List<ForgeChunkInfo> getChunksFromCache(EntitySectionStorage<Entity> cache) {
-        LongSet loadedChunks = cache.getAllChunksWithExistingSections();
-        List<ForgeChunkInfo> list = new ArrayList<>(loadedChunks.size());
-
-        for (LongIterator iterator = loadedChunks.iterator(); iterator.hasNext(); ) {
-            long chunkPos = iterator.nextLong();
-            Stream<EntitySection<Entity>> sections = cache.getExistingSectionsInChunk(chunkPos);
-
-            list.add(new ForgeChunkInfo(chunkPos, sections));
-        }
-
-        return list;
-    }
-
     public static final class Server extends ForgeWorldInfoProvider {
         private final MinecraftServer server;
 
@@ -89,11 +70,15 @@ public abstract class ForgeWorldInfoProvider implements WorldInfoProvider {
             ChunksResult<ForgeChunkInfo> data = new ChunksResult<>();
 
             for (ServerLevel level : this.server.getAllLevels()) {
-                PersistentEntitySectionManager<Entity> entityManager = level.entityManager;
-                EntitySectionStorage<Entity> cache = entityManager.sectionStorage;
+                Long2ObjectOpenHashMap<ForgeChunkInfo> levelInfos = new Long2ObjectOpenHashMap<>();
 
-                List<ForgeChunkInfo> list = getChunksFromCache(cache);
-                data.put(level.dimension().location().getPath(), list);
+                for (Entity entity : level.getEntities().getAll()) {
+                    ForgeChunkInfo info = levelInfos.computeIfAbsent(
+                        entity.chunkPosition().toLong(), ForgeChunkInfo::new);
+                    info.entityCounts.increment(entity.getType());
+                }
+
+                data.put(level.dimension().location().getPath(), List.copyOf(levelInfos.values()));
             }
 
             return data;
@@ -155,11 +140,14 @@ public abstract class ForgeWorldInfoProvider implements WorldInfoProvider {
 
             ChunksResult<ForgeChunkInfo> data = new ChunksResult<>();
 
-            TransientEntitySectionManager<Entity> entityManager = level.entityStorage;
-            EntitySectionStorage<Entity> cache = entityManager.sectionStorage;
+            Long2ObjectOpenHashMap<ForgeChunkInfo> levelInfos = new Long2ObjectOpenHashMap<>();
 
-            List<ForgeChunkInfo> list = getChunksFromCache(cache);
-            data.put(level.dimension().location().getPath(), list);
+            for (Entity entity : level.getEntities().getAll()) {
+                ForgeChunkInfo info = levelInfos.computeIfAbsent(entity.chunkPosition().toLong(), ForgeChunkInfo::new);
+                info.entityCounts.increment(entity.getType());
+            }
+
+            data.put(level.dimension().location().getPath(), List.copyOf(levelInfos.values()));
 
             return data;
         }
@@ -191,20 +179,13 @@ public abstract class ForgeWorldInfoProvider implements WorldInfoProvider {
         }
     }
 
-    static final class ForgeChunkInfo extends AbstractChunkInfo<EntityType<?>> {
+    public static final class ForgeChunkInfo extends AbstractChunkInfo<EntityType<?>> {
         private final CountMap<EntityType<?>> entityCounts;
 
-        ForgeChunkInfo(long chunkPos, Stream<EntitySection<Entity>> entities) {
+        ForgeChunkInfo(long chunkPos) {
             super(ChunkPos.getX(chunkPos), ChunkPos.getZ(chunkPos));
 
             this.entityCounts = new CountMap.Simple<>(new HashMap<>());
-            entities.forEach(section -> {
-                if (section.getStatus().isAccessible()) {
-                    section.getEntities().forEach(entity ->
-                            this.entityCounts.increment(entity.getType())
-                    );
-                }
-            });
         }
 
         @Override
diff --git a/spark-forge/src/main/resources/META-INF/accesstransformer.cfg b/spark-forge/src/main/resources/META-INF/accesstransformer.cfg
index 2699a0e..3772274 100644
--- a/spark-forge/src/main/resources/META-INF/accesstransformer.cfg
+++ b/spark-forge/src/main/resources/META-INF/accesstransformer.cfg
@@ -5,3 +5,4 @@ public net.minecraft.client.multiplayer.ClientLevel f_171631_ # entityStorage
 public net.minecraft.world.level.entity.TransientEntitySectionManager f_157638_ # sectionStorage
 public net.minecraft.world.level.entity.TransientEntitySectionManager f_157637_ # entityStorage
 public net.minecraft.client.Minecraft f_91018_ # gameThread
+public net.minecraft.client.multiplayer.ClientLevel m_142646_()Lnet/minecraft/world/level/entity/LevelEntityGetter; # getEntities