Make DynamicSynchableGroup implement Int2ReferenceMap

what am i doing

src/main/kotlin/ru/dbotthepony/mc/otm/network/syncher/DynamicSynchableGroup.kt

@@ -1,8 +1,19 @@
 package ru.dbotthepony.mc.otm.network.syncher
 
 import it.unimi.dsi.fastutil.ints.Int2ObjectOpenHashMap
+import it.unimi.dsi.fastutil.ints.Int2ReferenceLinkedOpenHashMap
+import it.unimi.dsi.fastutil.ints.Int2ReferenceMap
+import it.unimi.dsi.fastutil.ints.Int2ReferenceOpenHashMap
 import it.unimi.dsi.fastutil.ints.IntArraySet
+import it.unimi.dsi.fastutil.ints.IntCollection
+import it.unimi.dsi.fastutil.ints.IntIterator
+import it.unimi.dsi.fastutil.ints.IntSet
+import it.unimi.dsi.fastutil.objects.ObjectIterator
+import it.unimi.dsi.fastutil.objects.ObjectSet
+import it.unimi.dsi.fastutil.objects.Reference2ObjectLinkedOpenHashMap
+import it.unimi.dsi.fastutil.objects.ReferenceCollection
 import it.unimi.dsi.fastutil.objects.ReferenceOpenHashSet
+import it.unimi.dsi.fastutil.objects.ReferenceSet
 import net.minecraft.network.RegistryFriendlyByteBuf
 import ru.dbotthepony.kommons.util.KOptional
 import ru.dbotthepony.mc.otm.util.IDAllocator
