From 389d0648b11d4798dceae16ce22da9437418b884 Mon Sep 17 00:00:00 2001
From: DBotThePony <dbotthepony@yandex.ru>
Date: Thu, 10 Apr 2025 10:12:01 +0700
Subject: [PATCH] Allow allocating indices sparsingly in IDAllocator

---
 .../ru/dbotthepony/mc/otm/util/IDAllocator.kt | 49 +++++++++++++------
 1 file changed, 35 insertions(+), 14 deletions(-)

diff --git a/src/main/kotlin/ru/dbotthepony/mc/otm/util/IDAllocator.kt b/src/main/kotlin/ru/dbotthepony/mc/otm/util/IDAllocator.kt
index 9fa5b05ff..cc197fb62 100644
--- a/src/main/kotlin/ru/dbotthepony/mc/otm/util/IDAllocator.kt
+++ b/src/main/kotlin/ru/dbotthepony/mc/otm/util/IDAllocator.kt
@@ -1,37 +1,58 @@
 package ru.dbotthepony.mc.otm.util
 
-import it.unimi.dsi.fastutil.ints.IntAVLTreeSet
+import it.unimi.dsi.fastutil.ints.IntRBTreeSet
 
 class IDAllocator {
-	private var highestID = 0
-	private val gaps = IntAVLTreeSet()
+	private var nextID = -1
+	private val gaps = IntRBTreeSet()
+	private val sparseAllocations = IntRBTreeSet()
+
+	private fun increaseHighestID(): Int {
+		while (sparseAllocations.remove(++nextID)) {}
+		return nextID
+	}
 
 	fun allocate(): Int {
 		if (gaps.isEmpty()) {
-			return highestID++
+			return increaseHighestID()
 		} else {
 			return gaps.removeFirst()
 		}
 	}
 
+	fun allocate(id: Int): Boolean {
+		if (id < 0) {
+			return true // not handling negative IDs
+		} else if (id > nextID) {
+			sparseAllocations.add(id)
+			return true
+		} else if (id == nextID) {
+			increaseHighestID()
+			return true
+		} else {
+			return gaps.remove(id)
+		}
+	}
+
 	fun release(id: Int) {
-		if (id >= 0) {
-			throw IllegalArgumentException("Invalid ID: $id")
-		} else if (id >= highestID) {
-			throw IllegalArgumentException("Not tracking ID: $id (highest known ID is ${highestID - 1})")
+		if (id < 0) {
+			return // not handling negative IDs
+		} else if (sparseAllocations.remove(id)) {
+			return // was allocated above highest ID
+		} else if (id > nextID) {
+			throw IllegalArgumentException("Not tracking ID: $id (highest known ID is ${nextID})")
 		} else if (id in gaps) {
 			throw IllegalArgumentException("ID is already free: $id")
-		}
-
-		if (id + 1 == highestID) {
-			highestID--
+		} else if (id == nextID) {
+			nextID--
 		} else {
-			gaps.add(id)
+			check(gaps.add(id))
 		}
 	}
 
 	fun reset() {
-		highestID = 0
+		nextID = -1
 		gaps.clear()
+		sparseAllocations.clear()
 	}
 }