From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: etil2jz <81570777+etil2jz@users.noreply.github.com>
Date: Fri, 24 Jun 2022 21:17:05 +0200
Subject: [PATCH] lithium: world.tick_scheduler

Original license: GPLv3
Original project: https://github.com/CaffeineMC/lithium-fabric (Yarn mappings)

diff --git a/src/main/java/me/jellysquid/mods/lithium/common/world/scheduler/OrderedTickQueue.java b/src/main/java/me/jellysquid/mods/lithium/common/world/scheduler/OrderedTickQueue.java
new file mode 100644
index 0000000000000000000000000000000000000000..52e90ea09b44543af661c214767073cf9d5f3d8f
--- /dev/null
+++ b/src/main/java/me/jellysquid/mods/lithium/common/world/scheduler/OrderedTickQueue.java
@@ -0,0 +1,192 @@
+package me.jellysquid.mods.lithium.common.world.scheduler;
+
+import it.unimi.dsi.fastutil.HashCommon;
+import java.util.*;
+import net.minecraft.world.ticks.ScheduledTick;
+
+/**
+
+ */
+public class OrderedTickQueue<T> extends AbstractQueue<ScheduledTick<T>> {
+    private static final int INITIAL_CAPACITY = 16;
+    private static final Comparator<ScheduledTick<?>> COMPARATOR = Comparator.comparingLong(ScheduledTick::subTickOrder);
+
+    private ScheduledTick<T>[] arr;
+
+    private int lastIndexExclusive;
+    private int firstIndex;
+
+    private long currentMaxSubTickOrder = Long.MIN_VALUE;
+    private boolean isSorted;
+    private ScheduledTick<T> unsortedPeekResult;
+
+    @SuppressWarnings("unchecked")
+    public OrderedTickQueue(int capacity) {
+        this.arr = (ScheduledTick<T>[]) new ScheduledTick[capacity];
+        this.lastIndexExclusive = 0;
+        this.isSorted = true;
+        this.unsortedPeekResult = null;
+        this.firstIndex = 0;
+    }
+
+    public OrderedTickQueue() {
+        this(INITIAL_CAPACITY);
+    }
+
+    @Override
+    public void clear() {
+        Arrays.fill(this.arr, null);
+        this.lastIndexExclusive = 0;
+        this.firstIndex = 0;
+        this.currentMaxSubTickOrder = Long.MIN_VALUE;
+        this.isSorted = true;
+        this.unsortedPeekResult = null;
+    }
+
+    @Override
+    public Iterator<ScheduledTick<T>> iterator() {
+        if (this.isEmpty()) {
+            return Collections.emptyIterator();
+        }
+        this.sort();
+        return new Iterator<>() {
+            int nextIndex = OrderedTickQueue.this.firstIndex;
+
+            @Override
+            public boolean hasNext() {
+                return this.nextIndex < OrderedTickQueue.this.lastIndexExclusive;
+            }
+
+            @Override
+            public ScheduledTick<T> next() {
+                return OrderedTickQueue.this.arr[this.nextIndex++];
+            }
+        };
+    }
+
+    @Override
+    public ScheduledTick<T> poll() {
+        if (this.isEmpty()) {
+            return null;
+        }
+        if (!this.isSorted) {
+            this.sort();
+        }
+        ScheduledTick<T> nextTick;
+        int polledIndex = this.firstIndex++;
+        ScheduledTick<T>[] ticks = this.arr;
+        nextTick = ticks[polledIndex];
+        ticks[polledIndex] = null;
+        return nextTick;
+    }
+
+    @Override
+    public ScheduledTick<T> peek() {
+        if (!this.isSorted) {
+            return this.unsortedPeekResult;
+        } else if (this.lastIndexExclusive > this.firstIndex) {
+            return this.getTickAtIndex(this.firstIndex);
+        }
+        return null;
+    }
+
+    public boolean offer(ScheduledTick<T> tick) {
+        if (this.lastIndexExclusive >= this.arr.length) {
+            //todo remove consumed elements first
+            this.arr = copyArray(this.arr, HashCommon.nextPowerOfTwo(this.arr.length + 1));
+        }
+        if (tick.subTickOrder() <= this.currentMaxSubTickOrder) {
+            //Set to unsorted instead of slowing down the insertion
+            //This is rare but may happen in bulk
+            //Sorting later needs O(n*log(n)) time, but it only needs to happen when unordered insertion needs to happen
+            //Therefore it is better than n times log(n) time of the PriorityQueue that happens on ordered insertion too
+            this.isSorted = false;
+            ScheduledTick<T> firstTick = this.size() > 0 ? this.arr[this.firstIndex] : null;
+            this.unsortedPeekResult = firstTick == null || tick.subTickOrder() < firstTick.subTickOrder() ? tick : firstTick;
+        } else {
+            this.currentMaxSubTickOrder = tick.subTickOrder();
+        }
+        this.arr[this.lastIndexExclusive++] = tick;
+        return true;
+    }
+
+    public int size() {
+        return this.lastIndexExclusive - this.firstIndex;
+    }
+
+    private void resize(int size) {
+        // Only compact the array if it is completely empty or is less than 50% filled
+        if (size == 0 || size < this.arr.length / 2) {
+            this.arr = copyArray(this.arr, size);
+        } else {
+            // Fill the unused array elements with nulls to release our references to the elements in it
+            for (int i = size; i < this.arr.length; i++) {
+                this.arr[i] = null;
+            }
+        }
+
+        this.firstIndex = 0;
+        this.lastIndexExclusive = size;
+
+        if (size == 0 || !this.isSorted) {
+            this.currentMaxSubTickOrder = Long.MIN_VALUE;
+        } else {
+            ScheduledTick<T> tick = this.arr[size - 1];
+            this.currentMaxSubTickOrder = tick == null ? Long.MIN_VALUE : tick.subTickOrder();
+        }
+    }
+
+    public void sort() {
+        if (this.isSorted) {
+            return;
+        }
+        this.removeNullsAndConsumed();
+        Arrays.sort(this.arr, this.firstIndex, this.lastIndexExclusive, COMPARATOR);
+        this.isSorted = true;
+        this.unsortedPeekResult = null;
+    }
+
+    public void removeNullsAndConsumed() {
+        int src = this.firstIndex;
+        int dst = 0;
+        while (src < this.lastIndexExclusive) {
+            ScheduledTick<T> orderedTick = this.arr[src];
+            if (orderedTick != null) {
+                this.arr[dst] = orderedTick;
+                dst++;
+            }
+            src++;
+        }
+        this.resize(dst);
+    }
+
+    public ScheduledTick<T> getTickAtIndex(int index) {
+        if (!this.isSorted) {
+            throw new IllegalStateException("Unexpected access on unsorted queue!");
+        }
+        return this.arr[index];
+    }
+
+    public void setTickAtIndex(int index, ScheduledTick<T> tick) {
+        if (!this.isSorted) {
+            throw new IllegalStateException("Unexpected access on unsorted queue!");
+        }
+        this.arr[index] = tick;
+    }
+
+    @SuppressWarnings("unchecked")
+    private static <T> ScheduledTick<T>[] copyArray(ScheduledTick<T>[] src, int size) {
+        final ScheduledTick<T>[] copy = new ScheduledTick[size];
+
+        if (size != 0) {
+            System.arraycopy(src, 0, copy, 0, Math.min(src.length, size));
+        }
+
+        return copy;
+    }
+
+    @Override
+    public boolean isEmpty() {
+        return this.lastIndexExclusive <= this.firstIndex;
+    }
+}
diff --git a/src/main/java/net/minecraft/world/ticks/LevelChunkTicks.java b/src/main/java/net/minecraft/world/ticks/LevelChunkTicks.java
index 9f6c2e5b5d9e8d714a47c770e255d06c0ef7c190..826ced345c97bd2eb04749f42744a086fafc4ce8 100644
--- a/src/main/java/net/minecraft/world/ticks/LevelChunkTicks.java
+++ b/src/main/java/net/minecraft/world/ticks/LevelChunkTicks.java
@@ -16,14 +16,37 @@ import javax.annotation.Nullable;
 import net.minecraft.core.BlockPos;
 import net.minecraft.nbt.ListTag;
 import net.minecraft.world.level.ChunkPos;