@@ -39,7 +50,7 @@ class DynamicSynchableGroup<T : Any>(
 	 * first-time networking of [T] to remote
 	 */
 	private val writer: T.(RegistryFriendlyByteBuf) -> Unit = {},
-) : ISynchable, MutableSet<T> {
+) : ISynchable, ReferenceSet<T>, Int2ReferenceMap<T> {
 	private inner class RemoteState(val listener: Runnable) : IRemoteState {
 		private inner class RemoteSlot(val slot: Slot<T>, fromConstructor: Boolean) : Runnable, Closeable {
 			val remoteState = synchable(slot.value).createRemoteState(this)
@@ -188,11 +199,12 @@ class DynamicSynchableGroup<T : Any>(
 		}
 	}
 
-	private data class Slot<T : Any>(val value: T, val id: Int)
+	private class Slot<T : Any>(val value: T, val id: Int)
 
+	private var defaultReturnValue: T? = null
 	private val remoteStates = ArrayList<RemoteState>()
-	private val value2slot = LinkedHashMap<T, Slot<T>>()
+	private val value2slot = Reference2ObjectLinkedOpenHashMap<T, Slot<T>>()
-	private val id2slot = Int2ObjectOpenHashMap<Slot<T>>()
+	private val id2slot = Int2ReferenceLinkedOpenHashMap<Slot<T>>()
 	private val idAllocator = IDAllocator()
 
 	override val hasRemotes: Boolean
@@ -210,7 +222,7 @@ class DynamicSynchableGroup<T : Any>(
 				REMOVE_ENTRY -> {
 					val id = stream.readVarInt()
 					val slot = checkNotNull(id2slot.remove(id)) { "No such slot with ID: $id" }
-					check(value2slot.remove(slot.value) == value2slot)
+					check(value2slot.remove(slot.value) == slot)
 					remoteStates.forEach { it.remove(slot) }
 				}
 
@@ -223,7 +235,7 @@ class DynamicSynchableGroup<T : Any>(
 
 					if (id2slot.containsKey(id)) {
 						val slot = id2slot.remove(id)!!
-						check(value2slot.remove(slot.value) == value2slot)
+						check(value2slot.remove(slot.value) == slot)
 						remoteStates.forEach { it.remove(slot) }
 					}
 
@@ -298,8 +310,8 @@ class DynamicSynchableGroup<T : Any>(
 		return value2slot.isEmpty()
 	}
 
-	override fun iterator(): MutableIterator<T> {
+	override fun iterator(): ObjectIterator<T> {
-		return object : MutableIterator<T> {
+		return object : ObjectIterator<T> {
 			private val parent = value2slot.values.iterator()
 			private var last: KOptional<Slot<T>> = KOptional()
 
@@ -323,11 +335,15 @@ class DynamicSynchableGroup<T : Any>(
 		}
 	}
 
-	override fun remove(element: T): Boolean {
+	private fun removeBySlot(slot: Slot<T>) {
-		val slot = value2slot.remove(element) ?: return false
+		value2slot.remove(slot.value)
-		checkNotNull(id2slot.remove(slot.id))
+		id2slot.remove(slot.id)
 		remoteStates.forEach { it.remove(slot) }
 		idAllocator.release(slot.id)
+	}
+
+	override fun remove(element: T): Boolean {
+		removeBySlot(value2slot.remove(element) ?: return false)
 		return true
 	}
 
@@ -351,6 +367,315 @@ class DynamicSynchableGroup<T : Any>(
 		return any
 	}
 
+	override fun put(key: Int, value: T): T? {
+		val slot = Slot(value, key)
+		val oldSlot = id2slot.put(key, slot)
+
+		if (oldSlot?.value === value)
+			return value
+
+		if (oldSlot != null) {
+			remoteStates.forEach { it.remove(oldSlot) }
+			value2slot.remove(oldSlot.value)
+		} else {
+			idAllocator.allocate(key)
+		}
+
+		value2slot[value] = slot
+		remoteStates.forEach { it.add(slot) }
+
+		return oldSlot?.value
+	}
+
+	override fun get(key: Int): T? {
+		return id2slot.get(key)?.value ?: defaultReturnValue
+	}
+
+	override fun containsKey(key: Int): Boolean {
+		return id2slot.containsKey(key)
+	}
+
+	override fun defaultReturnValue(rv: T?) {
+		defaultReturnValue = rv
+	}
+
+	override fun defaultReturnValue(): T? {
+		return defaultReturnValue
+	}
+
+	override fun putAll(from: Map<out Int, T>) {
+		from.forEach { (k, v) ->
+			put(k, v)
+		}
+	}
+
+	override fun containsValue(value: T): Boolean {
+		return value2slot.containsKey(value)
+	}
+
+	private val entrySet: ObjectSet<Int2ReferenceMap.Entry<T>> by lazy(LazyThreadSafetyMode.PUBLICATION) {
+		object : ObjectSet<Int2ReferenceMap.Entry<T>> {
+			inner class Entry(private val ikey: Int) : Int2ReferenceMap.Entry<T> {
+				override val value: T
+					get() = id2slot[ikey]?.value ?: throw NoSuchElementException()
+
+				override fun setValue(newValue: T): T? {
+					return put(ikey, newValue)
+				}
+
+				override fun getIntKey(): Int {
+					return ikey
+				}
+
+				override fun equals(other: Any?): Boolean {
+					return this === other || other is Int2ReferenceMap.Entry<*> && other.intKey == ikey && other.value === id2slot[ikey]?.value
+				}
+
+				override fun hashCode(): Int {
+					return ikey * 31 + (id2slot[ikey] as Slot<T>?).hashCode()
+				}
+
+				override fun toString(): String {
+					return "Map.Entry[$ikey, ${id2slot[ikey]?.value}]"
+				}
+			}
+
+			override val size: Int
+				get() = id2slot.size
+
+			override fun add(element: Int2ReferenceMap.Entry<T>): Boolean {
+				return put(element.intKey, element.value) != null
+			}
+
+			override fun addAll(elements: Collection<Int2ReferenceMap.Entry<T>>): Boolean {
+				var any = false
+				elements.forEach { any = add(it) || any }
+				return any
+			}
+
+			override fun clear() {
+				this@DynamicSynchableGroup.clear()
+			}
+
+			override fun iterator(): ObjectIterator<Int2ReferenceMap.Entry<T>> {
+				return object : ObjectIterator<Int2ReferenceMap.Entry<T>> {
+					private val parent = id2slot.int2ReferenceEntrySet().iterator()
+					private var last: Slot<T>? = null
+
+					override fun hasNext(): Boolean {
+						return parent.hasNext()
+					}
+
+					override fun remove() {
+						if (last == null)
+							throw NoSuchElementException()
+
+						parent.remove()
+						removeBySlot(last!!)
+						last = null
+					}
+
+					override fun next(): Int2ReferenceMap.Entry<T> {
+						val value = parent.next()
+						last = value.value
+						return Entry(value.intKey)
+					}
+				}
+			}
+
+			override fun remove(element: Int2ReferenceMap.Entry<T>): Boolean {
+				val slot = id2slot.get(element.intKey)
+
+				if (slot?.value === element.value) {
+					this@DynamicSynchableGroup.removeBySlot(slot)
+					return true
+				}
+
+				return false
+			}
+
+			override fun removeAll(elements: Collection<Int2ReferenceMap.Entry<T>>): Boolean {
+				var any = false
+				elements.forEach { any = remove(it) || any }
+				return any
+			}
+
+			override fun retainAll(elements: Collection<Int2ReferenceMap.Entry<T>>): Boolean {
+				if (elements.isEmpty()) {
+					val isNotEmpty = id2slot.isNotEmpty()
+					this@DynamicSynchableGroup.clear()
+					return isNotEmpty
+				}
+
+				var any = false
+				val itr = id2slot.int2ReferenceEntrySet().iterator()
+
+				for (entry in itr) {
+					if (elements.none { it.intKey == entry.intKey && it.value === entry.value.value }) {
+						any = true
+						itr.remove()
+						removeBySlot(entry.value)
+					}
+				}
+
+				return any
+			}
+
+			override fun contains(element: Int2ReferenceMap.Entry<T>): Boolean {
+				return id2slot.get(element.intKey)?.value === element.value
+			}
+
+			override fun containsAll(elements: Collection<Int2ReferenceMap.Entry<T>>): Boolean {
+				return elements.all { contains(it) }
+			}
+
+			override fun isEmpty(): Boolean {
+				return id2slot.isEmpty()
+			}
+		}
+	}
+
+	override fun int2ReferenceEntrySet(): ObjectSet<Int2ReferenceMap.Entry<T>> {
+		return entrySet
+	}
+
+	override val keys: IntSet by lazy(LazyThreadSafetyMode.PUBLICATION) {
+		object : IntSet {
+			override fun add(key: Int): Boolean {
+				throw UnsupportedOperationException()
+			}
+
+			override fun addAll(c: IntCollection): Boolean {
+				throw UnsupportedOperationException()
+			}
+
+			override fun addAll(elements: Collection<Int>): Boolean {
+				throw UnsupportedOperationException()
+			}
+
+			override fun clear() {
+				this@DynamicSynchableGroup.clear()
+			}
+
+			override fun iterator(): IntIterator {
+				return object : IntIterator {
+					private val parent = id2slot.keys.iterator()
+					private var last = -1
+					private var hasLast = false
+
+					override fun hasNext(): Boolean {
+						return parent.hasNext()
+					}
+
+					override fun remove() {
+						if (!hasLast)
+							throw NoSuchElementException()
+
+						hasLast = false
+						val slot = id2slot.get(last)
+						parent.remove()
+						removeBySlot(slot)
+					}
+
+					override fun nextInt(): Int {
+						last = parent.nextInt()
+						hasLast = true
+						return last
+					}
+				}
+			}
+
+			override fun remove(k: Int): Boolean {
+				removeBySlot(id2slot.get(k) ?: return false)
+				return true
+			}
+
+			override fun removeAll(c: IntCollection): Boolean {
+				val itr = c.iterator()
+				var any = false
+
+				while (itr.hasNext()) {
+					any = remove(itr.nextInt()) || any
+				}
+
+				return any
+			}
+
+			override fun removeAll(elements: Collection<Int>): Boolean {
+				val itr = elements.iterator()
+				var any = false
+
+				while (itr.hasNext()) {
+					any = remove(itr.next()) || any
+				}
+
+				return any
+			}
+
+			override fun retainAll(c: IntCollection): Boolean {
+				val itr = id2slot.values.iterator()
+				var any = false
+
+				for (slot in itr) {
+					if (!c.contains(slot.id)) {
+						itr.remove()
+						check(value2slot.remove(slot.value) == slot)
+						remoteStates.forEach { it.remove(slot) }
+						any = true
+					}
+				}
+
+				return any
+			}
+
+			override fun retainAll(elements: Collection<Int>): Boolean {
+				val itr = id2slot.values.iterator()
+				var any = false
+
+				for (slot in itr) {
+					if (!elements.contains(slot.id)) {
+						itr.remove()
+						check(value2slot.remove(slot.value) == slot)
+						remoteStates.forEach { it.remove(slot) }
+						any = true
+					}
+				}
+
+				return any
+			}
+
+			override fun contains(key: Int): Boolean {
+				return id2slot.containsKey(key)
+			}
+
+			override fun containsAll(c: IntCollection): Boolean {
+				return id2slot.keys.containsAll(c)
+			}
+
+			override fun containsAll(elements: Collection<Int>): Boolean {
+				return id2slot.keys.containsAll(elements)
+			}
+
+			override fun isEmpty(): Boolean {
+				return id2slot.isEmpty()
+			}
+
+			override fun toArray(a: IntArray?): IntArray {
+				return id2slot.keys.toArray(a)
+			}
+
+			override fun toIntArray(): IntArray {
+				return id2slot.keys.toIntArray()
+			}
+
+			override val size: Int
+				get() = id2slot.size
+		}
+	}
+
+	override val values: ReferenceCollection<T>
+		get() = this
+
 	companion object {
 		private const val END = 0
 		private const val ADD_ENTRY = 1