github.com/lrita/cmap@v0.0.0-20231108122212-cb084a67f554/cmap_reference_test.go (about)

     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  
     5  package cmap_test
     6  
     7  import (
     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.
    15  type mapInterface interface {
    16  	Load(interface{}) (interface{}, bool)
    17  	Store(key, value interface{})
    18  	LoadOrStore(key, value interface{}) (actual interface{}, loaded bool)
    19  	Delete(interface{})
    20  	Range(func(key, value interface{}) (shouldContinue bool))
    21  }
    22  
    23  // RWMutexMap is an implementation of mapInterface using a sync.RWMutex.
    24  type RWMutexMap struct {
    25  	mu    sync.RWMutex
    26  	dirty map[interface{}]interface{}
    27  }
    28  
    29  func (m *RWMutexMap) Load(key interface{}) (value interface{}, ok bool) {
    30  	m.mu.RLock()
    31  	value, ok = m.dirty[key]
    32  	m.mu.RUnlock()
    33  	return
    34  }
    35  
    36  func (m *RWMutexMap) Store(key, value interface{}) {
    37  	m.mu.Lock()
    38  	if m.dirty == nil {
    39  		m.dirty = make(map[interface{}]interface{})
    40  	}
    41  	m.dirty[key] = value
    42  	m.mu.Unlock()
    43  }
    44  
    45  func (m *RWMutexMap) LoadOrStore(key, value interface{}) (actual interface{}, loaded bool) {
    46  	m.mu.Lock()
    47  	actual, loaded = m.dirty[key]
    48  	if !loaded {
    49  		actual = value
    50  		if m.dirty == nil {
    51  			m.dirty = make(map[interface{}]interface{})
    52  		}
    53  		m.dirty[key] = value
    54  	}
    55  	m.mu.Unlock()
    56  	return actual, loaded
    57  }
    58  
    59  func (m *RWMutexMap) Delete(key interface{}) {
    60  	m.mu.Lock()
    61  	delete(m.dirty, key)
    62  	m.mu.Unlock()
    63  }
    64  
    65  func (m *RWMutexMap) Range(f func(key, value interface{}) (shouldContinue bool)) {
    66  	m.mu.RLock()
    67  	keys := make([]interface{}, 0, len(m.dirty))
    68  	for k := range m.dirty {
    69  		keys = append(keys, k)
    70  	}
    71  	m.mu.RUnlock()
    72  
    73  	for _, k := range keys {
    74  		v, ok := m.Load(k)
    75  		if !ok {
    76  			continue
    77  		}
    78  		if !f(k, v) {
    79  			break
    80  		}
    81  	}
    82  }
    83  
    84  // DeepCopyMap is an implementation of mapInterface using a Mutex and
    85  // atomic.Value.  It makes deep copies of the map on every write to avoid
    86  // acquiring the Mutex in Load.
    87  type DeepCopyMap struct {
    88  	mu    sync.Mutex
    89  	clean atomic.Value
    90  }
    91  
    92  func (m *DeepCopyMap) Load(key interface{}) (value interface{}, ok bool) {
    93  	clean, _ := m.clean.Load().(map[interface{}]interface{})
    94  	value, ok = clean[key]
    95  	return value, ok
    96  }
    97  
    98  func (m *DeepCopyMap) Store(key, value interface{}) {
    99  	m.mu.Lock()
   100  	dirty := m.dirty()
   101  	dirty[key] = value
   102  	m.clean.Store(dirty)
   103  	m.mu.Unlock()
   104  }
   105  
   106  func (m *DeepCopyMap) LoadOrStore(key, value interface{}) (actual interface{}, loaded bool) {
   107  	clean, _ := m.clean.Load().(map[interface{}]interface{})
   108  	actual, loaded = clean[key]
   109  	if loaded {
   110  		return actual, loaded
   111  	}
   112  
   113  	m.mu.Lock()
   114  	// Reload clean in case it changed while we were waiting on m.mu.
   115  	clean, _ = m.clean.Load().(map[interface{}]interface{})
   116  	actual, loaded = clean[key]
   117  	if !loaded {
   118  		dirty := m.dirty()
   119  		dirty[key] = value
   120  		actual = value
   121  		m.clean.Store(dirty)
   122  	}
   123  	m.mu.Unlock()
   124  	return actual, loaded
   125  }
   126  
   127  func (m *DeepCopyMap) Delete(key interface{}) {
   128  	m.mu.Lock()
   129  	dirty := m.dirty()
   130  	delete(dirty, key)
   131  	m.clean.Store(dirty)
   132  	m.mu.Unlock()
   133  }
   134  
   135  func (m *DeepCopyMap) Range(f func(key, value interface{}) (shouldContinue bool)) {
   136  	clean, _ := m.clean.Load().(map[interface{}]interface{})
   137  	for k, v := range clean {
   138  		if !f(k, v) {
   139  			break
   140  		}
   141  	}
   142  }
   143  
   144  func (m *DeepCopyMap) dirty() map[interface{}]interface{} {
   145  	clean, _ := m.clean.Load().(map[interface{}]interface{})
   146  	dirty := make(map[interface{}]interface{}, len(clean)+1)
   147  	for k, v := range clean {
   148  		dirty[k] = v
   149  	}
   150  	return dirty
   151  }