1package csync
2
3import (
4 "encoding/json"
5 "iter"
6 "maps"
7 "sync"
8)
9
10// Map is a concurrent map implementation that provides thread-safe access.
11type Map[K comparable, V any] struct {
12 inner map[K]V
13 mu sync.RWMutex
14}
15
16// NewMap creates a new thread-safe map with the specified key and value types.
17func NewMap[K comparable, V any]() *Map[K, V] {
18 return &Map[K, V]{
19 inner: make(map[K]V),
20 }
21}
22
23// NewMapFrom creates a new thread-safe map from an existing map.
24func NewMapFrom[K comparable, V any](m map[K]V) *Map[K, V] {
25 return &Map[K, V]{
26 inner: m,
27 }
28}
29
30// NewMapFromSeq creates a new thread-safe map from an iter.Seq2 of key-value pairs.
31func NewMapFromSeq[k comparable, v any](seq iter.Seq2[k, v]) *Map[k, v] {
32 m := make(map[k]v)
33 seq(func(kk k, vv v) bool {
34 m[kk] = vv
35 return true
36 })
37 return NewMapFrom(m)
38}
39
40// NewLazyMap creates a new lazy-loaded map. The provided load function is
41// executed in a separate goroutine to populate the map.
42func NewLazyMap[K comparable, V any](load func() map[K]V) *Map[K, V] {
43 m := &Map[K, V]{}
44 m.mu.Lock()
45 go func() {
46 m.inner = load()
47 m.mu.Unlock()
48 }()
49 return m
50}
51
52// Set sets the value for the specified key in the map.
53func (m *Map[K, V]) Set(key K, value V) {
54 m.mu.Lock()
55 defer m.mu.Unlock()
56 m.inner[key] = value
57}
58
59// Del deletes the specified key from the map.
60func (m *Map[K, V]) Del(key K) {
61 m.mu.Lock()
62 defer m.mu.Unlock()
63 delete(m.inner, key)
64}
65
66// Get gets the value for the specified key from the map.
67func (m *Map[K, V]) Get(key K) (V, bool) {
68 m.mu.RLock()
69 defer m.mu.RUnlock()
70 v, ok := m.inner[key]
71 return v, ok
72}
73
74// Len returns the number of items in the map.
75func (m *Map[K, V]) Len() int {
76 m.mu.RLock()
77 defer m.mu.RUnlock()
78 return len(m.inner)
79}
80
81// GetOrSet gets and returns the key if it exists, otherwise, it executes the
82// given function, set its return value for the given key, and returns it.
83func (m *Map[K, V]) GetOrSet(key K, fn func() V) V {
84 got, ok := m.Get(key)
85 if ok {
86 return got
87 }
88 value := fn()
89 m.Set(key, value)
90 return value
91}
92
93// Take gets an item and then deletes it.
94func (m *Map[K, V]) Take(key K) (V, bool) {
95 m.mu.Lock()
96 defer m.mu.Unlock()
97 v, ok := m.inner[key]
98 delete(m.inner, key)
99 return v, ok
100}
101
102// Seq2 returns an iter.Seq2 that yields key-value pairs from the map.
103func (m *Map[K, V]) Seq2() iter.Seq2[K, V] {
104 dst := make(map[K]V)
105 m.mu.RLock()
106 maps.Copy(dst, m.inner)
107 m.mu.RUnlock()
108 return func(yield func(K, V) bool) {
109 for k, v := range dst {
110 if !yield(k, v) {
111 return
112 }
113 }
114 }
115}
116
117// Seq returns an iter.Seq that yields values from the map.
118func (m *Map[K, V]) Seq() iter.Seq[V] {
119 return func(yield func(V) bool) {
120 for _, v := range m.Seq2() {
121 if !yield(v) {
122 return
123 }
124 }
125 }
126}
127
128var (
129 _ json.Unmarshaler = &Map[string, any]{}
130 _ json.Marshaler = &Map[string, any]{}
131)
132
133func (Map[K, V]) JSONSchemaAlias() any { //nolint
134 m := map[K]V{}
135 return m
136}
137
138// UnmarshalJSON implements json.Unmarshaler.
139func (m *Map[K, V]) UnmarshalJSON(data []byte) error {
140 m.mu.Lock()
141 defer m.mu.Unlock()
142 m.inner = make(map[K]V)
143 return json.Unmarshal(data, &m.inner)
144}
145
146// MarshalJSON implements json.Marshaler.
147func (m *Map[K, V]) MarshalJSON() ([]byte, error) {
148 m.mu.RLock()
149 defer m.mu.RUnlock()
150 return json.Marshal(m.inner)
151}