test/typeparam/orderedmap.go

   1 // run
   2
   3 // Copyright 2021 The Go Authors. All rights reserved.
   4 // Use of this source code is governed by a BSD-style
   5 // license that can be found in the LICENSE file.
   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         }()
 125         return &_Iterator[K, V]{receiver}
 126 }
 127
 128 // _Iterator is used to iterate over the map.
 129 type _Iterator[K, V any] struct {
 130         r *_Receiver[keyValue[K, V]]
 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
 225 // the receiver stops reading them.
 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         }
 233         r := &_Receiver[Elem]{
 234                 values: c,
 235                 done:   d,
 236         }
 237         runtime.SetFinalizer(r, (*_Receiver[Elem]).finalize)
 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
 267 // A _Receiver receives values from a _Sender.
 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 }