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

2 files changed, 96 insertions, 42 deletions

diff --git a/spark-fabric/src/main/java/me/lucko/spark/fabric/FabricWorldInfoProvider.java b/spark-fabric/src/main/java/me/lucko/spark/fabric/FabricWorldInfoProvider.java
index d7dc7a7..eae9c86 100644
--- a/spark-fabric/src/main/java/me/lucko/spark/fabric/FabricWorldInfoProvider.java
+++ b/spark-fabric/src/main/java/me/lucko/spark/fabric/FabricWorldInfoProvider.java
@@ -20,8 +20,8 @@
 
 package me.lucko.spark.fabric;
 
-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;
@@ -29,6 +29,8 @@ import me.lucko.spark.fabric.mixin.ClientEntityManagerAccessor;
 import me.lucko.spark.fabric.mixin.ClientWorldAccessor;
 import me.lucko.spark.fabric.mixin.ServerEntityManagerAccessor;
 import me.lucko.spark.fabric.mixin.ServerWorldAccessor;
+import me.lucko.spark.fabric.mixin.WorldAccessor;
+import net.fabricmc.loader.api.FabricLoader;
 import net.minecraft.client.MinecraftClient;
 import net.minecraft.client.world.ClientWorld;
 import net.minecraft.entity.Entity;
@@ -40,30 +42,13 @@ import net.minecraft.util.math.ChunkPos;
 import net.minecraft.world.GameRules;
 import net.minecraft.world.entity.ClientEntityManager;
 import net.minecraft.world.entity.EntityIndex;
-import net.minecraft.world.entity.EntityTrackingSection;
-import net.minecraft.world.entity.SectionedEntityCache;
+import net.minecraft.world.entity.EntityLookup;
 
