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. 6package orderedmap 7 8// TODO(gri) fix imports for tests 9import "chans" // ERROR could not import 10 11// Map is an ordered map. 12type 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. 18type 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. 25func 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. 32func (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. 50func (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. 62func (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. 72type 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. 78func (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. 99type 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. 105func (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