exist.go

package exist

import (
	"sync"
	"sync/atomic"
)

type ExistCache struct {
	growLock  sync.RWMutex
	indexMask uint64
	reprobe   uint64
	values    []atomic.Value
}

type node struct {
	key       uint64
	tombstone bool
}

func NewExistCache(size uint64) *ExistCache {
	// NOTE: upsize must be greater than size *and* power of 2
	var upsize uint64 = 8
	for size > upsize {
		upsize *= 2
	}
	return &ExistCache{
		indexMask: upsize - 1,
		reprobe:   8,
		values:    make([]atomic.Value, upsize),
	}
}

func (ec *ExistCache) Check(key uint64) bool {
	ec.growLock.RLock()
	defer ec.growLock.RUnlock()

	index := key
	for distance := uint64(0); distance < ec.reprobe; distance++ {
		index &= ec.indexMask
		maybe := ec.values[index].Load()

		value, ok := maybe.(node)
		if !ok {
			return false
		}
		if value.key == key {
			return !value.tombstone
		}

		index++
	}
	return false
}

func (ec *ExistCache) Clear(key uint64) {
	index := key
	ec.growLock.RLock()
	defer ec.growLock.RUnlock()

	for distance := uint64(0); distance < ec.reprobe; distance++ {
		index &= ec.indexMask
		maybe := ec.values[index].Load()
		value, ok := maybe.(node)
		if !ok {
			return
		}
		if value.key == key {
			value.tombstone = true
			ec.values[index].Store(value)
			return
		}

		index++
	}
}

func (ec *ExistCache) Set(key uint64) {
	ec.growLock.RLock()

	index := key
	for distance := uint64(0); distance < ec.reprobe; distance++ {
		index &= ec.indexMask
		maybe := ec.values[index].Load()
		value, ok := maybe.(node)
		if !ok {
			ec.values[index].Store(node{key: key})
			ec.growLock.RUnlock()
			return
		}
		if value.key == key {
			if value.tombstone {
				value.tombstone = false
				ec.values[index].Store(value)
			}
			ec.growLock.RUnlock()
			return
		}

		index++
	}

	ec.growLock.RUnlock()
	ec.growAndSet(key)
}

func (ec *ExistCache) growAndSet(key uint64) {
	ec.growLock.Lock()

	nec := NewExistCache((ec.indexMask + 1) << 1)

	for index := uint64(0); index < ec.indexMask; index++ {
		maybe := ec.values[index].Load()
		value, ok := maybe.(node)
		if ok && !value.tombstone {
			nec.Set(value.key)
		}
	}
	nec.Set(key)

	ec.indexMask = nec.indexMask
	ec.values = nec.values
	ec.reprobe = nec.reprobe
	ec.growLock.Unlock()
}