Use AVLTreeMap

1 files changed, 34 insertions, 20 deletions

diff --git a/src/main/java/me/xmrvizzy/skyblocker/utils/scheduler/Scheduler.java b/src/main/java/me/xmrvizzy/skyblocker/utils/scheduler/Scheduler.java
index 3f578899..1d61edd3 100644
--- a/src/main/java/me/xmrvizzy/skyblocker/utils/scheduler/Scheduler.java
+++ b/src/main/java/me/xmrvizzy/skyblocker/utils/scheduler/Scheduler.java
@@ -1,6 +1,10 @@
 package me.xmrvizzy.skyblocker.utils.scheduler;
 
 import com.mojang.brigadier.Command;
+import it.unimi.dsi.fastutil.ints.AbstractInt2ObjectSortedMap;
+import it.unimi.dsi.fastutil.ints.Int2ObjectAVLTreeMap;
+import it.unimi.dsi.fastutil.ints.Int2ObjectLinkedOpenHashMap;
+import it.unimi.dsi.fastutil.ints.Int2ObjectRBTreeMap;
 import me.xmrvizzy.skyblocker.SkyblockerMod;
 import net.fabricmc.fabric.api.client.command.v2.FabricClientCommandSource;
 import net.minecraft.client.MinecraftClient;
@@ -8,7 +12,8 @@ import net.minecraft.client.gui.screen.Screen;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.util.PriorityQueue;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.function.Supplier;
 
 /**
@@ -17,7 +22,7 @@ import java.util.function.Supplier;
 public class Scheduler {
     private static final Logger LOGGER = LoggerFactory.getLogger(Scheduler.class);
     private int currentTick = 0;
-    private final PriorityQueue<ScheduledTask> tasks = new PriorityQueue<>();
+    private final AbstractInt2ObjectSortedMap<List<ScheduledTask>> tasks = new Int2ObjectAVLTreeMap<>();
 
     /**
      * Do not instantiate this class. Use {@link SkyblockerMod#scheduler} instead.
@@ -37,7 +42,7 @@ public class Scheduler {
         if (delay < 0) {
             LOGGER.warn("Scheduled a task with negative delay");
         }
-        tasks.add(new ScheduledTask(task, currentTick + delay));
+        tasks.computeIfAbsent(currentTick + delay, key -> new ArrayList<>()).add(new ScheduledTask(task));
     }
 
     /**
@@ -50,7 +55,7 @@ public class Scheduler {
         if (period <= 0) {
             LOGGER.error("Attempted to schedule a cyclic task with period lower than 1");
         } else {
-            tasks.add(new CyclicTask(this, task, period));
+            tasks.computeIfAbsent(currentTick, key -> new ArrayList<>()).add(new CyclicTask(this, task, period));
         }
     }
 
@@ -71,10 +76,26 @@ public class Scheduler {
 
     public void tick() {
         currentTick += 1;
-        ScheduledTask task;
-        while ((task = tasks.peek()) != null && task.schedule <= currentTick && runTask(task)) {
-            tasks.poll();
-            if (task instanceof CyclicTask) tasks.add(task);
+        if (tasks.containsKey(currentTick - 1)) {
+            List<ScheduledTask> currentTickTasks = tasks.get(currentTick - 1);
+            for (int i = 0; i < currentTickTasks.size(); i++) {
+                ScheduledTask task = currentTickTasks.get(i);
+                if (!runTask(task)) {
+                    tasks.computeIfAbsent(currentTick + 1, key -> new ArrayList<>()).add(task);
+                }
+            }
+            tasks.remove(currentTick - 1);
+        }
+
+        if (tasks.containsKey(currentTick)) {
+            List<ScheduledTask> currentTickTasks = tasks.get(currentTick);
+            for (int i = 0; i < currentTickTasks.size(); i++) {
+                ScheduledTask task = currentTickTasks.get(i);
+                if (!runTask(task)) {
+                    tasks.computeIfAbsent(currentTick + 1, key -> new ArrayList<>()).add(task);
+                }
+            }
+            tasks.remove(currentTick);
         }
     }
 
@@ -92,12 +113,12 @@ public class Scheduler {
     /**
      * A task that runs every period ticks. More specifically, this task reschedules itself to run again after period ticks every time it runs.
      */
-    protected static class CyclicTask extends ScheduledTask {
+    protected class CyclicTask extends ScheduledTask {
         private final Scheduler scheduler;
         private final int period;
 
         CyclicTask(Scheduler scheduler, Runnable inner, int period) {
-            super(inner, scheduler.currentTick);
+            super(inner);
             this.scheduler = scheduler;
             this.period = period;
         }
@@ -105,25 +126,18 @@ public class Scheduler {
         @Override
         public void run() {
             super.run();
-            schedule = scheduler.currentTick + period;
+            tasks.computeIfAbsent(scheduler.currentTick + period, key -> new ArrayList<>()).add(this);
         }
     }
 
     /**
      * A task that runs at a specific tick, relative to {@link #currentTick}.
      */
-    protected static class ScheduledTask implements Comparable<ScheduledTask>, Runnable {
+    protected static class ScheduledTask implements Runnable {
         private final Runnable inner;
-        protected int schedule;
 
-        public ScheduledTask(Runnable inner, int schedule) {
+        public ScheduledTask(Runnable inner) {
             this.inner = inner;
-            this.schedule = schedule;
-        }
-
-        @Override
-        public int compareTo(ScheduledTask o) {
-            return schedule - o.schedule;
         }
 
         @Override