From 65f9460a1a27e930b3749525766fd44d57b65300 Mon Sep 17 00:00:00 2001
From: Luck <git@lucko.me>
Date: Sat, 26 Nov 2022 23:00:53 +0000
Subject: Include player/entity/chunk counts in window statistics

---
 .../spark/fabric/FabricWorldInfoProvider.java      | 42 +++++++++++++++++++---
 .../fabric/mixin/ClientEntityManagerAccessor.java  |  4 +++
 .../fabric/mixin/ServerEntityManagerAccessor.java  |  4 +++
 3 files changed, 46 insertions(+), 4 deletions(-)

(limited to 'spark-fabric/src/main/java/me/lucko')

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();
+
 }
-- 
cgit