github.com/ydb-platform/ydb-go-sdk/v3@v3.89.2/internal/xsync/map.go (about)

     1  package xsync
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  	"sync/atomic"
     7  )
     8  
     9  type Map[K comparable, V any] struct {
    10  	m    sync.Map
    11  	size atomic.Int32
    12  }
    13  
    14  func (m *Map[K, V]) Get(key K) (value V, ok bool) {
    15  	v, ok := m.m.Load(key)
    16  	if !ok {
    17  		return value, false
    18  	}
    19  	value, ok = v.(V)
    20  
    21  	return value, ok
    22  }
    23  
    24  func (m *Map[K, V]) Must(key K) (value V) {
    25  	v, ok := m.m.Load(key)
    26  	if !ok {
    27  		panic(fmt.Sprintf("unexpected key = %v", key))
    28  	}
    29  	value, ok = v.(V)
    30  	if !ok {
    31  		panic(fmt.Sprintf("unexpected type of value = %T", v))
    32  	}
    33  
    34  	return value
    35  }
    36  
    37  func (m *Map[K, V]) Has(key K) bool {
    38  	_, ok := m.m.Load(key)
    39  
    40  	return ok
    41  }
    42  
    43  func (m *Map[K, V]) Set(key K, value V) {
    44  	_, exists := m.m.LoadOrStore(key, value)
    45  
    46  	if !exists {
    47  		m.size.Add(1)
    48  	}
    49  }
    50  
    51  func (m *Map[K, V]) Delete(key K) bool {
    52  	_, exists := m.Extract(key)
    53  
    54  	if !exists {
    55  		m.size.Add(-1)
    56  	}
    57  
    58  	return exists
    59  }
    60  
    61  func (m *Map[K, V]) Len() int {
    62  	return int(m.size.Load())
    63  }
    64  
    65  func (m *Map[K, V]) Extract(key K) (value V, ok bool) {
    66  	v, exists := m.m.LoadAndDelete(key)
    67  	if !exists {
    68  		return value, false
    69  	}
    70  
    71  	m.size.Add(-1)
    72  
    73  	value, ok = v.(V)
    74  
    75  	return value, ok
    76  }
    77  
    78  func (m *Map[K, V]) Range(f func(key K, value V) bool) {
    79  	m.m.Range(func(k, v any) bool {
    80  		return f(k.(K), v.(V)) //nolint:forcetypeassert
    81  	})
    82  }
    83  
    84  func (m *Map[K, V]) Clear() (removed int) {
    85  	m.m.Range(func(k, v any) bool {
    86  		removed++
    87  
    88  		m.m.Delete(k)
    89  
    90  		return true
    91  	})
    92  
    93  	m.size.Add(int32(-removed))
    94  
    95  	return removed
    96  }