+// Mirai start - lithium: world.tick_scheduler
+import it.unimi.dsi.fastutil.ints.IntOpenHashSet;
+import it.unimi.dsi.fastutil.longs.Long2ReferenceAVLTreeMap;
+import it.unimi.dsi.fastutil.objects.ObjectIterator;
+import it.unimi.dsi.fastutil.objects.Reference2IntOpenHashMap;
+import me.jellysquid.mods.lithium.common.world.scheduler.OrderedTickQueue;
+import net.minecraft.world.ticks.SavedTick;
+import net.minecraft.world.ticks.ScheduledTick;
+import net.minecraft.world.ticks.TickPriority;
+import java.util.Collection;
+// Mirai end
 
 public class LevelChunkTicks<T> implements SerializableTickContainer<T>, TickContainerAccess<T> {
-    private final Queue<ScheduledTick<T>> tickQueue = new PriorityQueue<>(ScheduledTick.DRAIN_ORDER);
+    private Queue<ScheduledTick<T>> tickQueue = new PriorityQueue<>(ScheduledTick.DRAIN_ORDER); // Mirai - remove final
     @Nullable
     private List<SavedTick<T>> pendingTicks;
-    private final Set<ScheduledTick<?>> ticksPerPosition = new ObjectOpenCustomHashSet<>(ScheduledTick.UNIQUE_TICK_HASH);
+    private Set<ScheduledTick<?>> ticksPerPosition = new ObjectOpenCustomHashSet<>(ScheduledTick.UNIQUE_TICK_HASH); // Mirai - remove final
     @Nullable
     private BiConsumer<LevelChunkTicks<T>, ScheduledTick<T>> onTickAdded;
+    // Mirai start - lithium: world.tick_scheduler
+    private static volatile Reference2IntOpenHashMap<Object> TYPE_2_INDEX;
+
+    static {
+        TYPE_2_INDEX = new Reference2IntOpenHashMap<>();
+        TYPE_2_INDEX.defaultReturnValue(-1);
+    }
+
+    private final Long2ReferenceAVLTreeMap<OrderedTickQueue<T>> tickQueuesByTimeAndPriority = new Long2ReferenceAVLTreeMap<>();
+    private OrderedTickQueue<T> nextTickQueue;
+    private final IntOpenHashSet allpendingTicks = new IntOpenHashSet();
+    // Mirai end
 
     public LevelChunkTicks() {
     }
@@ -35,34 +58,133 @@ public class LevelChunkTicks<T> implements SerializableTickContainer<T>, TickCon
             this.ticksPerPosition.add(ScheduledTick.probe(savedTick.type(), savedTick.pos()));
         }
 
