org.jetbrains.kotlin.fir.util.Multimap.kt Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of kotlin-compiler Show documentation
the Kotlin compiler
There is a newer version: 2.1.0-RC
/*
 * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
 * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
 */

package org.jetbrains.kotlin.fir.util

interface Multimap> {
    operator fun get(key: K): C
    operator fun contains(key: K): Boolean
    val keys: Set
    val values: Collection
}

interface MutableMultimap> : Multimap {
    fun put(key: K, value: V)
    fun putAll(key: K, values: Collection) {
        values.forEach { put(key, it) }
    }

    fun remove(key: K, value: V)
    fun removeKey(key: K): C

    fun clear()
}

abstract class BaseMultimap, MC : MutableCollection> : MutableMultimap {
    private val map: MutableMap = mutableMapOf()
    protected abstract fun createContainer(): MC
    protected abstract fun createEmptyContainer(): C

    override fun get(key: K): C {
        @Suppress("UNCHECKED_CAST")
        return map[key] as C? ?: createEmptyContainer()
    }

    override operator fun contains(key: K): Boolean {
        return key in map
    }

    override val keys: Set
        get() = map.keys

    override val values: Collection
        get() = object : AbstractCollection() {
            override val size: Int
                get() = map.values.sumOf { it.size }

            override fun iterator(): Iterator {
                return ChainedIterator(map.values.map { it.iterator() })
            }
        }

    override fun put(key: K, value: V) {
        val container = map.getOrPut(key) { createContainer() }
        container.add(value)
    }

    override fun remove(key: K, value: V) {
        val collection = map[key] ?: return
        collection.remove(value)
        if (collection.isEmpty()) {
            map.remove(key)
        }
    }

    override fun removeKey(key: K): C {
        @Suppress("UNCHECKED_CAST")
        return map.remove(key) as C? ?: createEmptyContainer()
    }

    override fun clear() {
        map.clear()
    }
}

class SetMultimap : BaseMultimap, MutableSet>() {
    override fun createContainer(): MutableSet {
        return mutableSetOf()
    }

    override fun createEmptyContainer(): Set {
        return emptySet()
    }
}

class ListMultimap : BaseMultimap, MutableList>() {
    override fun createContainer(): MutableList {
        return mutableListOf()
    }

    override fun createEmptyContainer(): List {
        return emptyList()
    }
}

fun  setMultimapOf(): SetMultimap = SetMultimap()
fun  listMultimapOf(): ListMultimap = ListMultimap()