a.go

     1  // Copyright 2021 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package a
     6  
     7  import (
     8  	"context"
     9  	"runtime"
    10  )
    11  
    12  type Ordered interface {
    13  	~int | ~int8 | ~int16 | ~int32 | ~int64 |
    14  		~uint | ~uint8 | ~uint16 | ~uint32 | ~uint64 | ~uintptr |
    15  		~float32 | ~float64 |
    16  		~string
    17  }
    18  
    19  // Map is an ordered map.
    20  type Map[K, V any] struct {
    21  	root    *node[K, V]
    22  	compare func(K, K) int
    23  }
    24  
    25  // node is the type of a node in the binary tree.
    26  type node[K, V any] struct {
    27  	key         K
    28  	val         V
    29  	left, right *node[K, V]
    30  }
    31  
    32  // New returns a new map. It takes a comparison function that compares two
    33  // keys and returns < 0 if the first is less, == 0 if they are equal,
    34  // > 0 if the first is greater.
    35  func New[K, V any](compare func(K, K) int) *Map[K, V] {
    36  	return &Map[K, V]{compare: compare}
    37  }
    38  
    39  // NewOrdered returns a new map whose key is an ordered type.
    40  // This is like New, but does not require providing a compare function.
    41  // The map compare function uses the obvious key ordering.
    42  func NewOrdered[K Ordered, V any]() *Map[K, V] {
    43  	return New[K, V](func(k1, k2 K) int {
    44  		switch {
    45  		case k1 < k2:
    46  			return -1
    47  		case k1 > k2:
    48  			return 1
    49  		default:
    50  			return 0
    51  		}
    52  	})
    53  }
    54  
    55  // find looks up key in the map, returning either a pointer to the slot of the
    56  // node holding key, or a pointer to the slot where a node would go.
    57  func (m *Map[K, V]) find(key K) **node[K, V] {
    58  	pn := &m.root
    59  	for *pn != nil {
    60  		switch cmp := m.compare(key, (*pn).key); {
    61  		case cmp < 0:
    62  			pn = &(*pn).left
    63  		case cmp > 0:
    64  			pn = &(*pn).right
    65  		default:
    66  			return pn
    67  		}
    68  	}
    69  	return pn
    70  }
    71  
    72  // Insert inserts a new key/value into the map.
    73  // If the key is already present, the value is replaced.
    74  // Reports whether this is a new key.
    75  func (m *Map[K, V]) Insert(key K, val V) bool {
    76  	pn := m.find(key)
    77  	if *pn != nil {
    78  		(*pn).val = val
    79  		return false
    80  	}
    81  	*pn = &node[K, V]{key: key, val: val}
    82  	return true
    83  }
    84  
    85  // Find returns the value associated with a key, or the zero value
    86  // if not present. The second result reports whether the key was found.
    87  func (m *Map[K, V]) Find(key K) (V, bool) {
    88  	pn := m.find(key)
    89  	if *pn == nil {
    90  		var zero V
    91  		return zero, false
    92  	}
    93  	return (*pn).val, true
    94  }
    95  
    96  // keyValue is a pair of key and value used while iterating.
    97  type keyValue[K, V any] struct {
    98  	key K
    99  	val V
   100  }
   101  
   102  // iterate returns an iterator that traverses the map.
   103  func (m *Map[K, V]) Iterate() *Iterator[K, V] {
   104  	sender, receiver := Ranger[keyValue[K, V]]()
   105  	var f func(*node[K, V]) bool
   106  	f = func(n *node[K, V]) bool {
   107  		if n == nil {
   108  			return true
   109  		}
   110  		// Stop the traversal if Send fails, which means that
   111  		// nothing is listening to the receiver.
   112  		return f(n.left) &&
   113  			sender.Send(context.Background(), keyValue[K, V]{n.key, n.val}) &&
   114  			f(n.right)
   115  	}
   116  	go func() {
   117  		f(m.root)
   118  		sender.Close()
   119  	}()
   120  	return &Iterator[K, V]{receiver}
   121  }
   122  
   123  // Iterator is used to iterate over the map.
   124  type Iterator[K, V any] struct {
   125  	r *Receiver[keyValue[K, V]]
   126  }
   127  
   128  // Next returns the next key and value pair, and a boolean that reports
   129  // whether they are valid. If not valid, we have reached the end of the map.
   130  func (it *Iterator[K, V]) Next() (K, V, bool) {
   131  	keyval, ok := it.r.Next(context.Background())
   132  	if !ok {
   133  		var zerok K
   134  		var zerov V
   135  		return zerok, zerov, false
   136  	}
   137  	return keyval.key, keyval.val, true
   138  }
   139  
   140  // Equal reports whether two slices are equal: the same length and all
   141  // elements equal. All floating point NaNs are considered equal.
   142  func SliceEqual[Elem comparable](s1, s2 []Elem) bool {
   143  	if len(s1) != len(s2) {
   144  		return false
   145  	}
   146  	for i, v1 := range s1 {
   147  		v2 := s2[i]
   148  		if v1 != v2 {
   149  			isNaN := func(f Elem) bool { return f != f }
   150  			if !isNaN(v1) || !isNaN(v2) {
   151  				return false
   152  			}
   153  		}
   154  	}
   155  	return true
   156  }
   157  
   158  // Ranger returns a Sender and a Receiver. The Receiver provides a
   159  // Next method to retrieve values. The Sender provides a Send method
   160  // to send values and a Close method to stop sending values. The Next
   161  // method indicates when the Sender has been closed, and the Send
   162  // method indicates when the Receiver has been freed.
   163  //
   164  // This is a convenient way to exit a goroutine sending values when
   165  // the receiver stops reading them.
   166  func Ranger[Elem any]() (*Sender[Elem], *Receiver[Elem]) {
   167  	c := make(chan Elem)
   168  	d := make(chan struct{})
   169  	s := &Sender[Elem]{
   170  		values: c,
   171  		done:   d,
   172  	}
   173  	r := &Receiver[Elem]{
   174  		values: c,
   175  		done:   d,
   176  	}
   177  	runtime.SetFinalizer(r, (*Receiver[Elem]).finalize)
   178  	return s, r
   179  }
   180  
   181  // A Sender is used to send values to a Receiver.
   182  type Sender[Elem any] struct {
   183  	values chan<- Elem
   184  	done   <-chan struct{}
   185  }
   186  
   187  // Send sends a value to the receiver. It reports whether the value was sent.
   188  // The value will not be sent if the context is closed or the receiver
   189  // is freed.
   190  func (s *Sender[Elem]) Send(ctx context.Context, v Elem) bool {
   191  	select {
   192  	case <-ctx.Done():
   193  		return false
   194  	case s.values <- v:
   195  		return true
   196  	case <-s.done:
   197  		return false
   198  	}
   199  }
   200  
   201  // Close tells the receiver that no more values will arrive.
   202  // After Close is called, the Sender may no longer be used.
   203  func (s *Sender[Elem]) Close() {
   204  	close(s.values)
   205  }
   206  
   207  // A Receiver receives values from a Sender.
   208  type Receiver[Elem any] struct {
   209  	values <-chan Elem
   210  	done   chan<- struct{}
   211  }
   212  
   213  // Next returns the next value from the channel. The bool result indicates
   214  // whether the value is valid.
   215  func (r *Receiver[Elem]) Next(ctx context.Context) (v Elem, ok bool) {
   216  	select {
   217  	case <-ctx.Done():
   218  	case v, ok = <-r.values:
   219  	}
   220  	return v, ok
   221  }
   222  
   223  // finalize is a finalizer for the receiver.
   224  func (r *Receiver[Elem]) finalize() {
   225  	close(r.done)
   226  }
   227
View as plain text