Concurrent (almost) constant O(1) HashMap for Go

themap.go

package themap

import (
	"runtime"
	"sync"
	"sync/atomic"
)

type Hashable interface {
	comparable
	Hash() uint64
}

type smallMapEntry[K Hashable, V any] struct {
	metadata uint64
	key      K
	value    V
}

type optLock uint32

func (lock *optLock) RLock() (uint32, bool) {
	val := atomic.LoadUint32((*uint32)(lock))
	notLocked := val&optLockBit == 0
	return val, notLocked
}

func (lock *optLock) RVerify(expected uint32) bool {
	return atomic.LoadUint32((*uint32)(lock)) == expected
}

const (
	optLockBit     = 1
	optLockMask    = ^uint32(1)
	optLockBackoff = 1
)

func (lock *optLock) Lock() {
	backoff := 1
	for {
		if lock.TryLock() {
			return
		}
		for i := 0; i < backoff; i++ {
			runtime.Gosched()
		}
		backoff <<= 1
		if backoff > optLockBackoff {
			backoff = optLockBackoff
		}
	}
}

func (lock *optLock) TryLock() bool {
	old := atomic.LoadUint32((*uint32)(lock)) & optLockMask
	new := old | optLockBit
	return atomic.CompareAndSwapUint32((*uint32)(lock), old, new)
}

func (lock *optLock) Unlock() {
	old := atomic.LoadUint32((*uint32)(lock))
	new := old & optLockMask
	check := old & optLockBit
	if check != optLockBit {
		panic("OptLock.Unlock() called without OptLock.Lock()")
	}
	atomic.StoreUint32((*uint32)(lock), new+2)
}

type smallMap[K Hashable, V any] struct {
	optLock optLock
	bits    uint8
	size    uint16
	// metadata [smallMapBucketCount]uint8
	entries [smallMapBucketCount]smallMapEntry[K, V]
}

const (
	smallMapBucketBits           = 5
	smallMapBucketCount          = 1 << smallMapBucketBits
	smallMapBucketMask           = smallMapBucketCount - 1
	smallMapMaxIter              = 8
	smallMapDistanceShift        = 59
	smallMapPresentFlag   uint64 = 1 << 63
	smallMapDistance             = 0b1111 << smallMapDistanceShift
	smallMapFingerprint          = ^(smallMapPresentFlag | smallMapDistance)
)

func (smallMap *smallMap[K, V]) setAndMeta(key K) (uint64, uint64) {
	hash := key.Hash() >> smallMap.bits
	set := (hash & smallMapBucketMask)
	compbyte := (hash >> smallMapBucketBits) & smallMapFingerprint
	metadatabyte := compbyte | smallMapPresentFlag
	return set, metadatabyte
}

func (smallMap *smallMap[K, V]) get(k K) (V, bool) {
	var v V
	set, _ := smallMap.setAndMeta(k)
	for i := uint64(0); i < smallMapMaxIter; i++ {
		idx := (set + i) & smallMapBucketMask
		if smallMap.entries[idx].metadata&smallMapPresentFlag == 0 {
			break
		}
		if smallMap.entries[idx].key == k {
			return smallMap.entries[idx].value, true
		}
	}
	return v, false
}

func distanceOf(meta uint64) uint64 {
	return (meta & smallMapDistance) >> smallMapDistanceShift
}

type smallMapInsertStatus = uint8

const (
	smallMapInsertSplit smallMapInsertStatus = iota
	smallMapInsertShouldSplit
	smallMapInsertVOK
	smallMapInsertVNil
)

