github.com/dolthub/dolt/go@v0.40.5-0.20240520175717-68db7794bea6/libraries/utils/concurrentmap/concurrentmap.go (about)

     1  // Copyright 2021 Dolthub, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package concurrentmap
    16  
    17  import (
    18  	"encoding/json"
    19  	"sync"
    20  )
    21  
    22  func New[K comparable, V any]() *Map[K, V] {
    23  	return &Map[K, V]{m: make(map[K]V)}
    24  }
    25  
    26  type Map[K comparable, V any] struct {
    27  	mu sync.RWMutex
    28  	m  map[K]V
    29  }
    30  
    31  // Get returns the value for the given key. If the key does not exist, the zero value for the value type will be returned.
    32  func (cm *Map[K, V]) Get(key K) (V, bool) {
    33  	cm.mu.RLock()
    34  	defer cm.mu.RUnlock()
    35  	if value, found := cm.m[key]; found {
    36  		return value, true
    37  	}
    38  	var zero V
    39  	return zero, false
    40  }
    41  
    42  // Set sets the value for the given key. If the key already exists, it will be overwritten.
    43  func (cm *Map[K, V]) Set(key K, value V) {
    44  	cm.mu.Lock()
    45  	defer cm.mu.Unlock()
    46  	cm.m[key] = value
    47  }
    48  
    49  // Delete removes the key from the map if it exists. If the key does not exist, this is a no-op.
    50  func (cm *Map[K, V]) Delete(key K) {
    51  	cm.mu.Lock()
    52  	defer cm.mu.Unlock()
    53  	delete(cm.m, key)
    54  }
    55  
    56  // Len returns the number of items in the map at the time of the call.
    57  func (cm *Map[K, V]) Len() int {
    58  	cm.mu.RLock()
    59  	defer cm.mu.RUnlock()
    60  	return len(cm.m)
    61  }
    62  
    63  // DeepCopy returns a deep copy of the concurrent map.
    64  func (cm *Map[K, V]) DeepCopy() *Map[K, V] {
    65  	cm.mu.RLock()
    66  	defer cm.mu.RUnlock()
    67  	newMap := make(map[K]V, len(cm.m))
    68  	for k, v := range cm.m {
    69  		newMap[k] = v
    70  	}
    71  	return &Map[K, V]{m: newMap}
    72  }
    73  
    74  // Iter iterates over the map, calling the provided function for each key/value pair. If the function returns false, the iteration stops.
    75  func (cm *Map[K, V]) Iter(f func(key K, value V) bool) {
    76  	cm.mu.RLock()
    77  	defer cm.mu.RUnlock()
    78  	for k, v := range cm.m {
    79  		if !f(k, v) {
    80  			break
    81  		}
    82  	}
    83  }
    84  
    85  // Snapshot returns a copy of the internal map at the time of the call. Returns a native map, not a concurrent one.
    86  func (cm *Map[K, V]) Snapshot() map[K]V {
    87  	cm.mu.RLock()
    88  	defer cm.mu.RUnlock()
    89  	newMap := make(map[K]V, len(cm.m))
    90  	for k, v := range cm.m {
    91  		newMap[k] = v
    92  	}
    93  	return newMap
    94  }
    95  
    96  func (cm *Map[K, V]) MarshalJSON() ([]byte, error) {
    97  	return json.Marshal(cm.Snapshot())
    98  }
    99  
   100  func (cm *Map[K, V]) UnmarshalJSON(data []byte) error {
   101  	return json.Unmarshal(data, &cm.m)
   102  }