+        // Mirai start - lithium: world.tick_scheduler
+        //Remove replaced collections
+        if (this.pendingTicks != null) {
+            for (SavedTick<?> orderedTick : this.pendingTicks) {
+                this.allpendingTicks.add(tickToInt(orderedTick.pos(), orderedTick.type()));
+            }
+        }
+        this.ticksPerPosition = null;
+        this.tickQueue = null;
+        // Mirai end
+    }
+
+    // Mirai start - lithium: world.tick_scheduler
+    private static int tickToInt(BlockPos pos, Object type) {
+        //Y coordinate is 12 bits (BlockPos.toLong)
+        //X and Z coordinate is 4 bits each (This scheduler is for a single chunk)
+        //20 bits are in use for pos
+        //12 bits remaining for the type, so up to 4096 different tickable blocks/fluids (not block states) -> can upgrade to long if needed
+        int typeIndex = TYPE_2_INDEX.getInt(type);
+        if (typeIndex == -1) {
+            typeIndex = fixMissingType2Index(type);
+        }
+
+        int ret = ((pos.getX() & 0xF) << 16) | ((pos.getY() & (0xfff)) << 4) | (pos.getZ() & 0xF);
+        ret |= typeIndex << 20;
+        return ret;
     }
 
+    //This method must be synchronized, otherwise type->int assignments can be overwritten and therefore change
+    //Uses clone and volatile store to ensure only fully initialized maps are used, all threads share the same mapping
+    private static synchronized int fixMissingType2Index(Object type) {
+        //check again, other thread might have replaced the collection
+        int typeIndex = TYPE_2_INDEX.getInt(type);
+        if (typeIndex == -1) {
+            Reference2IntOpenHashMap<Object> clonedType2Index = TYPE_2_INDEX.clone();
+            clonedType2Index.put(type, typeIndex = clonedType2Index.size());
+            TYPE_2_INDEX = clonedType2Index;
+            if (typeIndex >= 4096) {
+                throw new IllegalStateException("Lithium Tick Scheduler assumes at most 4096 different block types that receive scheduled pendingTicks exist! Open an issue on GitHub if you see this error!");
+            }
+        }
+        return typeIndex;
+    }
+    // Mirai end
+
     public void setOnTickAdded(@Nullable BiConsumer<LevelChunkTicks<T>, ScheduledTick<T>> tickConsumer) {
         this.onTickAdded = tickConsumer;
     }
 
