|
| 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 | +// TODO(gri) fix imports for tests |
| 9 | +import "chans" // ERROR could not import |
| 10 | + |
| 11 | +// Map is an ordered map. |
| 12 | +type Map[type K, V] 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[type K, V] struct { |
| 19 | + key K |
| 20 | + val V |
| 21 | + left, right *node[K, V] |
| 22 | +} |
| 23 | + |
| 24 | +// New returns a new map. |
| 25 | +func New[type K, V](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[type K, V] 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[type K, V] 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 | +} |
0 commit comments