mirror of https://github.com/golang/go.git
114 lines
2.8 KiB
Plaintext
114 lines
2.8 KiB
Plaintext
// Copyright 2019 The Go Authors. All rights reserved.
|
|
// Use of this source code is governed by a BSD-style
|
|
// license that can be found in the LICENSE file.
|
|
|
|
// Package orderedmap provides an ordered map, implemented as a binary tree.
|
|
package orderedmap
|
|
|
|
// TODO(gri) fix imports for tests
|
|
import "chans" // ERROR could not import
|
|
|
|
// Map is an ordered map.
|
|
type Map[K, V any] struct {
|
|
root *node[K, V]
|
|
compare func(K, K) int
|
|
}
|
|
|
|
// node is the type of a node in the binary tree.
|
|
type node[K, V any] struct {
|
|
key K
|
|
val V
|
|
left, right *node[K, V]
|
|
}
|
|
|
|
// New returns a new map.
|
|
func New[K, V any](compare func(K, K) int) *Map[K, V] {
|
|
return &Map[K, V]{compare: compare}
|
|
}
|
|
|
|
// find looks up key in the map, and returns either a pointer
|
|
// to the node holding key, or a pointer to the location where
|
|
// such a node would go.
|
|
func (m *Map[K, V]) find(key K) **node[K, V] {
|
|
pn := &m.root
|
|
for *pn != nil {
|
|
switch cmp := m.compare(key, (*pn).key); {
|
|
case cmp < 0:
|
|
pn = &(*pn).left
|
|
case cmp > 0:
|
|
pn = &(*pn).right
|
|
default:
|
|
return pn
|
|
}
|
|
}
|
|
return pn
|
|
}
|
|
|
|
// Insert inserts a new key/value into the map.
|
|
// If the key is already present, the value is replaced.
|
|
// Returns true if this is a new key, false if already present.
|
|
func (m *Map[K, V]) Insert(key K, val V) bool {
|
|
pn := m.find(key)
|
|
if *pn != nil {
|
|
(*pn).val = val
|
|
return false
|
|
}
|
|
*pn = &node[K, V]{key: key, val: val}
|
|
return true
|
|
}
|
|
|
|
// Find returns the value associated with a key, or zero if not present.
|
|
// The found result reports whether the key was found.
|
|
func (m *Map[K, V]) Find(key K) (V, bool) {
|
|
pn := m.find(key)
|
|
if *pn == nil {
|
|
var zero V // see the discussion of zero values, above
|
|
return zero, false
|
|
}
|
|
return (*pn).val, true
|
|
}
|
|
|
|
// keyValue is a pair of key and value used when iterating.
|
|
type keyValue[K, V any] struct {
|
|
key K
|
|
val V
|
|
}
|
|
|
|
// InOrder returns an iterator that does an in-order traversal of the map.
|
|
func (m *Map[K, V]) InOrder() *Iterator[K, V] {
|
|
sender, receiver := chans.Ranger[keyValue[K, V]]()
|
|
var f func(*node[K, V]) bool
|
|
f = func(n *node[K, V]) bool {
|
|
if n == nil {
|
|
return true
|
|
}
|
|
// Stop sending values if sender.Send returns false,
|
|
// meaning that nothing is listening at the receiver end.
|
|
return f(n.left) &&
|
|
sender.Send(keyValue[K, V]{n.key, n.val}) &&
|
|
f(n.right)
|
|
}
|
|
go func() {
|
|
f(m.root)
|
|
sender.Close()
|
|
}()
|
|
return &Iterator[K, V]{receiver}
|
|
}
|
|
|
|
// Iterator is used to iterate over the map.
|
|
type Iterator[K, V any] struct {
|
|
r *chans.Receiver[keyValue[K, V]]
|
|
}
|
|
|
|
// Next returns the next key and value pair, and a boolean indicating
|
|
// whether they are valid or whether we have reached the end.
|
|
func (it *Iterator[K, V]) Next() (K, V, bool) {
|
|
keyval, ok := it.r.Next()
|
|
if !ok {
|
|
var zerok K
|
|
var zerov V
|
|
return zerok, zerov, false
|
|
}
|
|
return keyval.key, keyval.val, true
|
|
}
|