# Source file test/typeparam/orderedmapsimp.dir/a.go

```     1  // Copyright 2021 The Go Authors. All rights reserved.
2  // Use of this source code is governed by a BSD-style
4
5  package a
6
7  import (
8  	"context"
9  	"runtime"
10  )
11
12  type Ordered interface {
13  	~int | ~int8 | ~int16 | ~int32 | ~int64 |
14  		~uint | ~uint8 | ~uint16 | ~uint32 | ~uint64 | ~uintptr |
15  		~float32 | ~float64 |
16  		~string
17  }
18
19  // Map is an ordered map.
20  type Map[K, V any] struct {
21  	root    *node[K, V]
22  	compare func(K, K) int
23  }
24
25  // node is the type of a node in the binary tree.
26  type node[K, V any] struct {
27  	key         K
28  	val         V
29  	left, right *node[K, V]
30  }
31
32  // New returns a new map. It takes a comparison function that compares two
33  // keys and returns < 0 if the first is less, == 0 if they are equal,
34  // > 0 if the first is greater.
35  func New[K, V any](compare func(K, K) int) *Map[K, V] {
36  	return &Map[K, V]{compare: compare}
37  }
38
39  // NewOrdered returns a new map whose key is an ordered type.
40  // This is like New, but does not require providing a compare function.
41  // The map compare function uses the obvious key ordering.
42  func NewOrdered[K Ordered, V any]() *Map[K, V] {
43  	return New[K, V](func(k1, k2 K) int {
44  		switch {
45  		case k1 < k2:
46  			return -1
47  		case k1 > k2:
48  			return 1
49  		default:
50  			return 0
51  		}
52  	})
53  }
54
55  // find looks up key in the map, returning either a pointer to the slot of the
56  // node holding key, or a pointer to the slot where a node would go.
57  func (m *Map[K, V]) find(key K) **node[K, V] {
58  	pn := &m.root
59  	for *pn != nil {
60  		switch cmp := m.compare(key, (*pn).key); {
61  		case cmp < 0:
62  			pn = &(*pn).left
63  		case cmp > 0:
64  			pn = &(*pn).right
65  		default:
66  			return pn
67  		}
68  	}
69  	return pn
70  }
71
72  // Insert inserts a new key/value into the map.
73  // If the key is already present, the value is replaced.
74  // Reports whether this is a new key.
75  func (m *Map[K, V]) Insert(key K, val V) bool {
76  	pn := m.find(key)
77  	if *pn != nil {
78  		(*pn).val = val
79  		return false
80  	}
81  	*pn = &node[K, V]{key: key, val: val}
82  	return true
83  }
84
85  // Find returns the value associated with a key, or the zero value
86  // if not present. The second result reports whether the key was found.
87  func (m *Map[K, V]) Find(key K) (V, bool) {
88  	pn := m.find(key)
89  	if *pn == nil {
90  		var zero V
91  		return zero, false
92  	}
93  	return (*pn).val, true
94  }
95
96  // keyValue is a pair of key and value used while iterating.
97  type keyValue[K, V any] struct {
98  	key K
99  	val V
100  }
101
102  // iterate returns an iterator that traverses the map.
103  func (m *Map[K, V]) Iterate() *Iterator[K, V] {
104  	sender, receiver := Ranger[keyValue[K, V]]()
105  	var f func(*node[K, V]) bool
106  	f = func(n *node[K, V]) bool {
107  		if n == nil {
108  			return true
109  		}
110  		// Stop the traversal if Send fails, which means that
111  		// nothing is listening to the receiver.
112  		return f(n.left) &&
113  			sender.Send(context.Background(), keyValue[K, V]{n.key, n.val}) &&
114  			f(n.right)
115  	}
116  	go func() {
117  		f(m.root)
118  		sender.Close()
119  	}()
121  }
122
123  // Iterator is used to iterate over the map.
124  type Iterator[K, V any] struct {
126  }
127
128  // Next returns the next key and value pair, and a boolean that reports
129  // whether they are valid. If not valid, we have reached the end of the map.
130  func (it *Iterator[K, V]) Next() (K, V, bool) {
131  	keyval, ok := it.r.Next(context.Background())
132  	if !ok {
133  		var zerok K
134  		var zerov V
135  		return zerok, zerov, false
136  	}
137  	return keyval.key, keyval.val, true
138  }
139
140  // Equal reports whether two slices are equal: the same length and all
141  // elements equal. All floating point NaNs are considered equal.
142  func SliceEqual[Elem comparable](s1, s2 []Elem) bool {
143  	if len(s1) != len(s2) {
144  		return false
145  	}
146  	for i, v1 := range s1 {
147  		v2 := s2[i]
148  		if v1 != v2 {
149  			isNaN := func(f Elem) bool { return f != f }
150  			if !isNaN(v1) || !isNaN(v2) {
151  				return false
152  			}
153  		}
154  	}
155  	return true
156  }
157
158  // Ranger returns a Sender and a Receiver. The Receiver provides a
159  // Next method to retrieve values. The Sender provides a Send method
160  // to send values and a Close method to stop sending values. The Next
161  // method indicates when the Sender has been closed, and the Send
162  // method indicates when the Receiver has been freed.
163  //
164  // This is a convenient way to exit a goroutine sending values when
166  func Ranger[Elem any]() (*Sender[Elem], *Receiver[Elem]) {
167  	c := make(chan Elem)
168  	d := make(chan struct{})
169  	s := &Sender[Elem]{
170  		values: c,
171  		done:   d,
172  	}
174  		values: c,
175  		done:   d,
176  	}
178  	return s, r
179  }
180
181  // A Sender is used to send values to a Receiver.
182  type Sender[Elem any] struct {
183  	values chan<- Elem
184  	done   <-chan struct{}
185  }
186
187  // Send sends a value to the receiver. It reports whether the value was sent.
188  // The value will not be sent if the context is closed or the receiver
189  // is freed.
190  func (s *Sender[Elem]) Send(ctx context.Context, v Elem) bool {
191  	select {
192  	case <-ctx.Done():
193  		return false
194  	case s.values <- v:
195  		return true
196  	case <-s.done:
197  		return false
198  	}
199  }
200
201  // Close tells the receiver that no more values will arrive.
202  // After Close is called, the Sender may no longer be used.
203  func (s *Sender[Elem]) Close() {
204  	close(s.values)
205  }
206
208  type Receiver[Elem any] struct {
209  	values <-chan Elem
210  	done   chan<- struct{}
211  }
212
213  // Next returns the next value from the channel. The bool result indicates
214  // whether the value is valid.
215  func (r *Receiver[Elem]) Next(ctx context.Context) (v Elem, ok bool) {
216  	select {
217  	case <-ctx.Done():
218  	case v, ok = <-r.values:
219  	}
220  	return v, ok
221  }
222
223  // finalize is a finalizer for the receiver.
224  func (r *Receiver[Elem]) finalize() {
225  	close(r.done)
226  }
227
```

View as plain text