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  }
   287
View as plain text