MapTree.kt

package com.cardlay.nebula.shared.query

class MapTree<K, V : Any> : MutableTree<K, V> {
    private class MutableNode<K, V : Any> : Tree.Node<K, V> {
        override val nodes = mutableMapOf<K, MutableNode<K, V>>()
        override var value: V? = null

        val size: Int
            get() {
                if (value != null) {
                    return 1 + nodes.values.sumOf { it.size }
                }

                return nodes.values.sumOf { it.size }
            }
    }

    private val root = MutableNode<K, V>()

    override val size: Int
        get() = root.size

    override operator fun get(keys: List<K>): V? {
        return getNode(keys)?.value
    }

    override fun containsKeys(keys: List<K>): Boolean {
        return get(keys) != null
    }

    override fun put(keys: List<K>, value: V): V? {
        val node = getNode(keys) ?: makeNode(keys)

        val currentValue = node.value
        node.value = value
        return currentValue
    }

    override operator fun set(keys: List<K>, value: V) {
        put(keys, value)
    }

    override fun remove(keys: List<K>) {
        when (keys.size) {
            0 -> {
                root.nodes.clear()
                root.value = null
            }
            1 -> {
                val key = keys.last()
                root.nodes.remove(key)
            }
            else -> {
                val suffix = keys.last()
                val prefix = keys.dropLast(1)
                val parent = getNode(prefix)
                parent?.nodes?.remove(suffix)
            }
        }
    }

    private fun getNode(keys: List<K>): MutableNode<K, V>? {
        return keys.fold(root as MutableNode<K, V>?) { node, key -> node?.nodes?.get(key) }
    }

    private fun makeNode(keys: List<K>): MutableNode<K, V> {
        return keys.fold(root) { node, key -> node.nodes.getOrPut(key) { MutableNode() } }
    }
}