github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/pkg/genutil/mapz/multimap.go (about)

     1  package mapz
     2  
     3  import (
     4  	"golang.org/x/exp/maps"
     5  )
     6  
     7  // ReadOnlyMultimap is a read-only multimap.
     8  type ReadOnlyMultimap[T comparable, Q any] interface {
     9  	// Has returns true if the key is found in the map.
    10  	Has(key T) bool
    11  
    12  	// Get returns the values for the given key in the map and whether the key
    13  	// existed.
    14  	// If the key does not exist, an empty slice is returned.
    15  	Get(key T) ([]Q, bool)
    16  
    17  	// IsEmpty returns true if the map is currently empty.
    18  	IsEmpty() bool
    19  
    20  	// Len returns the length of the map, e.g. the number of *keys* present.
    21  	Len() int
    22  
    23  	// Keys returns the keys of the map.
    24  	Keys() []T
    25  
    26  	// Values returns all values in the map.
    27  	Values() []Q
    28  }
    29  
    30  // NewMultiMap initializes a new MultiMap.
    31  func NewMultiMap[T comparable, Q any]() *MultiMap[T, Q] {
    32  	return &MultiMap[T, Q]{items: map[T][]Q{}}
    33  }
    34  
    35  // NewMultiMapWithCap initializes with the provided capacity for the top-level
    36  // map.
    37  func NewMultiMapWithCap[T comparable, Q any](capacity uint32) *MultiMap[T, Q] {
    38  	return &MultiMap[T, Q]{items: make(map[T][]Q, capacity)}
    39  }
    40  
    41  // MultiMap represents a map that can contain 1 or more values for each key.
    42  type MultiMap[T comparable, Q any] struct {
    43  	items map[T][]Q
    44  }
    45  
    46  // Clear clears all entries in the map.
    47  func (mm *MultiMap[T, Q]) Clear() {
    48  	mm.items = map[T][]Q{}
    49  }
    50  
    51  // Add inserts the value into the map at the given key.
    52  //
    53  // If there exists an existing value, then this value is appended
    54  // *without comparison*. Put another way, a value can be added twice, if this
    55  // method is called twice for the same value.
    56  func (mm *MultiMap[T, Q]) Add(key T, item Q) {
    57  	if _, ok := mm.items[key]; !ok {
    58  		mm.items[key] = []Q{}
    59  	}
    60  
    61  	mm.items[key] = append(mm.items[key], item)
    62  }
    63  
    64  // RemoveKey removes the given key from the map.
    65  func (mm *MultiMap[T, Q]) RemoveKey(key T) {
    66  	delete(mm.items, key)
    67  }
    68  
    69  // Has returns true if the key is found in the map.
    70  func (mm *MultiMap[T, Q]) Has(key T) bool {
    71  	_, ok := mm.items[key]
    72  	return ok
    73  }
    74  
    75  // Get returns the values stored in the map for the provided key and whether
    76  // the key existed.
    77  //
    78  // If the key does not exist, an empty slice is returned.
    79  func (mm *MultiMap[T, Q]) Get(key T) ([]Q, bool) {
    80  	found, ok := mm.items[key]
    81  	if !ok {
    82  		return []Q{}, false
    83  	}
    84  
    85  	return found, true
    86  }
    87  
    88  // IsEmpty returns true if the map is currently empty.
    89  func (mm *MultiMap[T, Q]) IsEmpty() bool { return len(mm.items) == 0 }
    90  
    91  // Len returns the length of the map, e.g. the number of *keys* present.
    92  func (mm *MultiMap[T, Q]) Len() int { return len(mm.items) }
    93  
    94  // Keys returns the keys of the map.
    95  func (mm *MultiMap[T, Q]) Keys() []T { return maps.Keys(mm.items) }
    96  
    97  // Values returns all values in the map.
    98  func (mm MultiMap[T, Q]) Values() []Q {
    99  	values := make([]Q, 0, len(mm.items)*2)
   100  	for _, valueSlice := range maps.Values(mm.items) {
   101  		values = append(values, valueSlice...)
   102  	}
   103  	return values
   104  }
   105  
   106  // AsReadOnly returns a read-only *copy* of the mulitmap.
   107  func (mm *MultiMap[T, Q]) AsReadOnly() ReadOnlyMultimap[T, Q] {
   108  	return readOnlyMultimap[T, Q]{
   109  		maps.Clone(mm.items),
   110  	}
   111  }
   112  
   113  type readOnlyMultimap[T comparable, Q any] struct {
   114  	items map[T][]Q
   115  }
   116  
   117  // Has returns true if the key is found in the map.
   118  func (mm readOnlyMultimap[T, Q]) Has(key T) bool {
   119  	_, ok := mm.items[key]
   120  	return ok
   121  }
   122  
   123  // Get returns the values for the given key in the map and whether the key existed. If the key
   124  // does not exist, an empty slice is returned.
   125  func (mm readOnlyMultimap[T, Q]) Get(key T) ([]Q, bool) {
   126  	found, ok := mm.items[key]
   127  	if !ok {
   128  		return []Q{}, false
   129  	}
   130  
   131  	return found, true
   132  }
   133  
   134  // IsEmpty returns true if the map is currently empty.
   135  func (mm readOnlyMultimap[T, Q]) IsEmpty() bool { return len(mm.items) == 0 }
   136  
   137  // Len returns the length of the map, e.g. the number of *keys* present.
   138  func (mm readOnlyMultimap[T, Q]) Len() int { return len(mm.items) }
   139  
   140  // Keys returns the keys of the map.
   141  func (mm readOnlyMultimap[T, Q]) Keys() []T { return maps.Keys(mm.items) }
   142  
   143  // Values returns all values in the map.
   144  func (mm readOnlyMultimap[T, Q]) Values() []Q {
   145  	values := make([]Q, 0, len(mm.items)*2)
   146  	for _, valueSlice := range maps.Values(mm.items) {
   147  		values = append(values, valueSlice...)
   148  	}
   149  	return values
   150  }