diff options
Diffstat (limited to 'test/typeparam/orderedmap.go')
-rw-r--r-- | test/typeparam/orderedmap.go | 286 |
1 files changed, 286 insertions, 0 deletions
diff --git a/test/typeparam/orderedmap.go b/test/typeparam/orderedmap.go new file mode 100644 index 0000000..1245669 --- /dev/null +++ b/test/typeparam/orderedmap.go @@ -0,0 +1,286 @@ +// run + +// Copyright 2021 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +// Package orderedmap provides an ordered map, implemented as a binary tree. +package main + +import ( + "bytes" + "context" + "fmt" + "runtime" +) + +type Ordered interface { + ~int | ~int8 | ~int16 | ~int32 | ~int64 | + ~uint | ~uint8 | ~uint16 | ~uint32 | ~uint64 | ~uintptr | + ~float32 | ~float64 | + ~string +} + +// _Map is an ordered map. +type _Map[K, V any] struct { + root *node[K, V] + compare func(K, K) int +} + +// node is the type of a node in the binary tree. +type node[K, V any] struct { + key K + val V + left, right *node[K, V] +} + +// _New returns a new map. It takes a comparison function that compares two +// keys and returns < 0 if the first is less, == 0 if they are equal, +// > 0 if the first is greater. +func _New[K, V any](compare func(K, K) int) *_Map[K, V] { + return &_Map[K, V]{compare: compare} +} + +// _NewOrdered returns a new map whose key is an ordered type. +// This is like _New, but does not require providing a compare function. +// The map compare function uses the obvious key ordering. +func _NewOrdered[K Ordered, V any]() *_Map[K, V] { + return _New[K, V](func(k1, k2 K) int { + switch { + case k1 < k2: + return -1 + case k1 == k2: + return 0 + default: + return 1 + } + }) +} + +// find looks up key in the map, returning either a pointer to the slot of the +// node holding key, or a pointer to the slot where should a node would go. +func (m *_Map[K, V]) find(key K) **node[K, V] { + pn := &m.root + for *pn != nil { + switch cmp := m.compare(key, (*pn).key); { + case cmp < 0: + pn = &(*pn).left + case cmp > 0: + pn = &(*pn).right + default: + return pn + } + } + return pn +} + +// Insert inserts a new key/value into the map. +// If the key is already present, the value is replaced. +// Reports whether this is a new key. +func (m *_Map[K, V]) Insert(key K, val V) bool { + pn := m.find(key) + if *pn != nil { + (*pn).val = val + return false + } + *pn = &node[K, V]{key: key, val: val} + return true +} + +// Find returns the value associated with a key, or the zero value +// if not present. The found result reports whether the key was found. +func (m *_Map[K, V]) Find(key K) (V, bool) { + pn := m.find(key) + if *pn == nil { + var zero V + return zero, false + } + return (*pn).val, true +} + +// keyValue is a pair of key and value used while iterating. +type keyValue[K, V any] struct { + key K + val V +} + +// iterate returns an iterator that traverses the map. +func (m *_Map[K, V]) Iterate() *_Iterator[K, V] { + sender, receiver := _Ranger[keyValue[K, V]]() + var f func(*node[K, V]) bool + f = func(n *node[K, V]) bool { + if n == nil { + return true + } + // Stop the traversal if Send fails, which means that + // nothing is listening to the receiver. + return f(n.left) && + sender.Send(context.Background(), keyValue[K, V]{n.key, n.val}) && + f(n.right) + } + go func() { + f(m.root) + sender.Close() + }() + return &_Iterator[K, V]{receiver} +} + +// _Iterator is used to iterate over the map. +type _Iterator[K, V any] struct { + r *_Receiver[keyValue[K, V]] +} + +// Next returns the next key and value pair, and a boolean that reports +// whether they are valid. If not valid, we have reached the end of the map. +func (it *_Iterator[K, V]) Next() (K, V, bool) { + keyval, ok := it.r.Next(context.Background()) + if !ok { + var zerok K + var zerov V + return zerok, zerov, false + } + return keyval.key, keyval.val, true +} + +func TestMap() { + m := _New[[]byte, int](bytes.Compare) + + if _, found := m.Find([]byte("a")); found { + panic(fmt.Sprintf("unexpectedly found %q in empty map", []byte("a"))) + } + if !m.Insert([]byte("a"), 'a') { + panic(fmt.Sprintf("key %q unexpectedly already present", []byte("a"))) + } + if !m.Insert([]byte("c"), 'c') { + panic(fmt.Sprintf("key %q unexpectedly already present", []byte("c"))) + } + if !m.Insert([]byte("b"), 'b') { + panic(fmt.Sprintf("key %q unexpectedly already present", []byte("b"))) + } + if m.Insert([]byte("c"), 'x') { + panic(fmt.Sprintf("key %q unexpectedly not present", []byte("c"))) + } + + if v, found := m.Find([]byte("a")); !found { + panic(fmt.Sprintf("did not find %q", []byte("a"))) + } else if v != 'a' { + panic(fmt.Sprintf("key %q returned wrong value %c, expected %c", []byte("a"), v, 'a')) + } + if v, found := m.Find([]byte("c")); !found { + panic(fmt.Sprintf("did not find %q", []byte("c"))) + } else if v != 'x' { + panic(fmt.Sprintf("key %q returned wrong value %c, expected %c", []byte("c"), v, 'x')) + } + + if _, found := m.Find([]byte("d")); found { + panic(fmt.Sprintf("unexpectedly found %q", []byte("d"))) + } + + gather := func(it *_Iterator[[]byte, int]) []int { + var r []int + for { + _, v, ok := it.Next() + if !ok { + return r + } + r = append(r, v) + } + } + got := gather(m.Iterate()) + want := []int{'a', 'b', 'x'} + if !_SliceEqual(got, want) { + panic(fmt.Sprintf("Iterate returned %v, want %v", got, want)) + } +} + +func main() { + TestMap() +} + +// _Equal reports whether two slices are equal: the same length and all +// elements equal. All floating point NaNs are considered equal. +func _SliceEqual[Elem comparable](s1, s2 []Elem) bool { + if len(s1) != len(s2) { + return false + } + for i, v1 := range s1 { + v2 := s2[i] + if v1 != v2 { + isNaN := func(f Elem) bool { return f != f } + if !isNaN(v1) || !isNaN(v2) { + return false + } + } + } + return true +} + +// Ranger returns a Sender and a Receiver. The Receiver provides a +// Next method to retrieve values. The Sender provides a Send method +// to send values and a Close method to stop sending values. The Next +// method indicates when the Sender has been closed, and the Send +// method indicates when the Receiver has been freed. +// +// This is a convenient way to exit a goroutine sending values when +// the receiver stops reading them. +func _Ranger[Elem any]() (*_Sender[Elem], *_Receiver[Elem]) { + c := make(chan Elem) + d := make(chan struct{}) + s := &_Sender[Elem]{ + values: c, + done: d, + } + r := &_Receiver[Elem]{ + values: c, + done: d, + } + runtime.SetFinalizer(r, (*_Receiver[Elem]).finalize) + return s, r +} + +// A _Sender is used to send values to a Receiver. +type _Sender[Elem any] struct { + values chan<- Elem + done <-chan struct{} +} + +// Send sends a value to the receiver. It reports whether the value was sent. +// The value will not be sent if the context is closed or the receiver +// is freed. +func (s *_Sender[Elem]) Send(ctx context.Context, v Elem) bool { + select { + case <-ctx.Done(): + return false + case s.values <- v: + return true + case <-s.done: + return false + } +} + +// Close tells the receiver that no more values will arrive. +// After Close is called, the _Sender may no longer be used. +func (s *_Sender[Elem]) Close() { + close(s.values) +} + +// A _Receiver receives values from a _Sender. +type _Receiver[Elem any] struct { + values <-chan Elem + done chan<- struct{} +} + +// Next returns the next value from the channel. The bool result indicates +// whether the value is valid. +func (r *_Receiver[Elem]) Next(ctx context.Context) (v Elem, ok bool) { + select { + case <-ctx.Done(): + case v, ok = <-r.values: + } + return v, ok +} + +// finalize is a finalizer for the receiver. +func (r *_Receiver[Elem]) finalize() { + close(r.done) +} |