index/src/main/kotlin/moe/odango/index/utils/MergeMap.kt

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()
}