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
|
// Copyright 2019 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
// TODO(gri) fix imports for tests
import "chans" // ERROR could not import
// 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.
func New[K, V any](compare func(K, K) int) *Map[K, V] {
return &Map[K, V]{compare: compare}
}
// find looks up key in the map, and returns either a pointer
// to the node holding key, or a pointer to the location where
// such 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.
// Returns true if this is a new key, false if already present.
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 zero 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 // see the discussion of zero values, above
return zero, false
}
return (*pn).val, true
}
// keyValue is a pair of key and value used when iterating.
type keyValue[K, V any] struct {
key K
val V
}
// InOrder returns an iterator that does an in-order traversal of the map.
func (m *Map[K, V]) InOrder() *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 sending values if sender.Send returns false,
// meaning that nothing is listening at the receiver end.
return f(n.left) &&
sender.Send(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 indicating
// whether they are valid or whether we have reached the end.
func (it *Iterator[K, V]) Next() (K, V, bool) {
keyval, ok := it.r.Next()
if !ok {
var zerok K
var zerov V
return zerok, zerov, false
}
return keyval.key, keyval.val, true
}
|