func (smallMap *smallMap[K, V]) insert(k K, v V) (V, smallMapInsertStatus) {
	var oldValue V
	if smallMap.size == smallMapBucketCount {
		return v, smallMapInsertSplit
	}
	smallMap.size++
	set, newMeta := smallMap.setAndMeta(k)
	distance := uint64(0)
	for {
		idx := (set + distance) & smallMapBucketMask
		slotMeta := smallMap.entries[idx].metadata
		if (slotMeta&smallMapPresentFlag != 0 &&
			slotMeta&smallMapFingerprint == newMeta&smallMapFingerprint &&
			smallMap.entries[idx].key == k) || (slotMeta&smallMapPresentFlag == 0) {

			oldValue = smallMap.entries[idx].value
			newMeta = (newMeta &^ smallMapDistance) | (distance << smallMapDistanceShift)
			smallMap.entries[idx] = smallMapEntry[K, V]{
				metadata: newMeta,
				key:      k,
				value:    v,
			}
			if slotMeta&smallMapPresentFlag != 0 {
				smallMap.size--
				return oldValue, smallMapInsertVOK
			}
			return oldValue, smallMapInsertVNil
		}

		if distanceOf(slotMeta) < distance {
			prev := smallMap.entries[idx]
			newMeta = (newMeta &^ smallMapDistance) | (distance << smallMapDistanceShift)
			smallMap.entries[idx] = smallMapEntry[K, V]{
				metadata: newMeta,
				key:      k,
				value:    v,
			}
			substitute := (idx + 1) & smallMapBucketMask
			returnCode := smallMapInsertVNil
			iter := 0
			for {
				meta := prev.metadata
				if meta&smallMapPresentFlag == 0 {
					break
				}
				tmp := smallMap.entries[substitute]
				distance := distanceOf(meta)
				if returnCode != smallMapInsertShouldSplit && (distance+1 >= smallMapMaxIter || iter >= smallMapBucketCount) {
					returnCode = smallMapInsertShouldSplit
				}
				prev.metadata = (meta &^ smallMapDistance) | ((distance + 1) << smallMapDistanceShift)
				smallMap.entries[substitute] = prev
				prev = tmp
				substitute = (substitute + 1) & smallMapBucketMask
				iter++
			}
			return oldValue, returnCode
		}
		distance++
	}
}

func (smallMap *smallMap[K, V]) delete(k K) (V, bool) {
	var oldValue V
	return oldValue, false
}

func (smallMap *smallMap[K, V]) reset() {
	smallMap.size = 0
	smallMap.entries = [smallMapBucketCount]smallMapEntry[K, V]{}
}

func (theMap *TheMap[K, V]) splitMap(sm *smallMap[K, V], bmask uint64, pool *sync.Pool) (*smallMap[K, V], bool, *smallMap[K, V], bool) {
	amap := pool.Get().(*smallMap[K, V])
	amap.optLock.Lock()
	amap.bits = sm.bits + 1

	bmap := pool.Get().(*smallMap[K, V])
	bmap.optLock.Lock()
	bmap.bits = sm.bits + 1

	splita, splitb := false, false
	for i := 0; i < smallMapBucketCount; i++ {
		if sm.entries[i].metadata&smallMapPresentFlag == 0 {
			continue
		}
		h := sm.entries[i].key.Hash() & bmask
		k, v := sm.entries[i].key, sm.entries[i].value
		if h != 0 {
			_, code := bmap.insert(k, v)
			if code == smallMapInsertShouldSplit {
				splitb = true
			}
		} else {
			_, code := amap.insert(k, v)
			if code == smallMapInsertShouldSplit {
				splita = true
			}
		}
	}
	return amap, splita, bmap, splitb
}

type TheMap[K Hashable, V any] struct {
	dictlock     optLock
	bits         uint8
	smaps        uint64
	dictionary   []atomic.Pointer[smallMap[K, V]]
	smallMapPool sync.Pool
}

