-
Notifications
You must be signed in to change notification settings - Fork 649
/
bimap.go
102 lines (89 loc) · 2.39 KB
/
bimap.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
// Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.
package bimap
import "github.com/ava-labs/avalanchego/utils"
type Entry[K, V any] struct {
Key K
Value V
}
// BiMap is a bi-directional map.
type BiMap[K, V comparable] struct {
keyToValue map[K]V
valueToKey map[V]K
}
// New creates a new empty bimap.
func New[K, V comparable]() *BiMap[K, V] {
return &BiMap[K, V]{
keyToValue: make(map[K]V),
valueToKey: make(map[V]K),
}
}
// Put the key value pair into the map. If either [key] or [val] was previously
// in the map, the previous entries will be removed and returned.
//
// Note: Unlike normal maps, it's possible that Put removes 0, 1, or 2 existing
// entries to ensure that mappings are one-to-one.
func (m *BiMap[K, V]) Put(key K, val V) []Entry[K, V] {
var removed []Entry[K, V]
oldVal, oldValDeleted := m.DeleteKey(key)
if oldValDeleted {
removed = append(removed, Entry[K, V]{
Key: key,
Value: oldVal,
})
}
oldKey, oldKeyDeleted := m.DeleteValue(val)
if oldKeyDeleted {
removed = append(removed, Entry[K, V]{
Key: oldKey,
Value: val,
})
}
m.keyToValue[key] = val
m.valueToKey[val] = key
return removed
}
// GetKey that maps to the provided value.
func (m *BiMap[K, V]) GetKey(val V) (K, bool) {
key, ok := m.valueToKey[val]
return key, ok
}
// GetValue that is mapped to the provided key.
func (m *BiMap[K, V]) GetValue(key K) (V, bool) {
val, ok := m.keyToValue[key]
return val, ok
}
// HasKey returns true if [key] is in the map.
func (m *BiMap[K, _]) HasKey(key K) bool {
_, ok := m.keyToValue[key]
return ok
}
// HasValue returns true if [val] is in the map.
func (m *BiMap[_, V]) HasValue(val V) bool {
_, ok := m.valueToKey[val]
return ok
}
// DeleteKey removes [key] from the map and returns the value it mapped to.
func (m *BiMap[K, V]) DeleteKey(key K) (V, bool) {
val, ok := m.keyToValue[key]
if !ok {
return utils.Zero[V](), false
}
delete(m.keyToValue, key)
delete(m.valueToKey, val)
return val, true
}
// DeleteValue removes [val] from the map and returns the key that mapped to it.
func (m *BiMap[K, V]) DeleteValue(val V) (K, bool) {
key, ok := m.valueToKey[val]
if !ok {
return utils.Zero[K](), false
}
delete(m.keyToValue, key)
delete(m.valueToKey, val)
return key, true
}
// Len return the number of entries in this map.
func (m *BiMap[K, V]) Len() int {
return len(m.keyToValue)
}