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