WeakReferenceMap.kt

package adb

import kotlinx.coroutines.CoroutineName
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.DelicateCoroutinesApi
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.isActive
import kotlinx.coroutines.launch
import java.lang.ref.Reference
import java.lang.ref.ReferenceQueue
import java.lang.ref.WeakReference
import java.util.Collections

/**
 * A WeakHashMap using reference equality.
 */
@DelicateCoroutinesApi
class WeakReferenceMap<K, V>(cleanupScope: CoroutineScope = CoroutineScope(Dispatchers.IO)) {
    private val map: MutableMap<Reference<out K>, V> = Collections.synchronizedMap(HashMap())

    private val queue = ReferenceQueue<K>()

    init {
        cleanupScope.launch(CoroutineName("WeakReferenceMap.Cleanup")) {
            while (isActive) {
                map.remove(queue.remove())
            }
        }
    }

    operator fun get(key: K): V? = map[IdentityRef(key)]

    operator fun set(key: K, value: V) {
        map[IdentityRef(key, queue)] = value
    }

    fun remove(key: K): V? = map.remove(IdentityRef(key))

    private class IdentityRef<T> : WeakReference<T> {
        private val hashCode: Int

        constructor(value: T) : super(value) {
            hashCode = System.identityHashCode(value)
        }

        constructor(value: T, queue: ReferenceQueue<T>) : super(value, queue) {
            hashCode = System.identityHashCode(value)
        }

        override fun hashCode(): Int = hashCode
        override fun equals(other: Any?) = other == get()
    }
}