github.com/machinefi/w3bstream@v1.6.5-rc9.0.20240426031326-b8c7c4876e72/pkg/depends/x/mapx/map.go (about)

     1  package mapx
     2  
     3  import "sync"
     4  
     5  type Map[K comparable, V any] struct {
     6  	val map[K]V
     7  	mtx *sync.RWMutex
     8  }
     9  
    10  func New[K comparable, V any]() *Map[K, V] {
    11  	return &Map[K, V]{val: make(map[K]V), mtx: &sync.RWMutex{}}
    12  }
    13  
    14  func (m *Map[K, V]) Load(k K) (v V, ok bool) {
    15  	m.mtx.RLock()
    16  	defer m.mtx.RUnlock()
    17  	v, ok = m.val[k]
    18  	return
    19  }
    20  
    21  func (m *Map[K, V]) Store(k K, v V) {
    22  	m.mtx.Lock()
    23  	defer m.mtx.Unlock()
    24  	m.val[k] = v
    25  }
    26  
    27  func (m *Map[K, V]) StoreNX(k K, v V) bool {
    28  	m.mtx.Lock()
    29  	defer m.mtx.Unlock()
    30  	_, ok := m.val[k]
    31  	if ok {
    32  		return false
    33  	}
    34  	m.val[k] = v
    35  	return true
    36  }
    37  
    38  func (m *Map[K, V]) LoadOrStore(k K, newv func() (V, error)) (V, error) {
    39  	m.mtx.Lock()
    40  	defer m.mtx.Unlock()
    41  	v, ok := m.val[k]
    42  	if ok {
    43  		return v, nil
    44  	}
    45  	v, err := newv()
    46  	if err != nil {
    47  		return v, err
    48  	}
    49  
    50  	m.val[k] = v
    51  	return v, nil
    52  }
    53  
    54  func (m *Map[K, V]) Remove(k K) {
    55  	m.mtx.Lock()
    56  	defer m.mtx.Unlock()
    57  	delete(m.val, k)
    58  }
    59  
    60  func (m *Map[K, V]) Clear() {
    61  	m.mtx.Lock()
    62  	defer m.mtx.Unlock()
    63  	m.val = make(map[K]V)
    64  }
    65  
    66  func (m *Map[K, V]) LoadAndRemove(k K) (v V, ok bool) {
    67  	m.mtx.Lock()
    68  	defer m.mtx.Unlock()
    69  	v, ok = m.val[k]
    70  	if ok {
    71  		defer delete(m.val, k)
    72  	}
    73  	return v, ok
    74  }
    75  
    76  func (m *Map[K, V]) Len() int {
    77  	m.mtx.Lock()
    78  	defer m.mtx.Unlock()
    79  	return len(m.val)
    80  }
    81  
    82  func (m *Map[K, V]) Range(f func(k K, v V) bool) {
    83  	m.mtx.Lock()
    84  	defer m.mtx.Unlock()
    85  	for k, v := range m.val {
    86  		if !f(k, v) {
    87  			break
    88  		}
    89  	}
    90  }