map2.go

     1  // Copyright 2019 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  // This file is like map.go, but instead of importing chans, it contains
     6  // the necessary functionality at the end of the file.
     7  
     8  // Package orderedmap provides an ordered map, implemented as a binary tree.
     9  package orderedmap
    10  
    11  // Map is an ordered map.
    12  type Map[K, V any] struct {
    13  	root    *node[K, V]
    14  	compare func(K, K) int
    15  }
    16  
    17  // node is the type of a node in the binary tree.
    18  type node[K, V any] struct {
    19  	key         K
    20  	val         V
    21  	left, right *node[K, V]
    22  }
    23  
    24  // New returns a new map.
    25  func New[K, V any](compare func(K, K) int) *Map[K, V] {
    26  	return &Map[K, V]{compare: compare}
    27  }
    28  
    29  // find looks up key in the map, and returns either a pointer
    30  // to the node holding key, or a pointer to the location where
    31  // such a node would go.
    32  func (m *Map[K, V]) find(key K) **node[K, V] {
    33  	pn := &m.root
    34  	for *pn != nil {
    35  		switch cmp := m.compare(key, (*pn).key); {
    36  		case cmp < 0:
    37  			pn = &(*pn).left
    38  		case cmp > 0:
    39  			pn = &(*pn).right
    40  		default:
    41  			return pn
    42  		}
    43  	}
    44  	return pn
    45  }
    46  
    47  // Insert inserts a new key/value into the map.
    48  // If the key is already present, the value is replaced.
    49  // Returns true if this is a new key, false if already present.
    50  func (m *Map[K, V]) Insert(key K, val V) bool {
    51  	pn := m.find(key)
    52  	if *pn != nil {
    53  		(*pn).val = val
    54  		return false
    55  	}
    56  	*pn = &node[K, V]{key: key, val: val}
    57  	return true
    58  }
    59  
    60  // Find returns the value associated with a key, or zero if not present.
    61  // The found result reports whether the key was found.
    62  func (m *Map[K, V]) Find(key K) (V, bool) {
    63  	pn := m.find(key)
    64  	if *pn == nil {
    65  		var zero V // see the discussion of zero values, above
    66  		return zero, false
    67  	}
    68  	return (*pn).val, true
    69  }
    70  
    71  // keyValue is a pair of key and value used when iterating.
    72  type keyValue[K, V any] struct {
    73  	key K
    74  	val V
    75  }
    76  
    77  // InOrder returns an iterator that does an in-order traversal of the map.
    78  func (m *Map[K, V]) InOrder() *Iterator[K, V] {
    79  	sender, receiver := chans_Ranger[keyValue[K, V]]()
    80  	var f func(*node[K, V]) bool
    81  	f = func(n *node[K, V]) bool {
    82  		if n == nil {
    83  			return true
    84  		}
    85  		// Stop sending values if sender.Send returns false,
    86  		// meaning that nothing is listening at the receiver end.
    87  		return f(n.left) &&
    88  			sender.Send(keyValue[K, V]{n.key, n.val}) &&
    89  			f(n.right)
    90  	}
    91  	go func() {
    92  		f(m.root)
    93  		sender.Close()
    94  	}()
    95  	return &Iterator[K, V]{receiver}
    96  }
    97  
    98  // Iterator is used to iterate over the map.
    99  type Iterator[K, V any] struct {
   100  	r *chans_Receiver[keyValue[K, V]]
   101  }
   102  
   103  // Next returns the next key and value pair, and a boolean indicating
   104  // whether they are valid or whether we have reached the end.
   105  func (it *Iterator[K, V]) Next() (K, V, bool) {
   106  	keyval, ok := it.r.Next()
   107  	if !ok {
   108  		var zerok K
   109  		var zerov V
   110  		return zerok, zerov, false
   111  	}
   112  	return keyval.key, keyval.val, true
   113  }
   114  
   115  // chans
   116  
   117  func chans_Ranger[T any]() (*chans_Sender[T], *chans_Receiver[T])
   118  
   119  // A sender is used to send values to a Receiver.
   120  type chans_Sender[T any] struct {
   121  	values chan<- T
   122  	done   <-chan bool
   123  }
   124  
   125  func (s *chans_Sender[T]) Send(v T) bool {
   126  	select {
   127  	case s.values <- v:
   128  		return true
   129  	case <-s.done:
   130  		return false
   131  	}
   132  }
   133  
   134  func (s *chans_Sender[T]) Close() {
   135  	close(s.values)
   136  }
   137  
   138  type chans_Receiver[T any] struct {
   139  	values <-chan T
   140  	done   chan<- bool
   141  }
   142  
   143  func (r *chans_Receiver[T]) Next() (T, bool) {
   144  	v, ok := <-r.values
   145  	return v, ok
   146  }
   147
View as plain text