+    // Mirai start - lithium: world.tick_scheduler
+    /**
+     * @author 2No2Name
+     * @reason use faster collections
+     */
     @Nullable
     public ScheduledTick<T> peek() {
-        return this.tickQueue.peek();
+        if (this.nextTickQueue == null) {
+            return null;
+        }
+        return this.nextTickQueue.peek();
     }
 
+    /**
+     * @author 2No2Name
+     * @reason use faster collections
+     */
     @Nullable
     public ScheduledTick<T> poll() {
-        ScheduledTick<T> scheduledTick = this.tickQueue.poll();
-        if (scheduledTick != null) {
-            this.ticksPerPosition.remove(scheduledTick);
+        ScheduledTick<T> orderedTick = this.nextTickQueue.poll();
+        if (orderedTick != null) {
+            if (this.nextTickQueue.isEmpty()) {
+                this.updateNextTickQueue(true);
+            }
+            this.allpendingTicks.remove(tickToInt(orderedTick.pos(), orderedTick.type()));
+            return orderedTick;
         }
-
-        return scheduledTick;
+        return null;
     }
 
+    /**
+     * @author 2No2Name
+     * @reason use faster collections
+     */
     @Override
     public void schedule(ScheduledTick<T> orderedTick) {
-        if (this.ticksPerPosition.add(orderedTick)) {
-            this.scheduleUnchecked(orderedTick);
+        int intTick = tickToInt(orderedTick.pos(), orderedTick.type());
+        if (this.allpendingTicks.add(intTick)) {
+            this.queueTick(orderedTick);
+        }
+    }
+
+    // Computes a timestamped key including the tick's priority
+    // Keys can be sorted in descending order to find what should be executed first
+    // 60 time bits, 4 priority bits
+    private static long getBucketKey(long time, TickPriority priority) {
+        //using priority.ordinal() as is not negative instead of priority.index
+        return (time << 4L) | (priority.ordinal() & 15);
+    }
+
+    private void updateNextTickQueue(boolean elementRemoved) {
+        if (elementRemoved && this.nextTickQueue != null && this.nextTickQueue.isEmpty()) {
+            OrderedTickQueue<T> removed = this.tickQueuesByTimeAndPriority.remove(this.tickQueuesByTimeAndPriority.firstLongKey());
+            if (removed != this.nextTickQueue) {
+                throw new IllegalStateException("Next tick queue doesn't have the lowest key!");
+            }
+        }
+        if (this.tickQueuesByTimeAndPriority.isEmpty()) {
+            this.nextTickQueue = null;
+            return;
         }
+        long firstKey = this.tickQueuesByTimeAndPriority.firstLongKey();
+        this.nextTickQueue = this.tickQueuesByTimeAndPriority.get(firstKey);
+    }
+
+    private void queueTick(ScheduledTick<T> orderedTick) {
+        OrderedTickQueue<T> tickQueue = this.tickQueuesByTimeAndPriority.computeIfAbsent(getBucketKey(orderedTick.triggerTick(), orderedTick.priority()), key -> new OrderedTickQueue<>());
+        if (tickQueue.isEmpty()) {
+            this.updateNextTickQueue(false);
+        }
+        tickQueue.offer(orderedTick);
 
+        if (this.onTickAdded != null) {
+            //noinspection unchecked
+            this.onTickAdded.accept((LevelChunkTicks<T>) (Object) this, orderedTick);
+        }
     }
+    // Mirai end
 
     private void scheduleUnchecked(ScheduledTick<T> orderedTick) {
         this.tickQueue.add(orderedTick);
@@ -72,60 +194,93 @@ public class LevelChunkTicks<T> implements SerializableTickContainer<T>, TickCon
 
     }
 
+    // Mirai start - lithium: world.tick_scheduler
+    /**
+     * @author 2No2Name
+     * @reason use faster collections
+     */
     @Override
     public boolean hasScheduledTick(BlockPos pos, T type) {
-        return this.ticksPerPosition.contains(ScheduledTick.probe(type, pos));
+        return this.allpendingTicks.contains(tickToInt(pos, type));
     }
 
+    /**
+     * @author 2No2Name
+     * @reason use faster collections
+     */
     public void removeIf(Predicate<ScheduledTick<T>> predicate) {
-        Iterator<ScheduledTick<T>> iterator = this.tickQueue.iterator();
-
-        while(iterator.hasNext()) {
-            ScheduledTick<T> scheduledTick = iterator.next();
-            if (predicate.test(scheduledTick)) {
-                iterator.remove();
-                this.ticksPerPosition.remove(scheduledTick);
+        for (ObjectIterator<OrderedTickQueue<T>> tickQueueIterator = this.tickQueuesByTimeAndPriority.values().iterator(); tickQueueIterator.hasNext(); ) {
+            OrderedTickQueue<T> nextTickQueue = tickQueueIterator.next();
+            nextTickQueue.sort();
+            boolean removed = false;
+            for (int i = 0; i < nextTickQueue.size(); i++) {
+                ScheduledTick<T> nextTick = nextTickQueue.getTickAtIndex(i);
+                if (predicate.test(nextTick)) {
+                    nextTickQueue.setTickAtIndex(i, null);
+                    this.allpendingTicks.remove(tickToInt(nextTick.pos(), nextTick.type()));
+                    removed = true;
+                }
+            }
+            if (removed) {
+                nextTickQueue.removeNullsAndConsumed();
+            }
+            if (nextTickQueue.isEmpty()) {
+                tickQueueIterator.remove();
             }
         }
-
     }
 
+    /**
+     * @author 2No2Name
+     * @reason use faster collections
+     */
     public Stream<ScheduledTick<T>> getAll() {
-        return this.tickQueue.stream();
+        return this.tickQueuesByTimeAndPriority.values().stream().flatMap(Collection::stream);
     }
 
+    /**
+     * @author 2No2Name
+     * @reason not use unused field
+     */
     @Override
     public int count() {
-        return this.tickQueue.size() + (this.pendingTicks != null ? this.pendingTicks.size() : 0);
+        return this.allpendingTicks.size();
     }
 
+    /**
+     * @author 2No2Name
+     * @reason not use unused field
+     */
     @Override
     public ListTag save(long l, Function<T, String> function) {
-        ListTag listTag = new ListTag();
+        ListTag nbtList = new ListTag();
         if (this.pendingTicks != null) {
-            for(SavedTick<T> savedTick : this.pendingTicks) {
-                listTag.add(savedTick.save(function));
+            for (SavedTick<T> tick : this.pendingTicks) {
+                nbtList.add(tick.save(function));
             }
         }
-
-        for(ScheduledTick<T> scheduledTick : this.tickQueue) {
-            listTag.add(SavedTick.saveTick(scheduledTick, function, l));
+        for (OrderedTickQueue<T> nextTickQueue : this.tickQueuesByTimeAndPriority.values()) {
+            for (ScheduledTick<T> orderedTick : nextTickQueue) {
+                nbtList.add(SavedTick.saveTick(orderedTick, function, l));
+            }
         }
-
-        return listTag;
+        return nbtList;
     }
 
+    /**
+     * @author 2No2Name
+     * @reason use our datastructures
+     */
     public void unpack(long time) {
         if (this.pendingTicks != null) {
             int i = -this.pendingTicks.size();
-
-            for(SavedTick<T> savedTick : this.pendingTicks) {
-                this.scheduleUnchecked(savedTick.unpack(time, (long)(i++)));
+            for (SavedTick<T> tick : this.pendingTicks) {
+                this.queueTick(tick.unpack(time, i++));
             }
         }
-
         this.pendingTicks = null;
     }
+    // Mirai end
 
     public static <T> LevelChunkTicks<T> load(ListTag tickQueue, Function<String, Optional<T>> nameToTypeFunction, ChunkPos pos) {
         ImmutableList.Builder<SavedTick<T>> builder = ImmutableList.builder();