Leaf/leaf-server/minecraft-patches/features/0130-Improve-sorting-in-SortedArraySet.patch

From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: Taiyou06 <kaandindar21@gmail.com>
Date: Thu, 20 Feb 2025 15:37:39 +0100
Subject: [PATCH] Improve sorting in SortedArraySet


diff --git a/net/minecraft/util/SortedArraySet.java b/net/minecraft/util/SortedArraySet.java
index 339b19e88567be382e550ed54477fabd58d51faa..bde5b4cb4cda003acd7343b16f09f915b71fe3f2 100644
--- a/net/minecraft/util/SortedArraySet.java
+++ b/net/minecraft/util/SortedArraySet.java
@@ -11,6 +11,7 @@ import javax.annotation.Nullable;
 public class SortedArraySet<T> extends AbstractSet<T> implements ca.spottedleaf.moonrise.patches.chunk_system.util.ChunkSystemSortedArraySet<T> { // Paper - rewrite chunk system
     private static final int DEFAULT_INITIAL_CAPACITY = 10;
     private final Comparator<T> comparator;
+    private final boolean isNaturalOrder;
     T[] contents;
     int size;
 
@@ -93,10 +94,11 @@ public class SortedArraySet<T> extends AbstractSet<T> implements ca.spottedleaf.
 
     private SortedArraySet(int initialCapacity, Comparator<T> comparator) {
         this.comparator = comparator;
+        this.isNaturalOrder = comparator == Comparator.naturalOrder();
         if (initialCapacity < 0) {
             throw new IllegalArgumentException("Initial capacity (" + initialCapacity + ") is negative");
         } else {
-            this.contents = (T[])castRawArray(new Object[initialCapacity]);
+            this.contents = (T[]) castRawArray(new Object[initialCapacity]);
         }
     }
 
@@ -121,7 +123,51 @@ public class SortedArraySet<T> extends AbstractSet<T> implements ca.spottedleaf.
     }
 
     private int findIndex(T object) {
-        return Arrays.binarySearch(this.contents, 0, this.size, object, this.comparator);
+        return isNaturalOrder ? naturalBinarySearch(object) : customBinarySearch(object);
+    }
+
+    private int naturalBinarySearch(T object) {
+        int low = 0;
+        int high = this.size - 1;
+        Comparable<? super T> key = (Comparable<? super T>) object;
+        T[] a = this.contents;
+
+        while (low <= high) {
+            int mid = (low + high) >>> 1;
+            T midVal = a[mid];
+            int cmp = key.compareTo(midVal);
+
+            if (cmp < 0) {
+                high = mid - 1;
+            } else if (cmp > 0) {
+                low = mid + 1;
+            } else {
+                return mid;
+            }
+        }
+        return -(low + 1);
+    }
+
+    private int customBinarySearch(T object) {
+        int low = 0;
+        int high = this.size - 1;
+        T[] a = this.contents;
+        Comparator<T> c = this.comparator;
+
+        while (low <= high) {
+            int mid = (low + high) >>> 1;
+            T midVal = a[mid];
+            int cmp = c.compare(midVal, object);
+
+            if (cmp < 0) {
+                low = mid + 1;
+            } else if (cmp > 0) {
+                high = mid - 1;
+            } else {
+                return mid;
+            }
+        }
+        return -(low + 1);
     }
 
     private static int getInsertionPosition(int index) {