codekingpro/portable-devtools

1// run2 3// Copyright 2021 The Go Authors. All rights reserved.4// Use of this source code is governed by a BSD-style5// license that can be found in the LICENSE file.6 7// Package orderedmap provides an ordered map, implemented as a binary tree.8package main9 10import (11	"bytes"12	"context"13	"fmt"14	"runtime"15)16 17type Ordered interface {18	~int | ~int8 | ~int16 | ~int32 | ~int64 |19		~uint | ~uint8 | ~uint16 | ~uint32 | ~uint64 | ~uintptr |20		~float32 | ~float64 |21		~string22}23 24// _Map is an ordered map.25type _Map[K, V any] struct {26	root    *node[K, V]27	compare func(K, K) int28}29 30// node is the type of a node in the binary tree.31type node[K, V any] struct {32	key         K33	val         V34	left, right *node[K, V]35}36 37// _New returns a new map. It takes a comparison function that compares two38// keys and returns < 0 if the first is less, == 0 if they are equal,39// > 0 if the first is greater.40func _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.47func _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 -152		case k1 == k2:53			return 054		default:55			return 156		}57	})58}59 60// find looks up key in the map, returning either a pointer to the slot of the61// node holding key, or a pointer to the slot where should a node would go.62func (m *_Map[K, V]) find(key K) **node[K, V] {63	pn := &m.root64	for *pn != nil {65		switch cmp := m.compare(key, (*pn).key); {66		case cmp < 0:67			pn = &(*pn).left68		case cmp > 0:69			pn = &(*pn).right70		default:71			return pn72		}73	}74	return pn75}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.80func (m *_Map[K, V]) Insert(key K, val V) bool {81	pn := m.find(key)82	if *pn != nil {83		(*pn).val = val84		return false85	}86	*pn = &node[K, V]{key: key, val: val}87	return true88}89 90// Find returns the value associated with a key, or the zero value91// if not present. The found result reports whether the key was found.92func (m *_Map[K, V]) Find(key K) (V, bool) {93	pn := m.find(key)94	if *pn == nil {95		var zero V96		return zero, false97	}98	return (*pn).val, true99}100 101// keyValue is a pair of key and value used while iterating.102type keyValue[K, V any] struct {103	key K104	val V105}106 107// iterate returns an iterator that traverses the map.108func (m *_Map[K, V]) Iterate() *_Iterator[K, V] {109	sender, receiver := _Ranger[keyValue[K, V]]()110	var f func(*node[K, V]) bool111	f = func(n *node[K, V]) bool {112		if n == nil {113			return true114		}115		// Stop the traversal if Send fails, which means that116		// 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.129type _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 reports134// whether they are valid. If not valid, we have reached the end of the map.135func (it *_Iterator[K, V]) Next() (K, V, bool) {136	keyval, ok := it.r.Next(context.Background())137	if !ok {138		var zerok K139		var zerov V140		return zerok, zerov, false141	}142	return keyval.key, keyval.val, true143}144 145func 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 []int181		for {182			_, v, ok := it.Next()183			if !ok {184				return r185			}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 196func main() {197	TestMap()198}199 200// _Equal reports whether two slices are equal: the same length and all201// elements equal. All floating point NaNs are considered equal.202func _SliceEqual[Elem comparable](s1, s2 []Elem) bool {203	if len(s1) != len(s2) {204		return false205	}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 false212			}213		}214	}215	return true216}217 218// Ranger returns a Sender and a Receiver. The Receiver provides a219// Next method to retrieve values. The Sender provides a Send method220// to send values and a Close method to stop sending values. The Next221// method indicates when the Sender has been closed, and the Send222// method indicates when the Receiver has been freed.223//224// This is a convenient way to exit a goroutine sending values when225// the receiver stops reading them.226func _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, r239}240 241// A _Sender is used to send values to a Receiver.242type _Sender[Elem any] struct {243	values chan<- Elem244	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 receiver249// is freed.250func (s *_Sender[Elem]) Send(ctx context.Context, v Elem) bool {251	select {252	case <-ctx.Done():253		return false254	case s.values <- v:255		return true256	case <-s.done:257		return false258	}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.263func (s *_Sender[Elem]) Close() {264	close(s.values)265}266 267// A _Receiver receives values from a _Sender.268type _Receiver[Elem any] struct {269	values <-chan Elem270	done   chan<- struct{}271}272 273// Next returns the next value from the channel. The bool result indicates274// whether the value is valid.275func (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, ok281}282 283// finalize is a finalizer for the receiver.284func (r *_Receiver[Elem]) finalize() {285	close(r.done)286}287