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// NewLazyMap creates a new lazy-loaded map. The provided load function is
31// executed in a separate goroutine to populate the map.
32func NewLazyMap[K comparable, V any](load func() map[K]V) *Map[K, V] {
33 m := &Map[K, V]{}
34 m.mu.Lock()
35 go func() {
36 m.inner = load()
37 m.mu.Unlock()
38 }()
39 return m
40}
41
42// Reset replaces the inner map with the new one.
43func (m *Map[K, V]) Reset(input map[K]V) {
44 m.mu.Lock()
45 defer m.mu.Unlock()
46 m.inner = input
47}
48
49// Set sets the value for the specified key in the map.
50func (m *Map[K, V]) Set(key K, value V) {
51 m.mu.Lock()
52 defer m.mu.Unlock()
53 m.inner[key] = value
54}
55
56// Del deletes the specified key from the map.
57func (m *Map[K, V]) Del(key K) {
58 m.mu.Lock()
59 defer m.mu.Unlock()
60 delete(m.inner, key)
61}
62
63// Get gets the value for the specified key from the map.
64func (m *Map[K, V]) Get(key K) (V, bool) {
65 m.mu.RLock()
66 defer m.mu.RUnlock()
67 v, ok := m.inner[key]
68 return v, ok
69}
70
71// Len returns the number of items in the map.
72func (m *Map[K, V]) Len() int {
73 m.mu.RLock()
74 defer m.mu.RUnlock()
75 return len(m.inner)
76}
77
78// GetOrSet gets and returns the key if it exists, otherwise, it executes the
79// given function, set its return value for the given key, and returns it.
80func (m *Map[K, V]) GetOrSet(key K, fn func() V) V {
81 got, ok := m.Get(key)
82 if ok {
83 return got
84 }
85 value := fn()
86 m.Set(key, value)
87 return value
88}
89
90// Take gets an item and then deletes it.
91func (m *Map[K, V]) Take(key K) (V, bool) {
92 m.mu.Lock()
93 defer m.mu.Unlock()
94 v, ok := m.inner[key]
95 delete(m.inner, key)
96 return v, ok
97}
98
99// Seq2 returns an iter.Seq2 that yields key-value pairs from the map.
100func (m *Map[K, V]) Seq2() iter.Seq2[K, V] {
101 dst := make(map[K]V)
102 m.mu.RLock()
103 maps.Copy(dst, m.inner)
104 m.mu.RUnlock()
105 return func(yield func(K, V) bool) {
106 for k, v := range dst {
107 if !yield(k, v) {
108 return
109 }
110 }
111 }
112}
113
114// Seq returns an iter.Seq that yields values from the map.
115func (m *Map[K, V]) Seq() iter.Seq[V] {
116 return func(yield func(V) bool) {
117 for _, v := range m.Seq2() {
118 if !yield(v) {
119 return
120 }
121 }
122 }
123}
124
125var (
126 _ json.Unmarshaler = &Map[string, any]{}
127 _ json.Marshaler = &Map[string, any]{}
128)
129
130func (Map[K, V]) JSONSchemaAlias() any { //nolint
131 m := map[K]V{}
132 return m
133}
134
135// UnmarshalJSON implements json.Unmarshaler.
136func (m *Map[K, V]) UnmarshalJSON(data []byte) error {
137 m.mu.Lock()
138 defer m.mu.Unlock()
139 m.inner = make(map[K]V)
140 return json.Unmarshal(data, &m.inner)
141}
142
143// MarshalJSON implements json.Marshaler.
144func (m *Map[K, V]) MarshalJSON() ([]byte, error) {
145 m.mu.RLock()
146 defer m.mu.RUnlock()
147 return json.Marshal(m.inner)
148}