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

2 files changed, 59 insertions, 42 deletions

diff --git a/spark-neoforge/src/main/java/me/lucko/spark/neoforge/NeoForgeWorldInfoProvider.java b/spark-neoforge/src/main/java/me/lucko/spark/neoforge/NeoForgeWorldInfoProvider.java
index fe6a6c8..d172ecf 100644
--- a/spark-neoforge/src/main/java/me/lucko/spark/neoforge/NeoForgeWorldInfoProvider.java
+++ b/spark-neoforge/src/main/java/me/lucko/spark/neoforge/NeoForgeWorldInfoProvider.java
@@ -20,8 +20,8 @@
 
 package me.lucko.spark.neoforge;
 
-import it.unimi.dsi.fastutil.longs.LongIterator;
-import it.unimi.dsi.fastutil.longs.LongSet;
+import it.unimi.dsi.fastutil.longs.Long2ObjectOpenHashMap;
+import java.lang.reflect.Method;
 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 +34,16 @@ 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.LevelEntityGetter;
 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;
+import net.neoforged.fml.ModList;
 
 public abstract class NeoForgeWorldInfoProvider 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 NeoForgeWorldInfoProvider {
         private final MinecraftServer server;
 
@@ -74,10 +58,15 @@ public abstract class NeoForgeWorldInfoProvider implements WorldInfoProvider {
             int chunks = 0;
 
             for (ServerLevel level : this.server.getAllLevels()) {
-                PersistentEntitySectionManager<Entity> entityManager = level.entityManager;
-                EntityLookup<Entity> entityIndex = entityManager.visibleEntityStorage;
 
-                entities += entityIndex.count();
+                if (ModList.get().isLoaded("moonrise")) {
+                    entities += MoonriseMethods.getEntityCount(level.getEntities());
+                } else {
+                    PersistentEntitySectionManager<Entity> entityManager = level.entityManager;
+                    EntityLookup<Entity> entityIndex = entityManager.visibleEntityStorage;
+                    entities += entityIndex.count();
+                }
+
                 chunks += level.getChunkSource().getLoadedChunksCount();
             }
 
@@ -89,11 +78,15 @@ public abstract class NeoForgeWorldInfoProvider 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<>();
+
+                for (Entity entity : level.getEntities().getAll()) {
+                    ForgeChunkInfo info = levelInfos.computeIfAbsent(
+                        entity.chunkPosition().toLong(), ForgeChunkInfo::new);
+                    info.entityCounts.increment(entity.getType());
+                }
 
-                List<ForgeChunkInfo> list = getChunksFromCache(cache);
-                data.put(level.dimension().location().getPath(), list);
+                data.put(level.dimension().location().getPath(), List.copyOf(levelInfos.values()));
             }
 
             return data;
@@ -137,10 +130,15 @@ public abstract class NeoForgeWorldInfoProvider implements WorldInfoProvider {
                 return null;
             }
 
-            TransientEntitySectionManager<Entity> entityManager = level.entityStorage;
-            EntityLookup<Entity> entityIndex = entityManager.entityStorage;
+            int entities;
+            if (ModList.get().isLoaded("moonrise")) {
+                entities = MoonriseMethods.getEntityCount(level.getEntities());
+            } else {
+                TransientEntitySectionManager<Entity> entityManager = level.entityStorage;
+                EntityLookup<Entity> entityIndex = entityManager.entityStorage;
+                entities = entityIndex.count();
+            }
 
-            int entities = entityIndex.count();
             int chunks = level.getChunkSource().getLoadedChunksCount();
 
             return new CountsResult(-1, entities, -1, chunks);
@@ -155,11 +153,14 @@ public abstract class NeoForgeWorldInfoProvider 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;
         }
@@ -194,17 +195,10 @@ public abstract class NeoForgeWorldInfoProvider implements WorldInfoProvider {
     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
@@ -218,5 +212,27 @@ public abstract class NeoForgeWorldInfoProvider implements WorldInfoProvider {
         }
     }
 
+    private static final class MoonriseMethods {
+        private static Method getEntityCount;
+
+        private static Method getEntityCountMethod(LevelEntityGetter<Entity> getter) {
+            if (getEntityCount == null) {
+                try {
+                    getEntityCount = getter.getClass().getMethod("getEntityCount");
+                } catch (final ReflectiveOperationException e) {
+                    throw new RuntimeException("Cannot find Moonrise getEntityCount method", e);
+                }
+            }
+            return getEntityCount;
+        }
+
+        private static int getEntityCount(LevelEntityGetter<Entity> getter) {
+            try {
+                return (int) getEntityCountMethod(getter).invoke(getter);
+            } catch (final ReflectiveOperationException e) {
+                throw new RuntimeException("Failed to invoke Moonrise getEntityCount method", e);
+            }
+        }
+    }
 
 }
diff --git a/spark-neoforge/src/main/resources/META-INF/accesstransformer.cfg b/spark-neoforge/src/main/resources/META-INF/accesstransformer.cfg
index 43d14fc..80f3f6e 100644
--- a/spark-neoforge/src/main/resources/META-INF/accesstransformer.cfg
+++ b/spark-neoforge/src/main/resources/META-INF/accesstransformer.cfg
@@ -2,6 +2,7 @@ public net.minecraft.server.level.ServerLevel entityManager # entityManager
 public net.minecraft.world.level.entity.PersistentEntitySectionManager sectionStorage # sectionStorage
 public net.minecraft.world.level.entity.PersistentEntitySectionManager visibleEntityStorage # visibleEntityStorage
 public net.minecraft.client.multiplayer.ClientLevel entityStorage # entityStorage
+public net.minecraft.client.multiplayer.ClientLevel getEntities()Lnet/minecraft/world/level/entity/LevelEntityGetter; # getEntities
 public net.minecraft.world.level.entity.TransientEntitySectionManager sectionStorage # sectionStorage
 public net.minecraft.world.level.entity.TransientEntitySectionManager entityStorage # entityStorage
 public net.minecraft.client.Minecraft gameThread # gameThread