From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: Samsuik <kfian294ma4@gmail.com>
Date: Fri, 13 Sep 2024 17:22:51 +0100
Subject: [PATCH] Optimise entity scheduler ticking


diff --git a/ca/spottedleaf/moonrise/patches/chunk_system/level/entity/server/ServerEntityLookup.java b/ca/spottedleaf/moonrise/patches/chunk_system/level/entity/server/ServerEntityLookup.java
index 26207443b1223119c03db478d7e816d9cdf8e618..1664830a49f37825c39fb6b436011d8149196e3a 100644
--- a/ca/spottedleaf/moonrise/patches/chunk_system/level/entity/server/ServerEntityLookup.java
+++ b/ca/spottedleaf/moonrise/patches/chunk_system/level/entity/server/ServerEntityLookup.java
@@ -18,6 +18,22 @@ public final class ServerEntityLookup extends EntityLookup {
 
     private final ServerLevel serverWorld;
     public final ReferenceList<Entity> trackerEntities = new ReferenceList<>(EMPTY_ENTITY_ARRAY); // Moonrise - entity tracker
+    // Sakura start - optimise entity scheduler ticking
+    public final ReferenceList<org.bukkit.craftbukkit.entity.CraftEntity> scheduledEntities = new ReferenceList<>();
+
+    public void entityStartScheduled(final Entity entity) {
+        org.bukkit.craftbukkit.entity.CraftEntity bukkitEntity = entity.getBukkitEntityRaw();
+        if (bukkitEntity != null && bukkitEntity.taskScheduler.hasTask()) {
+            this.scheduledEntities.add(bukkitEntity);
+        }
+    }
+
+    public void entityEndScheduled(final Entity entity) {
+        if (entity.getBukkitEntityRaw() != null) {
+            this.scheduledEntities.remove(entity.getBukkitEntityRaw());
+        }
+    }
+    // Sakura end - optimise entity scheduler ticking
 
     public ServerEntityLookup(final ServerLevel world, final LevelCallback<Entity> worldCallback) {
         super(world, worldCallback);
@@ -89,6 +105,7 @@ public final class ServerEntityLookup extends EntityLookup {
         // Moonrise start - entity tracker
         this.trackerEntities.add(entity);
         // Moonrise end - entity tracker
+        this.entityStartScheduled(entity); // Sakura - optimise entity scheduler ticking
     }
 
     @Override
@@ -96,6 +113,7 @@ public final class ServerEntityLookup extends EntityLookup {
         // Moonrise start - entity tracker
         this.trackerEntities.remove(entity);
         // Moonrise end - entity tracker
+        this.entityEndScheduled(entity); // Sakura - optimise entity scheduler ticking
     }
 
     @Override
diff --git a/net/minecraft/server/MinecraftServer.java b/net/minecraft/server/MinecraftServer.java
index b234215e59ac303ef7979031ab1b042ff6692b7a..4898079bb4cceeb1106e1eac08f92417bdb2daf7 100644
--- a/net/minecraft/server/MinecraftServer.java
+++ b/net/minecraft/server/MinecraftServer.java
@@ -1676,7 +1676,17 @@ public abstract class MinecraftServer extends ReentrantBlockableEventLoop<TickTa
         // Paper start - Folia scheduler API
         ((io.papermc.paper.threadedregions.scheduler.FoliaGlobalRegionScheduler) org.bukkit.Bukkit.getGlobalRegionScheduler()).tick();
         getAllLevels().forEach(level -> {
-            for (final net.minecraft.world.entity.Entity entity : level.getEntities().getAll()) {
+            // Sakura start - optimise entity scheduler ticking
+            final ca.spottedleaf.moonrise.patches.chunk_system.level.entity.server.ServerEntityLookup entityLookup = (ca.spottedleaf.moonrise.patches.chunk_system.level.entity.server.ServerEntityLookup) level.moonrise$getEntityLookup();
+            final java.util.Iterator<org.bukkit.craftbukkit.entity.CraftEntity> entityIterator = entityLookup.scheduledEntities.iterator();
+            while (entityIterator.hasNext()) {
+                final org.bukkit.craftbukkit.entity.CraftEntity scheduledEntity = entityIterator.next();
+                final net.minecraft.world.entity.Entity entity = scheduledEntity.getHandle();
+                if (!scheduledEntity.taskScheduler.hasTask()) {
+                    entityIterator.remove();
+                    continue;
+                }
+                // Sakura end - optimise entity scheduler ticking
                 if (entity.isRemoved()) {
                     continue;
                 }