Include player/entity/chunk counts in window statistics

3 files changed, 46 insertions, 4 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 f2f7b96..156db89 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
@@ -40,6 +40,7 @@ import net.minecraft.server.MinecraftServer;
 import net.minecraft.server.world.ServerEntityManager;
 import net.minecraft.server.world.ServerWorld;
 import net.minecraft.util.math.ChunkPos;
+import net.minecraft.world.entity.EntityIndex;
 import net.minecraft.world.entity.EntityTrackingSection;
 import net.minecraft.world.entity.SectionedEntityCache;
 
@@ -72,8 +73,25 @@ public abstract class FabricWorldInfoProvider implements WorldInfoProvider {
         }
 
         @Override
-        public Result<FabricChunkInfo> poll() {
-            Result<FabricChunkInfo> data = new Result<>();
+        public CountsResult pollCounts() {
+            int players = this.server.getCurrentPlayerCount();
+            int entities = 0;
+            int chunks = 0;
+
+            for (ServerWorld world : this.server.getWorlds()) {
+                ServerEntityManager<Entity> entityManager = ((ServerWorldAccessor) world).getEntityManager();
+                EntityIndex<?> entityIndex = ((ServerEntityManagerAccessor) entityManager).getIndex();
+
+                entities += entityIndex.size();
+                chunks += world.getChunkManager().getLoadedChunkCount();
+            }
+
+            return new CountsResult(players, entities, -1, chunks);
+        }
+
+        @Override
+        public ChunksResult<FabricChunkInfo> pollChunks() {
+            ChunksResult<FabricChunkInfo> data = new ChunksResult<>();
 
             for (ServerWorld world : this.server.getWorlds()) {
                 ServerEntityManager<Entity> entityManager = ((ServerWorldAccessor) world).getEntityManager();
@@ -95,8 +113,24 @@ public abstract class FabricWorldInfoProvider implements WorldInfoProvider {
         }
 
         @Override
-        public Result<FabricChunkInfo> poll() {
-            Result<FabricChunkInfo> data = new Result<>();
+        public CountsResult pollCounts() {
+            ClientWorld world = this.client.world;
+            if (world == null) {
+                return null;
+            }
+
+            ClientEntityManager<Entity> entityManager = ((ClientWorldAccessor) world).getEntityManager();
+            EntityIndex<?> entityIndex = ((ClientEntityManagerAccessor) entityManager).getIndex();
+
+            int entities = entityIndex.size();
+            int chunks = world.getChunkManager().getLoadedChunkCount();
+
+            return new CountsResult(-1, entities, -1, chunks);
+        }
+
+        @Override
+        public ChunksResult<FabricChunkInfo> pollChunks() {
+            ChunksResult<FabricChunkInfo> data = new ChunksResult<>();
 
             ClientWorld world = this.client.world;
             if (world == null) {
diff --git a/spark-fabric/src/main/java/me/lucko/spark/fabric/mixin/ClientEntityManagerAccessor.java b/spark-fabric/src/main/java/me/lucko/spark/fabric/mixin/ClientEntityManagerAccessor.java
index 88c9521..994c9a3 100644
--- a/spark-fabric/src/main/java/me/lucko/spark/fabric/mixin/ClientEntityManagerAccessor.java
+++ b/spark-fabric/src/main/java/me/lucko/spark/fabric/mixin/ClientEntityManagerAccessor.java
@@ -22,6 +22,7 @@ package me.lucko.spark.fabric.mixin;
 
 import net.minecraft.client.world.ClientEntityManager;
 import net.minecraft.entity.Entity;
+import net.minecraft.world.entity.EntityIndex;
 import net.minecraft.world.entity.SectionedEntityCache;
 
 import org.spongepowered.asm.mixin.Mixin;
@@ -33,4 +34,7 @@ public interface ClientEntityManagerAccessor {
     @Accessor
     SectionedEntityCache<Entity> getCache();
 
+    @Accessor
+    EntityIndex<?> getIndex();
+
 }
diff --git a/spark-fabric/src/main/java/me/lucko/spark/fabric/mixin/ServerEntityManagerAccessor.java b/spark-fabric/src/main/java/me/lucko/spark/fabric/mixin/ServerEntityManagerAccessor.java
index 160a12b..2c67502 100644
--- a/spark-fabric/src/main/java/me/lucko/spark/fabric/mixin/ServerEntityManagerAccessor.java
+++ b/spark-fabric/src/main/java/me/lucko/spark/fabric/mixin/ServerEntityManagerAccessor.java
@@ -22,6 +22,7 @@ package me.lucko.spark.fabric.mixin;
 
 import net.minecraft.entity.Entity;
 import net.minecraft.server.world.ServerEntityManager;
+import net.minecraft.world.entity.EntityIndex;
 import net.minecraft.world.entity.SectionedEntityCache;
 
 import org.spongepowered.asm.mixin.Mixin;
@@ -33,4 +34,7 @@ public interface ServerEntityManagerAccessor {
     @Accessor
     SectionedEntityCache<Entity> getCache();
 
+    @Accessor
+    EntityIndex<?> getIndex();
+
 }