EnumMap.kt

open class EnumMap<E : Enum<E>, V : Any>(
    meta: Class<E>
) : MutableMap<E, V> {
    
    private val universe = meta.enumConstants
    private val storage = Array.newInstance(java.lang.Object::class.java, universe.size) as kotlin.Array<V?>

    override fun get(key: E): V? {
        return storage[key.ordinal]
    }

    override fun containsKey(key: E): Boolean {
        return storage[key.ordinal] != null
    }

    override fun containsValue(value: V): Boolean {
        return storage.contains(value)
    }

    override val entries: MutableSet<MutableMap.MutableEntry<E, V>>
        get() = object : MutableSet<MutableMap.MutableEntry<E, V>> {
            override fun add(element: MutableMap.MutableEntry<E, V>): Boolean {
                this@EnumMap[element.key] = element.value
                return true
            }

            override val size: Int get() = [email protected]

            override fun clear() = [email protected]()

            override fun isEmpty(): Boolean = [email protected]()

            override fun iterator(): MutableIterator<MutableMap.MutableEntry<E, V>> {
                return EntryItr()
            }

            override fun retainAll(elements: Collection<MutableMap.MutableEntry<E, V>>): Boolean {
                val retained = BitSet(universe.size)
                elements.forEach {
                    if (it.value == storage[it.key.ordinal])
                        retained.set(it.key.ordinal)
                }
                var deletedAny = false
                for (index in universe.indices) {
                    if (!retained.get(index)) {
                        [email protected][index] = null
                        deletedAny = true
                    }
                }
                return deletedAny
            }

            override fun removeAll(elements: Collection<MutableMap.MutableEntry<E, V>>): Boolean {
                var removedAny = false
                for (element in elements) {
                    removedAny = remove(element) || removedAny
                }
                return removedAny
            }

            override fun remove(element: MutableMap.MutableEntry<E, V>): Boolean {
                val original = storage[element.key.ordinal]
                if (original != null && element.value == original) {
                    storage[element.key.ordinal] = null
                    return true
                }
                return false
            }

            override fun containsAll(elements: Collection<MutableMap.MutableEntry<E, V>>): Boolean {
                for (element in elements) {
                    if (!contains(element))
                        return false
                }
                return true
            }

            override fun contains(element: MutableMap.MutableEntry<E, V>): Boolean {
                return storage[element.key.ordinal] == element.value
            }

            override fun addAll(elements: Collection<MutableMap.MutableEntry<E, V>>): Boolean {
                var anyChanged = false
                for (element in elements) {
                    if (storage[element.key.ordinal] != element.value) {
                        storage[element.key.ordinal] = element.value
                        anyChanged = true
                    }
                }
                return anyChanged
            }
        }

    private abstract inner class Itr<VV> : MutableIterator<VV> {
        var index = -1
        var foundAny = true
        override fun hasNext(): Boolean {
            if (foundAny) return true
            while (++index < universe.size) {
                if (storage[index] != null) {
                    foundAny = false
                    return true
                }
            }
            return false
        }

        override fun remove() {
            storage[index] = null
            foundAny = false
        }
    }

    private inner class KeyItr : Itr<E>() {
        override fun next(): E {
            if (!foundAny) throw NoSuchElementException()
            foundAny = false
            return universe[index]
        }
    }

    private inner class EntryItr : Itr<MutableMap.MutableEntry<E, V>>() {
        override fun next(): MutableMap.MutableEntry<E, V> {
            if (!foundAny) throw NoSuchElementException()
            val index = this.index
            return object : MutableMap.MutableEntry<E, V> {
                override val key: E
                    get() = universe[index]
                override val value: V
                    get() = storage[index] as V

                override fun setValue(newValue: V): V {
                    val oldValue = storage[index] as V
                    storage[index] = newValue
                    return oldValue
                }
            }
        }
    }

    private inner class ValueItr : Itr<V>() {
        override fun next(): V {
            if (!foundAny) throw NoSuchElementException()
            foundAny = false
            return storage[index] as V
        }
    }

    override val keys: MutableSet<E>
        get() = object : MutableSet<E> {
            override val size: Int
                get() = [email protected]

            override fun add(element: E): Boolean {
                throw NotImplementedError()
            }

            override fun addAll(elements: Collection<E>): Boolean {
                throw NotImplementedError()
            }

            override fun clear() {
                [email protected]()
            }

            override fun contains(element: E): Boolean {
                return storage[element.ordinal] != null
            }

            override fun containsAll(elements: Collection<E>): Boolean {
                return elements.all { contains(it) }
            }

            override fun isEmpty(): Boolean {
                return [email protected]()
            }

            override fun iterator(): MutableIterator<E> {
                return KeyItr()
            }

            override fun remove(element: E): Boolean {
                val removedAny = storage[element.ordinal] != null
                storage[element.ordinal] = null
                return removedAny
            }

            override fun removeAll(elements: Collection<E>): Boolean {
                var removedAny = false
                for (element in elements) {
                    removedAny = removedAny || storage[element.ordinal] != null
                    storage[element.ordinal] = null
                }
                return removedAny
            }

            override fun retainAll(elements: Collection<E>): Boolean {
                val retained = BitSet(universe.size)
                for (element in elements) {
                    retained.set(element.ordinal)
                }
                var deletedAny = false
                var index = 0
                while (true) {
                    index = retained.nextClearBit(index)
                    if (index == -1) break
                    deletedAny = deletedAny || storage[index] != null
                    storage[index] = null
                }
                return deletedAny
            }
        }

    override val size: Int
        get() = storage.count { it != null }
    override val values: MutableCollection<V>
        get() = object : MutableCollection<V> {
            override val size: Int
                get() = [email protected]

            override fun clear() {
                [email protected]()
            }

            override fun isEmpty(): Boolean {
                return [email protected]()
            }

            override fun iterator(): MutableIterator<V> {
                return ValueItr()
            }

            override fun retainAll(elements: Collection<V>): Boolean {
                var deletedAny = false
                for (index in storage.indices) {
                    if (storage[index] !in elements) {
                        storage[index] = null
                        deletedAny = true
                    }
                }
                return deletedAny
            }

            override fun removeAll(elements: Collection<V>): Boolean {
                var deletedAny = false
                for (index in storage.indices) {
                    if (storage[index] in elements) {
                        storage[index] = null
                        deletedAny = true
                    }
                }
                return deletedAny

            }

            override fun remove(element: V): Boolean {
                for (index in storage.indices) {
                    if (storage[index] == element) {
                        storage[index] = null
                        return true
                    }
                }
                return false
            }

            override fun containsAll(elements: Collection<V>): Boolean {
                val allElements = this.toSet()
                return elements.any { it !in allElements }
            }

            override fun contains(element: V): Boolean {
                return storage.contains(element)
            }

            override fun addAll(elements: Collection<V>): Boolean {
                throw NotImplementedError()
            }

            override fun add(element: V): Boolean {
                throw NotImplementedError()
            }
        }

    override fun clear() {
        Arrays.fill(storage, null)
    }

    override fun isEmpty(): Boolean {
        return storage.none { it != null }
    }

    override fun remove(key: E): V? {
        val removed = storage[key.ordinal]
        storage[key.ordinal] = null
        return removed
    }

    override fun putAll(from: Map<out E, V>) {
        for (entry in from) {
            put(entry.key, entry.value)
        }
    }

    override fun put(key: E, value: V): V? {
        val old = storage[key.ordinal]
        storage[key.ordinal] = value
        return old
    }
}