github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/utils/unittest/protected_map.go (about)

     1  package unittest
     2  
     3  import "sync"
     4  
     5  // ProtectedMap is a thread-safe map.
     6  type ProtectedMap[K comparable, V any] struct {
     7  	mu sync.RWMutex
     8  	m  map[K]V
     9  }
    10  
    11  // NewProtectedMap returns a new ProtectedMap with the given types
    12  func NewProtectedMap[K comparable, V any]() *ProtectedMap[K, V] {
    13  	return &ProtectedMap[K, V]{
    14  		m: make(map[K]V),
    15  	}
    16  }
    17  
    18  // Add adds a key-value pair to the map
    19  func (p *ProtectedMap[K, V]) Add(key K, value V) {
    20  	p.mu.Lock()
    21  	defer p.mu.Unlock()
    22  	p.m[key] = value
    23  }
    24  
    25  // Remove removes a key-value pair from the map
    26  func (p *ProtectedMap[K, V]) Remove(key K) {
    27  	p.mu.Lock()
    28  	defer p.mu.Unlock()
    29  	delete(p.m, key)
    30  }
    31  
    32  // Has returns true if the map contains the given key
    33  func (p *ProtectedMap[K, V]) Has(key K) bool {
    34  	p.mu.RLock()
    35  	defer p.mu.RUnlock()
    36  	_, ok := p.m[key]
    37  	return ok
    38  }
    39  
    40  // Get returns the value for the given key and a boolean indicating if the key was found
    41  func (p *ProtectedMap[K, V]) Get(key K) (V, bool) {
    42  	p.mu.RLock()
    43  	defer p.mu.RUnlock()
    44  	value, ok := p.m[key]
    45  	return value, ok
    46  }
    47  
    48  // ForEach iterates over the map and calls the given function for each key-value pair.
    49  // If the function returns an error, the iteration is stopped and the error is returned.
    50  func (p *ProtectedMap[K, V]) ForEach(fn func(k K, v V) error) error {
    51  	p.mu.RLock()
    52  	defer p.mu.RUnlock()
    53  	for k, v := range p.m {
    54  		if err := fn(k, v); err != nil {
    55  			return err
    56  		}
    57  	}
    58  	return nil
    59  }
    60  
    61  // Size returns the size of the map.
    62  func (p *ProtectedMap[K, V]) Size() int {
    63  	p.mu.RLock()
    64  	defer p.mu.RUnlock()
    65  	return len(p.m)
    66  }