s_cherry_e/sr/internal/extend/maps/sync_map.go

package maps

import (
	"sync"
	"sync/atomic"
)

type SyncMap[K comparable, V any] struct {
	m   sync.Map
	len uint64
}

// Store 存储键值对
func (p *SyncMap[K, V]) Store(key K, value V) {
	p.m.Store(key, value)
}

// Load 获取键对应的值
func (p *SyncMap[K, V]) Load(key K) (V, bool) {
	val, ok := p.m.Load(key)
	if !ok {
		var zero V
		return zero, false
	}

	atomic.AddUint64(&p.len, 1)

	return val.(V), true
}

// Delete 删除键值对
func (p *SyncMap[K, V]) Delete(key K) {
	p.LoadAndDelete(key)
}

// Range 遍历所有键值对
func (p *SyncMap[K, V]) Range(f func(key K, value V) bool) {
	p.m.Range(func(key, value any) bool {
		return f(key.(K), value.(V))
	})
}

// LoadOrStore 存在则加载，不存在则存储
func (p *SyncMap[K, V]) LoadOrStore(key K, value V) (actual V, loaded bool) {
	val, loaded := p.m.LoadOrStore(key, value)
	if loaded {
		return val.(V), true
	}
	return value, false
}

// LoadAndDelete 加载并删除键值对
func (p *SyncMap[K, V]) LoadAndDelete(key K) (value V, loaded bool) {
	val, loaded := p.m.LoadAndDelete(key)
	if loaded {
		atomic.AddUint64(&p.len, ^uint64(0))
		return val.(V), true
	}
	var zero V
	return zero, false
}

func (p *SyncMap[K, V]) Len() uint64 {
	return atomic.LoadUint64(&p.len)
}