github.com/MetalBlockchain/metalgo@v1.11.9/utils/setmap/setmap.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package setmap
     5  
     6  import (
     7  	"github.com/MetalBlockchain/metalgo/utils"
     8  	"github.com/MetalBlockchain/metalgo/utils/set"
     9  )
    10  
    11  type Entry[K any, V comparable] struct {
    12  	Key K
    13  	Set set.Set[V]
    14  }
    15  
    16  // SetMap is a map to a set where all sets are non-overlapping.
    17  type SetMap[K, V comparable] struct {
    18  	keyToSet   map[K]set.Set[V]
    19  	valueToKey map[V]K
    20  }
    21  
    22  // New creates a new empty setmap.
    23  func New[K, V comparable]() *SetMap[K, V] {
    24  	return &SetMap[K, V]{
    25  		keyToSet:   make(map[K]set.Set[V]),
    26  		valueToKey: make(map[V]K),
    27  	}
    28  }
    29  
    30  // Put the new entry into the map. Removes and returns:
    31  // * The existing entry for [key].
    32  // * Existing entries where the set overlaps with the [set].
    33  func (m *SetMap[K, V]) Put(key K, set set.Set[V]) []Entry[K, V] {
    34  	removed := m.DeleteOverlapping(set)
    35  	if removedSet, ok := m.DeleteKey(key); ok {
    36  		removed = append(removed, Entry[K, V]{
    37  			Key: key,
    38  			Set: removedSet,
    39  		})
    40  	}
    41  
    42  	m.keyToSet[key] = set
    43  	for val := range set {
    44  		m.valueToKey[val] = key
    45  	}
    46  	return removed
    47  }
    48  
    49  // GetKey that maps to the provided value.
    50  func (m *SetMap[K, V]) GetKey(val V) (K, bool) {
    51  	key, ok := m.valueToKey[val]
    52  	return key, ok
    53  }
    54  
    55  // GetSet that is mapped to by the provided key.
    56  func (m *SetMap[K, V]) GetSet(key K) (set.Set[V], bool) {
    57  	val, ok := m.keyToSet[key]
    58  	return val, ok
    59  }
    60  
    61  // HasKey returns true if [key] is in the map.
    62  func (m *SetMap[K, _]) HasKey(key K) bool {
    63  	_, ok := m.keyToSet[key]
    64  	return ok
    65  }
    66  
    67  // HasValue returns true if [val] is in a set in the map.
    68  func (m *SetMap[_, V]) HasValue(val V) bool {
    69  	_, ok := m.valueToKey[val]
    70  	return ok
    71  }
    72  
    73  // HasOverlap returns true if [set] overlaps with any of the sets in the map.
    74  func (m *SetMap[_, V]) HasOverlap(set set.Set[V]) bool {
    75  	if set.Len() < len(m.valueToKey) {
    76  		for val := range set {
    77  			if _, ok := m.valueToKey[val]; ok {
    78  				return true
    79  			}
    80  		}
    81  	} else {
    82  		for val := range m.valueToKey {
    83  			if set.Contains(val) {
    84  				return true
    85  			}
    86  		}
    87  	}
    88  	return false
    89  }
    90  
    91  // DeleteKey removes [key] from the map and returns the set it mapped to.
    92  func (m *SetMap[K, V]) DeleteKey(key K) (set.Set[V], bool) {
    93  	set, ok := m.keyToSet[key]
    94  	if !ok {
    95  		return nil, false
    96  	}
    97  
    98  	delete(m.keyToSet, key)
    99  	for val := range set {
   100  		delete(m.valueToKey, val)
   101  	}
   102  	return set, true
   103  }
   104  
   105  // DeleteValue removes and returns the entry that contained [val].
   106  func (m *SetMap[K, V]) DeleteValue(val V) (K, set.Set[V], bool) {
   107  	key, ok := m.valueToKey[val]
   108  	if !ok {
   109  		return utils.Zero[K](), nil, false
   110  	}
   111  	set, _ := m.DeleteKey(key)
   112  	return key, set, true
   113  }
   114  
   115  // DeleteOverlapping removes and returns all the entries where the set overlaps
   116  // with [set].
   117  func (m *SetMap[K, V]) DeleteOverlapping(set set.Set[V]) []Entry[K, V] {
   118  	var removed []Entry[K, V]
   119  	for val := range set {
   120  		if k, removedSet, ok := m.DeleteValue(val); ok {
   121  			removed = append(removed, Entry[K, V]{
   122  				Key: k,
   123  				Set: removedSet,
   124  			})
   125  		}
   126  	}
   127  	return removed
   128  }
   129  
   130  // Len return the number of sets in the map.
   131  func (m *SetMap[K, V]) Len() int {
   132  	return len(m.keyToSet)
   133  }
   134  
   135  // LenValues return the total number of values across all sets in the map.
   136  func (m *SetMap[K, V]) LenValues() int {
   137  	return len(m.valueToKey)
   138  }