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 (
// 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
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 :=, (*pn).key); {
case cmp < 0:
pn = &(*pn).left
case cmp > 0:
pn = &(*pn).right
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}) &&
go func() {
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