Kotlin WeakMap that uses identity comparison rather than equality comparison

WeakIdentityMap.kt

package org.mypackage.util

import java.lang.ref.Reference
import java.lang.ref.ReferenceQueue
import java.lang.ref.WeakReference

/**
 * WeakMap that uses identity comparison rather than equality comparison ([java.util.WeakHashMap]).
 * Useful for instance mappings.
 *
 * Notice that keys and entries getters are not implemented.
 */
class WeakIdentityMap<K, V> : MutableMap<K, V> {

    private val map = mutableMapOf<IdentityWeakReference<K>, V>()

    private val queue = ReferenceQueue<Any?>()

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

    override fun containsKey(key: K): Boolean {
        reap()
        return map.containsKey(IdentityWeakReference(key, null))
    }

    override fun containsValue(value: V): Boolean {
        reap()
        return map.containsValue(value)
    }

    override operator fun get(key: K): V? {
        reap()
        return map[IdentityWeakReference(key, null)]
    }

    override fun put(key: K, value: V): V? {
        reap()
        return map.put(IdentityWeakReference(key, queue), value)
    }

    override fun putAll(from: Map<out K, V>) {
        reap()
        from.forEach { (key, value) -> map[IdentityWeakReference(key, queue)] = value }
    }

    override fun isEmpty(): Boolean {
        return size == 0
    }

    override fun remove(key: K): V? {
        reap()
        return map.remove(IdentityWeakReference(key, null))
    }

    override val size: Int
        get() = doGetSize()

    private fun doGetSize(): Int {
        return if (map.isEmpty())
            0
        else {
            reap()
            map.size
        }
    }

    override val values: MutableCollection<V>
        get() = doGetValues()

    private fun doGetValues(): MutableCollection<V> {
        reap()
        return map.values
    }

    // Too much work to implement.
    override val keys: MutableSet<K>
        get() = throw NotImplementedError()

    // Too much work to implement.
    override val entries: MutableSet<MutableMap.MutableEntry<K, V>>
        get() = throw NotImplementedError()

    private fun reap() {
        var zombie: Reference<*>?
        while (queue.poll().also { zombie = it } != null) {
            map.remove(zombie)
        }
    }

    private class IdentityWeakReference<K>(obj: K?, queue: ReferenceQueue<in K?>?) : WeakReference<K>(obj, queue) {

        private val hash: Int = System.identityHashCode(obj)

        override fun hashCode(): Int {
            return hash
        }

        override fun equals(other: Any?): Boolean {
            if (this === other) {
                return true
            }
            if (other is IdentityWeakReference<*>) {
                if (this.get() === other.get()) {
                    return true
                }
            }
            return false
        }
    }

}