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
|
// Copyright 2020 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// Package orderedmap provides an ordered map, implemented as a binary tree.
package orderedmap
// FIXME: This should probably be container/orderedmap.
import (
"context"
"chans"
"constraints"
)
// Map is an ordered map.
type Map[K, V any] struct {
root *node[K, V]
compare func(K, K) int
}
// node is the type of a node in the binary tree.
type node[K, V any] struct {
key K
val V
left, right *node[K, V]
}
// New returns a new map. It takes a comparison function that compares two
// keys and returns < 0 if the first is less, == 0 if they are equal,
// > 0 if the first is greater.
func New[K, V any](compare func(K, K) int) *Map[K, V] {
return &Map[K, V]{compare: compare}
}
// NewOrdered returns a new map whose key is an ordered type.
// This is like New, but does not require providing a compare function.
// The map compare function uses the obvious key ordering.
func NewOrdered[K constraints.Ordered, V any]() *Map[K, V] {
return New[K, V](func(k1, k2 K) int {
switch {
case k1 < k2:
return -1
case k1 == k2:
return 0
default:
return 1
}
})
}
// find looks up key in the map, returning either a pointer to the slot of the
// node holding key, or a pointer to the slot where should a node would go.
func (m *Map[K, V]) find(key K) **node[K, V] {
pn := &m.root
for *pn != nil {
switch cmp := m.compare(key, (*pn).key); {
case cmp < 0:
pn = &(*pn).left
case cmp > 0:
pn = &(*pn).right
default:
return pn
}
}
return pn
}
// Insert inserts a new key/value into the map.
// If the key is already present, the value is replaced.
// Reports whether this is a new key.
func (m *Map[K, V]) Insert(key K, val V) bool {
pn := m.find(key)
if *pn != nil {
(*pn).val = val
return false
}
*pn = &node[K, V]{key: key, val: val}
return true
}
// Find returns the value associated with a key, or the zero value
// if not present. The found result reports whether the key was found.
func (m *Map[K, V]) Find(key K) (V, bool) {
pn := m.find(key)
if *pn == nil {
var zero V
return zero, false
}
return (*pn).val, true
}
// keyValue is a pair of key and value used while iterating.
type keyValue[K, V any] struct {
key K
val V
}
// iterate returns an iterator that traverses the map.
func (m *Map[K, V]) Iterate() *Iterator[K, V] {
sender, receiver := chans.Ranger(keyValue[K, V])()
var f func(*node[K, V]) bool
f = func(n *node[K, V]) bool {
if n == nil {
return true
}
// Stop the traversal if Send fails, which means that
// nothing is listening to the receiver.
return f(n.left) &&
sender.Send(context.Background(), keyValue[K, V]{n.key, n.val}) &&
f(n.right)
}
go func() {
f(m.root)
sender.Close()
}()
return &Iterator[K, V]{receiver}
}
// Iterator is used to iterate over the map.
type Iterator[K, V any] struct {
r *chans.Receiver[keyValue[K, V]]
}
// Next returns the next key and value pair, and a boolean that reports
// whether they are valid. If not valid, we have reached the end of the map.
func (it *Iterator[K, V]) Next() (K, V, bool) {
keyval, ok := it.r.Next(context.Background())
if !ok {
var zerok K
var zerov V
return zerok, zerov, false
}
return keyval.key, keyval.val, true
}
|
