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