Add providers for world (entity/chunk) statistics

7 files changed, 313 insertions, 0 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
new file mode 100644
index 0000000..fddcf58
--- /dev/null
+++ b/spark-fabric/src/main/java/me/lucko/spark/fabric/FabricWorldInfoProvider.java
@@ -0,0 +1,145 @@
+/*
+ * 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;
+
+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 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 net.minecraft.client.MinecraftClient;
+import net.minecraft.client.world.ClientEntityManager;
+import net.minecraft.client.world.ClientWorld;
+import net.minecraft.entity.Entity;
+import net.minecraft.entity.EntityType;
+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.EntityTrackingSection;
+import net.minecraft.world.entity.SectionedEntityCache;
+import net.minecraft.world.level.ServerWorldProperties;
+
+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;
+
+        public Server(MinecraftServer server) {
+            this.server = server;
+        }
+
+        @Override
+        public Result<FabricChunkInfo> poll() {
+            Result<FabricChunkInfo> data = new Result<>();
+
+            for (ServerWorld world : this.server.getWorlds()) {
+                ServerEntityManager<Entity> entityManager = ((ServerWorldAccessor) world).getEntityManager();
+                SectionedEntityCache<Entity> cache = ((ServerEntityManagerAccessor) entityManager).getCache();
+
+                List<FabricChunkInfo> list = getChunksFromCache(cache);
+                data.put(((ServerWorldProperties) world.getLevelProperties()).getLevelName(), list);
+            }
+
+            return data;
+        }
+    }
+
+    public static final class Client extends FabricWorldInfoProvider {
+        private final MinecraftClient client;
+
+        public Client(MinecraftClient client) {
+            this.client = client;
+        }
+
+        @Override
+        public Result<FabricChunkInfo> poll() {
+            Result<FabricChunkInfo> data = new Result<>();
+
+            ClientWorld world = this.client.world;
+            if (world == null) {
+                return null;
+            }
+
+            ClientEntityManager<Entity> entityManager = ((ClientWorldAccessor) world).getEntityManager();
+            SectionedEntityCache<Entity> cache = ((ClientEntityManagerAccessor) entityManager).getCache();
+
+            List<FabricChunkInfo> list = getChunksFromCache(cache);
+            data.put(world.getDimensionKey().getValue().getPath(), list);
+
+            return data;
+        }
+    }
+
+    static final class FabricChunkInfo extends AbstractChunkInfo<EntityType<?>> {
+        private final CountMap<EntityType<?>> entityCounts;
+
+        FabricChunkInfo(long chunkPos, Stream<EntityTrackingSection<Entity>> entities) {
+            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
+        public CountMap<EntityType<?>> getEntityCounts() {
+            return this.entityCounts;
+        }
+
+        @Override
+        public String entityTypeName(EntityType<?> type) {
+            return EntityType.getId(type).toString();
+        }
+    }
+
+}
+
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
new file mode 100644
index 0000000..88c9521
--- /dev/null
+++ b/spark-fabric/src/main/java/me/lucko/spark/fabric/mixin/ClientEntityManagerAccessor.java
@@ -0,0 +1,36 @@
+/*
+ * 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.client.world.ClientEntityManager;
+import net.minecraft.entity.Entity;
+import net.minecraft.world.entity.SectionedEntityCache;
+
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.gen.Accessor;
+
+@Mixin(ClientEntityManager.class)
+public interface ClientEntityManagerAccessor {
+
+    @Accessor
+    SectionedEntityCache<Entity> getCache();
+
+}
diff --git a/spark-fabric/src/main/java/me/lucko/spark/fabric/mixin/ClientWorldAccessor.java b/spark-fabric/src/main/java/me/lucko/spark/fabric/mixin/ClientWorldAccessor.java
new file mode 100644
index 0000000..01562ef
--- /dev/null
+++ b/spark-fabric/src/main/java/me/lucko/spark/fabric/mixin/ClientWorldAccessor.java
@@ -0,0 +1,36 @@
+/*
+ * 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.client.world.ClientEntityManager;
+import net.minecraft.client.world.ClientWorld;
+import net.minecraft.entity.Entity;
+
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.gen.Accessor;
+
+@Mixin(ClientWorld.class)
+public interface ClientWorldAccessor {
+
+    @Accessor
+    ClientEntityManager<Entity> getEntityManager();
+
+}
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
new file mode 100644
index 0000000..160a12b
--- /dev/null
+++ b/spark-fabric/src/main/java/me/lucko/spark/fabric/mixin/ServerEntityManagerAccessor.java
@@ -0,0 +1,36 @@
+/*
+ * 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.server.world.ServerEntityManager;
+import net.minecraft.world.entity.SectionedEntityCache;
+
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.gen.Accessor;
+
+@Mixin(ServerEntityManager.class)
+public interface ServerEntityManagerAccessor {
+
+    @Accessor
+    SectionedEntityCache<Entity> getCache();
+
+}
diff --git a/spark-fabric/src/main/java/me/lucko/spark/fabric/mixin/ServerWorldAccessor.java b/spark-fabric/src/main/java/me/lucko/spark/fabric/mixin/ServerWorldAccessor.java
new file mode 100644
index 0000000..cf2e7e8
--- /dev/null
+++ b/spark-fabric/src/main/java/me/lucko/spark/fabric/mixin/ServerWorldAccessor.java
@@ -0,0 +1,36 @@
+/*
+ * 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.server.world.ServerEntityManager;
+import net.minecraft.server.world.ServerWorld;
+
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.gen.Accessor;
+
+@Mixin(ServerWorld.class)
+public interface ServerWorldAccessor {
+
+    @Accessor
+    ServerEntityManager<Entity> getEntityManager();
+
+}
diff --git a/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricClientSparkPlugin.java b/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricClientSparkPlugin.java
index e94d697..1876658 100644
--- a/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricClientSparkPlugin.java
+++ b/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricClientSparkPlugin.java
@@ -29,6 +29,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.fabric.FabricCommandSender;
@@ -36,6 +37,7 @@ import me.lucko.spark.fabric.FabricPlatformInfo;
 import me.lucko.spark.fabric.FabricSparkMod;
 import me.lucko.spark.fabric.FabricTickHook;
 import me.lucko.spark.fabric.FabricTickReporter;
+import me.lucko.spark.fabric.FabricWorldInfoProvider;
 
 import net.fabricmc.fabric.api.client.command.v2.ClientCommandRegistrationCallback;
 import net.fabricmc.fabric.api.client.command.v2.FabricClientCommandSource;
@@ -113,6 +115,11 @@ public class FabricClientSparkPlugin extends FabricSparkPlugin implements Comman
     }
 
     @Override
+    public void executeSync(Runnable task) {
+        this.minecraft.executeSync(task);
+    }
+
+    @Override
     public TickHook createTickHook() {
         return new FabricTickHook.Client();
     }
@@ -123,6 +130,11 @@ public class FabricClientSparkPlugin extends FabricSparkPlugin implements Comman
     }
 
     @Override
+    public WorldInfoProvider createWorldInfoProvider() {
+        return new FabricWorldInfoProvider.Client(this.minecraft);
+    }
+
+    @Override
     public PlatformInfo getPlatformInfo() {
         return new FabricPlatformInfo(PlatformInfo.Type.CLIENT);
     }
diff --git a/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricServerSparkPlugin.java b/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricServerSparkPlugin.java
index 3d1a0e7..2283a84 100644
--- a/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricServerSparkPlugin.java
+++ b/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricServerSparkPlugin.java
@@ -31,6 +31,7 @@ import com.mojang.brigadier.suggestion.SuggestionsBuilder;
 import me.lucko.fabric.api.permissions.v0.Permissions;
 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.fabric.FabricCommandSender;
@@ -39,6 +40,7 @@ import me.lucko.spark.fabric.FabricPlayerPingProvider;
 import me.lucko.spark.fabric.FabricSparkMod;
 import me.lucko.spark.fabric.FabricTickHook;
 import me.lucko.spark.fabric.FabricTickReporter;
+import me.lucko.spark.fabric.FabricWorldInfoProvider;
 import me.lucko.spark.fabric.placeholder.SparkFabricPlaceholderApi;
 
 import net.fabricmc.loader.api.FabricLoader;
@@ -127,6 +129,11 @@ public class FabricServerSparkPlugin extends FabricSparkPlugin implements Comman
     }
 
     @Override
+    public void executeSync(Runnable task) {
+        this.server.executeSync(task);
+    }
+
+    @Override
     public TickHook createTickHook() {
         return new FabricTickHook.Server();
     }
@@ -142,6 +149,11 @@ public class FabricServerSparkPlugin extends FabricSparkPlugin implements Comman
     }
 
     @Override
+    public WorldInfoProvider createWorldInfoProvider() {
+        return new FabricWorldInfoProvider.Server(this.server);
+    }
+
+    @Override
     public PlatformInfo getPlatformInfo() {
         return new FabricPlatformInfo(PlatformInfo.Type.SERVER);
     }