1 // Copyright 2019 The Go Authors. All rights reserved.
2 // Use of this source code is governed by a BSD-style
3 // license that can be found in the LICENSE file.
4
5 // Package orderedmap provides an ordered map, implemented as a binary tree.
6 package orderedmap
7
8 // TODO(gri) fix imports for tests
9 import "chans" // ERROR could not import
10
11 // Map is an ordered map.
12 type Map[K, V any] struct {
13 root *node[K, V]
14 compare func(K, K) int
15 }
16
17 // node is the type of a node in the binary tree.
18 type node[K, V any] struct {
19 key K
20 val V
21 left, right *node[K, V]
22 }
23
24 // New returns a new map.
25 func New[K, V any](compare func(K, K) int) *Map[K, V] {
26 return &Map[K, V]{compare: compare}
27 }
28
29 // find looks up key in the map, and returns either a pointer
30 // to the node holding key, or a pointer to the location where
31 // such a node would go.
32 func (m *Map[K, V]) find(key K) **node[K, V] {
33 pn := &m.root
34 for *pn != nil {
35 switch cmp := m.compare(key, (*pn).key); {
36 case cmp < 0:
37 pn = &(*pn).left
38 case cmp > 0:
39 pn = &(*pn).right
40 default:
41 return pn
42 }
43 }
44 return pn
45 }
46
47 // Insert inserts a new key/value into the map.
48 // If the key is already present, the value is replaced.
49 // Returns true if this is a new key, false if already present.
50 func (m *Map[K, V]) Insert(key K, val V) bool {
51 pn := m.find(key)
52 if *pn != nil {
53 (*pn).val = val
54 return false
55 }
56 *pn = &node[K, V]{key: key, val: val}
57 return true
58 }
59
60 // Find returns the value associated with a key, or zero if not present.
61 // The found result reports whether the key was found.
62 func (m *Map[K, V]) Find(key K) (V, bool) {
63 pn := m.find(key)
64 if *pn == nil {
65 var zero V // see the discussion of zero values, above
66 return zero, false
67 }
68 return (*pn).val, true
69 }
70
71 // keyValue is a pair of key and value used when iterating.
72 type keyValue[K, V any] struct {
73 key K
74 val V
75 }
76
77 // InOrder returns an iterator that does an in-order traversal of the map.
78 func (m *Map[K, V]) InOrder() *Iterator[K, V] {
79 sender, receiver := chans.Ranger[keyValue[K, V]]()
80 var f func(*node[K, V]) bool
81 f = func(n *node[K, V]) bool {
82 if n == nil {
83 return true
84 }
85 // Stop sending values if sender.Send returns false,
86 // meaning that nothing is listening at the receiver end.
87 return f(n.left) &&
88 sender.Send(keyValue[K, V]{n.key, n.val}) &&
89 f(n.right)
90 }
91 go func() {
92 f(m.root)
93 sender.Close()
94 }()
95 return &Iterator[K, V]{receiver}
96 }
97
98 // Iterator is used to iterate over the map.
99 type Iterator[K, V any] struct {
100 r *chans.Receiver[keyValue[K, V]]
101 }
102
103 // Next returns the next key and value pair, and a boolean indicating
104 // whether they are valid or whether we have reached the end.
105 func (it *Iterator[K, V]) Next() (K, V, bool) {
106 keyval, ok := it.r.Next()
107 if !ok {
108 var zerok K
109 var zerov V
110 return zerok, zerov, false
111 }
112 return keyval.key, keyval.val, true
113 }
114
View as plain text