1// Package orderedmap provides an ordered map, implemented as a binary tree.
2package orderedmap
3
4import "chans"
5
6// Map is an ordered map.
7type 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.
13type 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.
20func 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.
27func (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.
45func (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.
57func (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.
67type 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.
73func (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.
95type 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.
101func (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}
View as plain text