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.
6package orderedmap
7
8// TODO(gri) fix imports for tests
9import "chans" // ERROR could not import
10
11// Map is an ordered map.
12type 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.
18type 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.
25func 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.
32func (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.
50func (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.
62func (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.
72type 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.
78func (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.
99type 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.
105func (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