func (theMap *TheMap[K, V]) Insert(key K, value V) (V, bool) {
	hash := key.Hash()
	for {
		idx := hash & ((uint64(1) << theMap.bits) - 1)
		ver, ok := theMap.dictlock.RLock()
		if !ok {
			runtime.Gosched()
			continue
		}
		mapRef := theMap.dictionary[idx].Load()
		mapRef.optLock.Lock()
		if !theMap.dictlock.RVerify(ver) || theMap.dictionary[idx].Load() != mapRef {
			mapRef.optLock.Unlock()
			continue
		}
		v, status := mapRef.insert(key, value)
		switch status {
		case smallMapInsertVOK:
			mapRef.optLock.Unlock()
			return v, true
		case smallMapInsertVNil:
			mapRef.optLock.Unlock()
			return v, false
		case smallMapInsertSplit:
			theMap.splitInsert(mapRef, idx, key, value)
		case smallMapInsertShouldSplit:
			theMap.split(mapRef, idx)
		}
		return v, false
	}
}

func (theMap *TheMap[K, V]) Get(key K) (V, bool) {
	hash := key.Hash()
	for {
		idx := hash & ((uint64(1) << theMap.bits) - 1)
		ver, ok := theMap.dictlock.RLock()
		if !ok {
			runtime.Gosched()
			continue
		}
		mapRef := theMap.dictionary[idx].Load()
		mapVer, ok := mapRef.optLock.RLock()
		if !theMap.dictlock.RVerify(ver) || !ok || theMap.dictionary[idx].Load() != mapRef {
			runtime.Gosched()
			continue
		}
		v, ok := mapRef.get(key)
		if !mapRef.optLock.RVerify(mapVer) {
			runtime.Gosched()
			continue
		}
		return v, ok
	}
}

func (theMap *TheMap[K, V]) split(mapRef *smallMap[K, V], idx uint64) {
	if mapRef.bits == theMap.bits {
		theMap.dictlock.Lock()
		if mapRef.bits == theMap.bits {
			len := len(theMap.dictionary)
			newDict := make([]atomic.Pointer[smallMap[K, V]], len*2)
			copy(newDict, theMap.dictionary)
			copy(newDict[len:], theMap.dictionary)
			theMap.bits++
			theMap.dictionary = newDict
		}
		theMap.dictlock.Unlock()
	}

	highMask := uint64(1) << uint64(mapRef.bits)
	lowerBits := idx & (highMask - 1)
	this, splitThis, other, splitOther := theMap.splitMap(mapRef, highMask, &theMap.smallMapPool)
	atomic.AddUint64(&theMap.smaps, 1)

	for {
		version, ok := theMap.dictlock.RLock()
		dict := theMap.dictionary
		dictLen := uint64(len(dict))
		if !ok {
			runtime.Gosched()
			continue
		}
		for i := lowerBits; i < dictLen; i += highMask << 1 {
			dict[i].Store(this)
		}
		for i := lowerBits + highMask; i < dictLen; i += highMask << 1 {
			dict[i].Store(other)
		}
		if theMap.dictlock.RVerify(version) {
			break
		}
	}
	if splitThis {
		theMap.split(this, lowerBits)
	} else {
		this.optLock.Unlock()
	}

	if splitOther {
		theMap.split(other, lowerBits+highMask)
	} else {
		other.optLock.Unlock()
	}
	mapRef.reset()
	mapRef.optLock.Unlock()
	theMap.smallMapPool.Put(mapRef)
}

func (theMap *TheMap[K, V]) splitInsert(mapRef *smallMap[K, V], idx uint64, key K, value V) {
	theMap.split(mapRef, idx)
	theMap.Insert(key, value)
}

func New[K Hashable, V any]() *TheMap[K, V] {
	sM := &smallMap[K, V]{
		bits: 0,
	}
	dict := []atomic.Pointer[smallMap[K, V]]{{}, {}}
	dict[0].Store(sM)
	dict[1].Store(sM)
	return &TheMap[K, V]{
		bits:       1,
		dictionary: dict,
		smallMapPool: sync.Pool{
			New: func() any {
				return &smallMap[K, V]{}
			},
		},
	}
}