ConcurrentLong2ObjectMap.kt

package cloud.rsps.util

import java.util.concurrent.locks.StampedLock

/**
 * A concurrent map implementation with `long` keys and object values.
 * * Uses lock striping and open addressing with linear probing.
 * * `null` values are not allowed.
 * * Supports up to 65,536 entries (index `0..65535`).
 *
 * @author Jire
 */
class ConcurrentLong2ObjectMap<V> {
    private val segments: Array<Segment<V?>?>

    init {
        @Suppress("UNCHECKED_CAST")
        segments = Array<Segment<V?>?>(SEGMENT_COUNT) {
            Segment<V?>()
        }
    }

    operator fun get(key: Long): V? {
        require(key != EMPTY_KEY) { "Key cannot be $EMPTY_KEY" }

        val hash = hash(key)
        val segment = segments[hash and SEGMENT_MASK]
            ?: error("Segment not found for key $key with hash $hash")
        return segment.get(key, hash)
    }

    operator fun set(key: Long, value: V): V? {
        require(key != EMPTY_KEY) { "Key cannot be $EMPTY_KEY" }

        val hash = hash(key)
        val segment = segments[hash and SEGMENT_MASK]
            ?: error("Segment not found for key $key with hash $hash")
        return segment.put(key, value, hash)
    }

    private class Segment<V> {
        private val lock = StampedLock()

        private val keys = LongArray(SEGMENT_SIZE) {
            EMPTY_KEY
        }
        private val values = arrayOfNulls<Any>(SEGMENT_SIZE)

        fun get(key: Long, hash: Int): V? {
            var idx = (hash ushr 8) and SEGMENT_MASK

            var stamp = lock.tryOptimisticRead()

            var k = keys[idx]
            var v = values[idx]

            if (!lock.validate(stamp)) {
                stamp = lock.readLock()
                try {
                    k = keys[idx]
                    v = values[idx]
                } finally {
                    lock.unlockRead(stamp)
                }
            }

            if (k == key) {
                @Suppress("UNCHECKED_CAST")
                return v as V?
            }

            repeat(SEGMENT_MASK) {
                idx = (idx + 1) and SEGMENT_MASK

                stamp = lock.tryOptimisticRead()

                k = keys[idx]
                v = values[idx]

                if (!lock.validate(stamp)) {
                    stamp = lock.readLock()
                    try {
                        k = keys[idx]
                        v = values[idx]
                    } finally {
                        lock.unlockRead(stamp)
                    }
                }

                if (k == key) {
                    @Suppress("UNCHECKED_CAST")
                    return v as V?
                }
                if (k == EMPTY_KEY) {
                    return null
                }
            }
            return null
        }

        fun put(key: Long, value: V, hash: Int): V? {
            val idx = (hash ushr 8) and SEGMENT_MASK

            val stamp = lock.writeLock()
            try {
                var emptyIdx = DEFAULT_EMPTY_IDX
                for (i in 0..SEGMENT_MASK) {
                    val currentIdx = (idx + i) and SEGMENT_MASK
                    val k = keys[currentIdx]
                    if (k == key) {
                        @Suppress("UNCHECKED_CAST")
                        val old = values[currentIdx] as V?
                        values[currentIdx] = value
                        return old
                    }
                    if (k == EMPTY_KEY) {
                        emptyIdx = currentIdx
                        break
                    }
                }

                if (emptyIdx == DEFAULT_EMPTY_IDX) {
                    error("Segment full")
                }

                keys[emptyIdx] = key
                values[emptyIdx] = value
                return null
            } finally {
                lock.unlockWrite(stamp)
            }
        }
    }

    private companion object {
        private const val SEGMENT_COUNT = 256
        private const val SEGMENT_SIZE = 256
        private const val SEGMENT_MASK = SEGMENT_COUNT - 1

        private const val DEFAULT_EMPTY_IDX = -1

        private const val EMPTY_KEY = Long.MIN_VALUE

        private fun hash(key: Long): Int {
            var hash = (key xor (key ushr 32)).toInt()
            hash = hash xor (hash ushr 16)
            return hash
        }
    }
}