blob: 5293b0e445fbbc823a77d31f1bbd929ee0d07c6d [file] [log] [blame]
 // Copyright 2020 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 // FIXME: This should probably be container/orderedmap. import ( "context" "chans" "constraints" ) // 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. It takes a comparison function that compares two // keys and returns < 0 if the first is less, == 0 if they are equal, // > 0 if the first is greater. func New[K, V any](compare func(K, K) int) *Map[K, V] { return &Map[K, V]{compare: compare} } // NewOrdered returns a new map whose key is an ordered type. // This is like New, but does not require providing a compare function. // The map compare function uses the obvious key ordering. func NewOrdered[K constraints.Ordered, V any]() *Map[K, V] { return New[K, V](func(k1, k2 K) int { switch { case k1 < k2: return -1 case k1 == k2: return 0 default: return 1 } }) } // find looks up key in the map, returning either a pointer to the slot of the // node holding key, or a pointer to the slot where should 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. // Reports whether this is a new key. 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 the zero value // 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 return zero, false } return (*pn).val, true } // keyValue is a pair of key and value used while iterating. type keyValue[K, V any] struct { key K val V } // iterate returns an iterator that traverses the map. func (m *Map[K, V]) Iterate() *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 the traversal if Send fails, which means that // nothing is listening to the receiver. return f(n.left) && sender.Send(context.Background(), 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 that reports // whether they are valid. If not valid, we have reached the end of the map. func (it *Iterator[K, V]) Next() (K, V, bool) { keyval, ok := it.r.Next(context.Background()) if !ok { var zerok K var zerov V return zerok, zerov, false } return keyval.key, keyval.val, true }