# Source file test/typeparam/orderedmap.go

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

View as plain text