index/src/main/kotlin/moe/odango/index/utils/MergeMap.kt
2020-06-16 21:03:09 +02:00

50 lines
1.5 KiB
Kotlin

package moe.odango.index.utils
class MergeMap<T> {
private var bucketIndex = 0L
private val buckets: MutableMap<Long, MutableSet<T>> = mutableMapOf()
private val bucketIndexByItem: MutableMap<T, Long> = mutableMapOf()
fun add(a: T, b: T) {
val bucketA = bucketIndexByItem[a]
val bucketB = bucketIndexByItem[b]
if (bucketA == null && bucketB == null) {
bucketIndexByItem[a] = bucketIndex
bucketIndexByItem[b] = bucketIndex
buckets[bucketIndex] = mutableSetOf(a, b)
bucketIndex++
return
}
if (bucketB == bucketA) {
return
}
if (bucketA == null && bucketB != null) {
bucketIndexByItem[a] = bucketB
buckets.getOrPut(bucketB, ::mutableSetOf).add(a)
return
}
if (bucketB == null && bucketA != null) {
bucketIndexByItem[b] = bucketA
buckets.getOrPut(bucketA, ::mutableSetOf).add(b)
return
}
// Always false :)
if (bucketA == null || bucketB == null) return
val bucket = buckets.getOrPut(bucketB, ::mutableSetOf)
buckets.getOrPut(bucketA, ::mutableSetOf).addAll(bucket)
for (item in bucket) {
bucketIndexByItem[item] = bucketA
}
}
operator fun get(item: T): Set<T>? {
return buckets[bucketIndexByItem[item] ?: return null]
}
operator fun iterator(): Iterator<Set<T>> = buckets.values.iterator()
}