1// Copyright 2016 The Go Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5package sync_test
6
7import (
8	"sync"
9	"sync/atomic"
10)
11
12// This file contains reference map implementations for unit-tests.
13
14// mapInterface is the interface Map implements.
15type mapInterface interface {
16	Load(key any) (value any, ok bool)
17	Store(key, value any)
18	LoadOrStore(key, value any) (actual any, loaded bool)
19	LoadAndDelete(key any) (value any, loaded bool)
20	Delete(any)
21	Swap(key, value any) (previous any, loaded bool)
22	CompareAndSwap(key, old, new any) (swapped bool)
23	CompareAndDelete(key, old any) (deleted bool)
24	Range(func(key, value any) (shouldContinue bool))
25	Clear()
26}
27
28var (
29	_ mapInterface = &RWMutexMap{}
30	_ mapInterface = &DeepCopyMap{}
31)
32
33// RWMutexMap is an implementation of mapInterface using a sync.RWMutex.
34type RWMutexMap struct {
35	mu    sync.RWMutex
36	dirty map[any]any
37}
38
39func (m *RWMutexMap) Load(key any) (value any, ok bool) {
40	m.mu.RLock()
41	value, ok = m.dirty[key]
42	m.mu.RUnlock()
43	return
44}
45
46func (m *RWMutexMap) Store(key, value any) {
47	m.mu.Lock()
48	if m.dirty == nil {
49		m.dirty = make(map[any]any)
50	}
51	m.dirty[key] = value
52	m.mu.Unlock()
53}
54
55func (m *RWMutexMap) LoadOrStore(key, value any) (actual any, loaded bool) {
56	m.mu.Lock()
57	actual, loaded = m.dirty[key]
58	if !loaded {
59		actual = value
60		if m.dirty == nil {
61			m.dirty = make(map[any]any)
62		}
63		m.dirty[key] = value
64	}
65	m.mu.Unlock()
66	return actual, loaded
67}
68
69func (m *RWMutexMap) Swap(key, value any) (previous any, loaded bool) {
70	m.mu.Lock()
71	if m.dirty == nil {
72		m.dirty = make(map[any]any)
73	}
74
75	previous, loaded = m.dirty[key]
76	m.dirty[key] = value
77	m.mu.Unlock()
78	return
79}
80
81func (m *RWMutexMap) LoadAndDelete(key any) (value any, loaded bool) {
82	m.mu.Lock()
83	value, loaded = m.dirty[key]
84	if !loaded {
85		m.mu.Unlock()
86		return nil, false
87	}
88	delete(m.dirty, key)
89	m.mu.Unlock()
90	return value, loaded
91}
92
93func (m *RWMutexMap) Delete(key any) {
94	m.mu.Lock()
95	delete(m.dirty, key)
96	m.mu.Unlock()
97}
98
99func (m *RWMutexMap) CompareAndSwap(key, old, new any) (swapped bool) {
100	m.mu.Lock()
101	defer m.mu.Unlock()
102	if m.dirty == nil {
103		return false
104	}
105
106	value, loaded := m.dirty[key]
107	if loaded && value == old {
108		m.dirty[key] = new
109		return true
110	}
111	return false
112}
113
114func (m *RWMutexMap) CompareAndDelete(key, old any) (deleted bool) {
115	m.mu.Lock()
116	defer m.mu.Unlock()
117	if m.dirty == nil {
118		return false
119	}
120
121	value, loaded := m.dirty[key]
122	if loaded && value == old {
123		delete(m.dirty, key)
124		return true
125	}
126	return false
127}
128
129func (m *RWMutexMap) Range(f func(key, value any) (shouldContinue bool)) {
130	m.mu.RLock()
131	keys := make([]any, 0, len(m.dirty))
132	for k := range m.dirty {
133		keys = append(keys, k)
134	}
135	m.mu.RUnlock()
136
137	for _, k := range keys {
138		v, ok := m.Load(k)
139		if !ok {
140			continue
141		}
142		if !f(k, v) {
143			break
144		}
145	}
146}
147
148func (m *RWMutexMap) Clear() {
149	m.mu.Lock()
150	defer m.mu.Unlock()
151
152	clear(m.dirty)
153}
154
155// DeepCopyMap is an implementation of mapInterface using a Mutex and
156// atomic.Value.  It makes deep copies of the map on every write to avoid
157// acquiring the Mutex in Load.
158type DeepCopyMap struct {
159	mu    sync.Mutex
160	clean atomic.Value
161}
162
163func (m *DeepCopyMap) Load(key any) (value any, ok bool) {
164	clean, _ := m.clean.Load().(map[any]any)
165	value, ok = clean[key]
166	return value, ok
167}
168
169func (m *DeepCopyMap) Store(key, value any) {
170	m.mu.Lock()
171	dirty := m.dirty()
172	dirty[key] = value
173	m.clean.Store(dirty)
174	m.mu.Unlock()
175}
176
177func (m *DeepCopyMap) LoadOrStore(key, value any) (actual any, loaded bool) {
178	clean, _ := m.clean.Load().(map[any]any)
179	actual, loaded = clean[key]
180	if loaded {
181		return actual, loaded
182	}
183
184	m.mu.Lock()
185	// Reload clean in case it changed while we were waiting on m.mu.
186	clean, _ = m.clean.Load().(map[any]any)
187	actual, loaded = clean[key]
188	if !loaded {
189		dirty := m.dirty()
190		dirty[key] = value
191		actual = value
192		m.clean.Store(dirty)
193	}
194	m.mu.Unlock()
195	return actual, loaded
196}
197
198func (m *DeepCopyMap) Swap(key, value any) (previous any, loaded bool) {
199	m.mu.Lock()
200	dirty := m.dirty()
201	previous, loaded = dirty[key]
202	dirty[key] = value
203	m.clean.Store(dirty)
204	m.mu.Unlock()
205	return
206}
207
208func (m *DeepCopyMap) LoadAndDelete(key any) (value any, loaded bool) {
209	m.mu.Lock()
210	dirty := m.dirty()
211	value, loaded = dirty[key]
212	delete(dirty, key)
213	m.clean.Store(dirty)
214	m.mu.Unlock()
215	return
216}
217
218func (m *DeepCopyMap) Delete(key any) {
219	m.mu.Lock()
220	dirty := m.dirty()
221	delete(dirty, key)
222	m.clean.Store(dirty)
223	m.mu.Unlock()
224}
225
226func (m *DeepCopyMap) CompareAndSwap(key, old, new any) (swapped bool) {
227	clean, _ := m.clean.Load().(map[any]any)
228	if previous, ok := clean[key]; !ok || previous != old {
229		return false
230	}
231
232	m.mu.Lock()
233	defer m.mu.Unlock()
234	dirty := m.dirty()
235	value, loaded := dirty[key]
236	if loaded && value == old {
237		dirty[key] = new
238		m.clean.Store(dirty)
239		return true
240	}
241	return false
242}
243
244func (m *DeepCopyMap) CompareAndDelete(key, old any) (deleted bool) {
245	clean, _ := m.clean.Load().(map[any]any)
246	if previous, ok := clean[key]; !ok || previous != old {
247		return false
248	}
249
250	m.mu.Lock()
251	defer m.mu.Unlock()
252
253	dirty := m.dirty()
254	value, loaded := dirty[key]
255	if loaded && value == old {
256		delete(dirty, key)
257		m.clean.Store(dirty)
258		return true
259	}
260	return false
261}
262
263func (m *DeepCopyMap) Range(f func(key, value any) (shouldContinue bool)) {
264	clean, _ := m.clean.Load().(map[any]any)
265	for k, v := range clean {
266		if !f(k, v) {
267			break
268		}
269	}
270}
271
272func (m *DeepCopyMap) dirty() map[any]any {
273	clean, _ := m.clean.Load().(map[any]any)
274	dirty := make(map[any]any, len(clean)+1)
275	for k, v := range clean {
276		dirty[k] = v
277	}
278	return dirty
279}
280
281func (m *DeepCopyMap) Clear() {
282	m.mu.Lock()
283	defer m.mu.Unlock()
284
285	m.clean.Store((map[any]any)(nil))
286}
287