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