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

Go to download

Show more of this group Show more artifacts with this name
Show all versions of kotlin-compiler-embeddable Show documentation

the Kotlin compiler embeddable

There is a newer version: 2.1.0-RC

Show newest version

/*
 * 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

import kotlinx.collections.immutable.*

class PersistentMultimap private constructor(private val map: PersistentMap>) {

    constructor() : this(persistentMapOf())

    fun put(key: K, value: V): PersistentMultimap {
        val collection = map[key] ?: persistentListOf()
        val newSet = collection.add(value)
        if (newSet === collection) return this
        val newMap = map.put(key, newSet)
        return PersistentMultimap(newMap)
    }

    fun remove(key: K, value: V): PersistentMultimap {
        val list = map.get(key) ?: return this
        val newSet = list.remove(value)
        if (list === newSet) return this
        val newMap = if (newSet.isEmpty()) {
            map.remove(key)
        } else {
            map.put(key, newSet)
        }
        return PersistentMultimap(newMap)
    }

    operator fun get(key: K): List {
        return map[key] ?: emptyList()
    }

    val keys: ImmutableSet get() = map.keys
}

class PersistentSetMultimap private constructor(private val map: PersistentMap>) {

    constructor() : this(persistentMapOf())

    fun put(key: K, value: V): PersistentSetMultimap {
        val set = map[key] ?: persistentSetOf()
        val newSet = set.add(value)
        if (newSet === set) return this
        val newMap = map.put(key, newSet)
        return PersistentSetMultimap(newMap)
    }

    fun remove(key: K, value: V): PersistentSetMultimap {
        val set = map.get(key) ?: return this
        val newSet = set.remove(value)
        if (set === newSet) return this
        val newMap = if (newSet.isEmpty()) {
            map.remove(key)
        } else {
            map.put(key, newSet)
        }
        return PersistentSetMultimap(newMap)
    }

    operator fun get(key: K): Set {
        return map[key] ?: emptySet()
    }
}