-import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
-import java.util.stream.Stream;
 
 public abstract class FabricWorldInfoProvider implements WorldInfoProvider {
 
-    protected List<FabricChunkInfo> getChunksFromCache(SectionedEntityCache<Entity> cache) {
-        LongSet loadedChunks = cache.getChunkPositions();
-        List<FabricChunkInfo> list = new ArrayList<>(loadedChunks.size());
-
-        for (LongIterator iterator = loadedChunks.iterator(); iterator.hasNext(); ) {
-            long chunkPos = iterator.nextLong();
-            Stream<EntityTrackingSection<Entity>> sections = cache.getTrackingSections(chunkPos);
-
-            list.add(new FabricChunkInfo(chunkPos, sections));
-        }
-
-        return list;
-    }
-
     public static final class Server extends FabricWorldInfoProvider {
         private final MinecraftServer server;
 
@@ -78,10 +63,15 @@ public abstract class FabricWorldInfoProvider implements WorldInfoProvider {
             int chunks = 0;
 
             for (ServerWorld world : this.server.getWorlds()) {
-                ServerEntityManager<Entity> entityManager = ((ServerWorldAccessor) world).getEntityManager();
-                EntityIndex<?> entityIndex = ((ServerEntityManagerAccessor) entityManager).getIndex();
 
-                entities += entityIndex.size();
+                if (FabricLoader.getInstance().isModLoaded("moonrise")) {
+                    entities += MoonriseMethods.getEntityCount(((WorldAccessor) world).spark$getEntityLookup());
+                } else {
+                    ServerEntityManager<Entity> entityManager = ((ServerWorldAccessor) world).getEntityManager();
+                    EntityIndex<?> entityIndex = ((ServerEntityManagerAccessor) entityManager).getIndex();
+                    entities += entityIndex.size();
+                }
+
                 chunks += world.getChunkManager().getLoadedChunkCount();
             }
 
@@ -93,11 +83,15 @@ public abstract class FabricWorldInfoProvider implements WorldInfoProvider {
             ChunksResult<FabricChunkInfo> data = new ChunksResult<>();
 
             for (ServerWorld world : this.server.getWorlds()) {
-                ServerEntityManager<Entity> entityManager = ((ServerWorldAccessor) world).getEntityManager();
-                SectionedEntityCache<Entity> cache = ((ServerEntityManagerAccessor) entityManager).getCache();
+                Long2ObjectOpenHashMap<FabricChunkInfo> worldInfos = new Long2ObjectOpenHashMap<>();
+
+                for (Entity entity : ((WorldAccessor) world).spark$getEntityLookup().iterate()) {
+                    FabricChunkInfo info = worldInfos.computeIfAbsent(
+                        entity.getChunkPos().toLong(), FabricChunkInfo::new);
+                    info.entityCounts.increment(entity.getType());
+                }
 
-                List<FabricChunkInfo> list = getChunksFromCache(cache);
-                data.put(world.getRegistryKey().getValue().getPath(), list);
+                data.put(world.getRegistryKey().getValue().getPath(), List.copyOf(worldInfos.values()));
             }
 
             return data;
@@ -140,10 +134,16 @@ public abstract class FabricWorldInfoProvider implements WorldInfoProvider {
                 return null;
             }
 
-            ClientEntityManager<Entity> entityManager = ((ClientWorldAccessor) world).getEntityManager();
-            EntityIndex<?> entityIndex = ((ClientEntityManagerAccessor) entityManager).getIndex();
+            int entities;
+
+            if (FabricLoader.getInstance().isModLoaded("moonrise")) {
+                entities = MoonriseMethods.getEntityCount(((WorldAccessor) world).spark$getEntityLookup());
+            } else {
+                ClientEntityManager<Entity> entityManager = ((ClientWorldAccessor) world).getEntityManager();
+                EntityIndex<?> entityIndex = ((ClientEntityManagerAccessor) entityManager).getIndex();
+                entities = entityIndex.size();
+            }
 
-            int entities = entityIndex.size();
             int chunks = world.getChunkManager().getLoadedChunkCount();
 
             return new CountsResult(-1, entities, -1, chunks);
@@ -158,11 +158,14 @@ public abstract class FabricWorldInfoProvider implements WorldInfoProvider {
 
             ChunksResult<FabricChunkInfo> data = new ChunksResult<>();
 
-            ClientEntityManager<Entity> entityManager = ((ClientWorldAccessor) world).getEntityManager();
-            SectionedEntityCache<Entity> cache = ((ClientEntityManagerAccessor) entityManager).getCache();
+            Long2ObjectOpenHashMap<FabricChunkInfo> worldInfos = new Long2ObjectOpenHashMap<>();
 
-            List<FabricChunkInfo> list = getChunksFromCache(cache);
-            data.put(world.getRegistryKey().getValue().getPath(), list);
+            for (Entity entity : ((WorldAccessor) world).spark$getEntityLookup().iterate()) {
+                FabricChunkInfo info = worldInfos.computeIfAbsent(entity.getChunkPos().toLong(), FabricChunkInfo::new);
+                info.entityCounts.increment(entity.getType());
+            }
+
+            data.put(world.getRegistryKey().getValue().getPath(), List.copyOf(worldInfos.values()));
 
             return data;
         }
@@ -197,17 +200,10 @@ public abstract class FabricWorldInfoProvider implements WorldInfoProvider {
     static final class FabricChunkInfo extends AbstractChunkInfo<EntityType<?>> {
         private final CountMap<EntityType<?>> entityCounts;
 
-        FabricChunkInfo(long chunkPos, Stream<EntityTrackingSection<Entity>> entities) {
+        FabricChunkInfo(long chunkPos) {
             super(ChunkPos.getPackedX(chunkPos), ChunkPos.getPackedZ(chunkPos));
 
             this.entityCounts = new CountMap.Simple<>(new HashMap<>());
-            entities.forEach(section -> {
-                if (section.getStatus().shouldTrack()) {
-                    section.stream().forEach(entity ->
-                            this.entityCounts.increment(entity.getType())
-                    );
-                }
-            });
         }
 
         @Override
@@ -221,5 +217,28 @@ public abstract class FabricWorldInfoProvider implements WorldInfoProvider {
         }
     }
 
+    private static final class MoonriseMethods {
+        private static Method getEntityCount;
+
+        private static Method getEntityCountMethod(EntityLookup<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(EntityLookup<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-fabric/src/main/java/me/lucko/spark/fabric/mixin/WorldAccessor.java b/spark-fabric/src/main/java/me/lucko/spark/fabric/mixin/WorldAccessor.java
new file mode 100644
index 0000000..e50dc20
--- /dev/null
+++ b/spark-fabric/src/main/java/me/lucko/spark/fabric/mixin/WorldAccessor.java
@@ -0,0 +1,35 @@
+/*
+ * This file is part of spark.
+ *
+ *  Copyright (c) lucko (Luck) <luck@lucko.me>
+ *  Copyright (c) contributors
+ *
+ *  This program is free software: you can redistribute it and/or modify
+ *  it under the terms of the GNU General Public License as published by
+ *  the Free Software Foundation, either version 3 of the License, or
+ *  (at your option) any later version.
+ *
+ *  This program is distributed in the hope that it will be useful,
+ *  but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *  GNU General Public License for more details.
+ *
+ *  You should have received a copy of the GNU General Public License
+ *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package me.lucko.spark.fabric.mixin;
+
+import net.minecraft.entity.Entity;
+import net.minecraft.world.World;
+import net.minecraft.world.entity.EntityLookup;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.gen.Invoker;
+
+@Mixin(World.class)
+public interface WorldAccessor {
+
+    @Invoker(value = "getEntityLookup")
+    EntityLookup<Entity> spark$getEntityLookup();
+
+}