Add providers for world (entity/chunk) statistics

3 files changed, 165 insertions, 0 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
new file mode 100644
index 0000000..b17dab5
--- /dev/null
+++ b/spark-forge/src/main/java/me/lucko/spark/forge/ForgeWorldInfoProvider.java
@@ -0,0 +1,141 @@
+/*
+ * 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.forge;
+
+import it.unimi.dsi.fastutil.longs.LongIterator;
+import it.unimi.dsi.fastutil.longs.LongSet;
+
+import me.lucko.spark.common.platform.world.AbstractChunkInfo;
+import me.lucko.spark.common.platform.world.CountMap;
+import me.lucko.spark.common.platform.world.WorldInfoProvider;
+
+import net.minecraft.client.Minecraft;
+import net.minecraft.client.multiplayer.ClientLevel;
+import net.minecraft.server.MinecraftServer;
+import net.minecraft.server.level.ServerLevel;
+import net.minecraft.world.entity.Entity;
+import net.minecraft.world.entity.EntityType;
+import net.minecraft.world.level.ChunkPos;
+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 net.minecraft.world.level.storage.ServerLevelData;
+
+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;
+
+        public Server(MinecraftServer server) {
+            this.server = server;
+        }
+
+        @Override
+        public Result<ForgeChunkInfo> poll() {
+            Result<ForgeChunkInfo> data = new Result<>();
+
+            for (ServerLevel level : this.server.getAllLevels()) {
+                PersistentEntitySectionManager<Entity> entityManager = level.entityManager;
+                EntitySectionStorage<Entity> cache = entityManager.sectionStorage;
+
+                List<ForgeChunkInfo> list = getChunksFromCache(cache);
+                data.put(((ServerLevelData) level.getLevelData()).getLevelName(), list);
+            }
+
+            return data;
+        }
+    }
+
+    public static final class Client extends ForgeWorldInfoProvider {
+        private final Minecraft client;
+
+        public Client(Minecraft client) {
+            this.client = client;
+        }
+
+        @Override
+        public Result<ForgeChunkInfo> poll() {
+            Result<ForgeChunkInfo> data = new Result<>();
+
+            ClientLevel level = this.client.level;
+            if (level == null) {
+                return null;
+            }
+
+            TransientEntitySectionManager<Entity> entityManager = level.entityStorage;
+            EntitySectionStorage<Entity> cache = entityManager.sectionStorage;
+
+            List<ForgeChunkInfo> list = getChunksFromCache(cache);
+            data.put(level.dimensionTypeId().location().getPath(), list);
+
+            return data;
+        }
+    }
+
+    static final class ForgeChunkInfo extends AbstractChunkInfo<EntityType<?>> {
+        private final CountMap<EntityType<?>> entityCounts;
+
+        ForgeChunkInfo(long chunkPos, Stream<EntitySection<Entity>> entities) {
+            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
+        public CountMap<EntityType<?>> getEntityCounts() {
+            return this.entityCounts;
+        }
+
+        @Override
+        public String entityTypeName(EntityType<?> type) {
+            return EntityType.getKey(type).toString();
+        }
+    }
+
+
+}
diff --git a/spark-forge/src/main/java/me/lucko/spark/forge/plugin/ForgeClientSparkPlugin.java b/spark-forge/src/main/java/me/lucko/spark/forge/plugin/ForgeClientSparkPlugin.java
index cf5c89b..04c8785 100644
--- a/spark-forge/src/main/java/me/lucko/spark/forge/plugin/ForgeClientSparkPlugin.java
+++ b/spark-forge/src/main/java/me/lucko/spark/forge/plugin/ForgeClientSparkPlugin.java
@@ -28,6 +28,7 @@ import com.mojang.brigadier.suggestion.Suggestions;
 import com.mojang.brigadier.suggestion.SuggestionsBuilder;
 
 import me.lucko.spark.common.platform.PlatformInfo;
+import me.lucko.spark.common.platform.world.WorldInfoProvider;
 import me.lucko.spark.common.tick.TickHook;
 import me.lucko.spark.common.tick.TickReporter;
 import me.lucko.spark.forge.ForgeCommandSender;
@@ -35,6 +36,7 @@ import me.lucko.spark.forge.ForgePlatformInfo;
 import me.lucko.spark.forge.ForgeSparkMod;
 import me.lucko.spark.forge.ForgeTickHook;
 import me.lucko.spark.forge.ForgeTickReporter;
+import me.lucko.spark.forge.ForgeWorldInfoProvider;
 
 import net.minecraft.client.Minecraft;
 import net.minecraft.commands.CommandSource;
@@ -108,6 +110,11 @@ public class ForgeClientSparkPlugin extends ForgeSparkPlugin implements Command<
     }
 
     @Override
+    public void executeSync(Runnable task) {
+        this.minecraft.executeIfPossible(task);
+    }
+
+    @Override
     public TickHook createTickHook() {
         return new ForgeTickHook(TickEvent.Type.CLIENT);
     }
@@ -118,6 +125,11 @@ public class ForgeClientSparkPlugin extends ForgeSparkPlugin implements Command<
     }
 
     @Override
+    public WorldInfoProvider createWorldInfoProvider() {
+        return new ForgeWorldInfoProvider.Client(this.minecraft);
+    }
+
+    @Override
     public PlatformInfo getPlatformInfo() {
         return new ForgePlatformInfo(PlatformInfo.Type.CLIENT);
     }
diff --git a/spark-forge/src/main/java/me/lucko/spark/forge/plugin/ForgeServerSparkPlugin.java b/spark-forge/src/main/java/me/lucko/spark/forge/plugin/ForgeServerSparkPlugin.java
index e341d6f..03f9952 100644
--- a/spark-forge/src/main/java/me/lucko/spark/forge/plugin/ForgeServerSparkPlugin.java
+++ b/spark-forge/src/main/java/me/lucko/spark/forge/plugin/ForgeServerSparkPlugin.java
@@ -31,6 +31,7 @@ import com.mojang.brigadier.suggestion.SuggestionsBuilder;
 
 import me.lucko.spark.common.monitor.ping.PlayerPingProvider;
 import me.lucko.spark.common.platform.PlatformInfo;
+import me.lucko.spark.common.platform.world.WorldInfoProvider;
 import me.lucko.spark.common.tick.TickHook;
 import me.lucko.spark.common.tick.TickReporter;
 import me.lucko.spark.forge.ForgeCommandSender;
@@ -39,6 +40,7 @@ import me.lucko.spark.forge.ForgePlayerPingProvider;
 import me.lucko.spark.forge.ForgeSparkMod;
 import me.lucko.spark.forge.ForgeTickHook;
 import me.lucko.spark.forge.ForgeTickReporter;
+import me.lucko.spark.forge.ForgeWorldInfoProvider;
 
 import net.minecraft.commands.CommandSource;
 import net.minecraft.commands.CommandSourceStack;
@@ -184,6 +186,11 @@ public class ForgeServerSparkPlugin extends ForgeSparkPlugin implements Command<
     }
 
     @Override
+    public void executeSync(Runnable task) {
+        this.server.executeIfPossible(task);
+    }
+
+    @Override
     public TickHook createTickHook() {
         return new ForgeTickHook(TickEvent.Type.SERVER);
     }
@@ -199,6 +206,11 @@ public class ForgeServerSparkPlugin extends ForgeSparkPlugin implements Command<
     }
 
     @Override
+    public WorldInfoProvider createWorldInfoProvider() {
+        return new ForgeWorldInfoProvider.Server(this.server);
+    }
+
+    @Override
     public PlatformInfo getPlatformInfo() {
         return new ForgePlatformInfo(PlatformInfo.Type.SERVER);
     }