WeakIdentityHashMap (not tested, HashMap-backed)

WeakIdentityHashMap.kt

private class WeakIdentityHashMap<K : Any, V> : MutableMap<K, V> {
    private val backed = HashMap<Key<K>, V>()
    private val queue = ReferenceQueue<K>()

    private fun clean() {
        while (true) {
            val v = queue.poll() ?: return
            backed.remove(v)
        }
    }

    private class Key<K : Any>(value: K, queue: ReferenceQueue<K>?) : WeakReference<K>(value, queue) {
        private val hash = System.identityHashCode(value)

        override fun equals(other: Any?): Boolean {
            if (this === other) return true
            if (javaClass != other?.javaClass) return false

            other as Key<*>

            if (hash != other.hash) return false

            val otherValue = other.get()
            val thisValue = get()

            return otherValue != null && otherValue === thisValue
        }

        override fun hashCode(): Int {
            return hash
        }
    }

    override val size: Int get() {
        clean()
        return backed.size
    }

    override fun containsKey(key: K): Boolean {
        clean()
        return backed.containsKey(Key(key, null))
    }

    override fun containsValue(value: V): Boolean {
        clean()
        return backed.containsValue(value)
    }

    override fun get(key: K): V? {
        clean()
        return backed[Key(key, null)]
    }

    override fun isEmpty(): Boolean {
        clean()
        return backed.isEmpty()
    }

    override val entries: MutableSet<MutableMap.MutableEntry<K, V>> = EntrySet(this)
    override val keys: MutableSet<K> = KeySet(this)
    override val values: MutableCollection<V>
        get() = backed.values

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

    override fun put(key: K, value: V): V? {
        return backed.put(Key(key, queue), value)
    }

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

    override fun remove(key: K): V? {
        return backed.remove(Key(key, null))
    }

    private class KeySet<K : Any>(val map: WeakIdentityHashMap<K, *>) : AbstractMutableSet<K>() {
        val backedSet = map.backed.keys

        override fun add(element: K): Boolean {
            throw UnsupportedOperationException("add")
        }

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

        override fun iterator() = IteratorImpl(backedSet.iterator())

        override fun remove(element: K): Boolean {
            return backedSet.remove(Key(element, null))
        }

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

        override fun contains(element: K): Boolean = map.containsKey(element)

        override fun isEmpty(): Boolean = map.isEmpty()

        private class IteratorImpl<K : Any>(val backed: MutableIterator<Key<K>>) : MutableIterator<K> {
            var cur: K? = null

            fun computeNext() {
                if (cur != null) return
                while (backed.hasNext()) {
                    cur = backed.next().get()
                    if (cur != null){
                        return
                    }
                }
            }

            override fun remove() {
                backed.remove()
            }

            override fun hasNext(): Boolean {
                computeNext()
                return cur != null
            }

            override fun next(): K {
                computeNext()
                return cur ?: throw NoSuchElementException()
            }
        }
    }

    private class EntrySet<K : Any, V>(val map: WeakIdentityHashMap<K, V>)
            : AbstractMutableSet<MutableMap.MutableEntry<K, V>>() {
        val backedSet = map.backed.entries

        override fun add(element: MutableMap.MutableEntry<K, V>): Boolean {
            throw UnsupportedOperationException("add")
        }

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

        override fun iterator(): MutableIterator<MutableMap.MutableEntry<K, V>> = IteratorImpl(backedSet.iterator())

        override fun remove(element: MutableMap.MutableEntry<K, V>): Boolean {
            return map.remove(element.key, element.value)
        }

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

        override fun contains(element: MutableMap.MutableEntry<K, V>): Boolean {
            return backedSet.contains(AbstractMap.SimpleEntry(Key(element.key, null), element.value))
        }

        override fun isEmpty(): Boolean = map.isEmpty()

        private class IteratorImpl<K: Any, V>(val backed: MutableIterator<MutableMap.MutableEntry<Key<K>, V>>)
            : MutableIterator<MutableMap.MutableEntry<K, V>> {
            var cur: EntryImpl<K, V>? = null

            fun computeNext() {
                if (cur != null) return
                while (backed.hasNext()) {
                    val backed = backed.next()
                    val key = backed.key.get()
                    if (key != null){
                        cur = EntryImpl(key, backed)
                        return
                    }
                }
            }

            override fun remove() {
                backed.remove()
            }

            override fun hasNext(): Boolean {
                computeNext()
                return cur != null
            }

            override fun next(): EntryImpl<K, V> {
                computeNext()
                return cur ?: throw NoSuchElementException()
            }

            private class EntryImpl<K: Any, V>(override val key: K, val backed: MutableMap.MutableEntry<*, V>)
                : MutableMap.MutableEntry<K, V> {

                override val value: V get() = backed.value

                override fun setValue(newValue: V): V {
                    return backed.setValue(newValue)
                }
            }
        }
    }
}