github.com/cockroachdb/tools@v0.0.0-20230222021103-a6d27438930d/internal/persistent/map.go

// Copyright 2022 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.

// The persistent package defines various persistent data structures;
// that is, data structures that can be efficiently copied and modified
// in sublinear time.
package persistent

import (
	"fmt"
	"math/rand"
	"strings"
	"sync/atomic"
)

// Implementation details:
// * Each value is reference counted by nodes which hold it.
// * Each node is reference counted by its parent nodes.
// * Each map is considered a top-level parent node from reference counting perspective.
// * Each change does always effectivelly produce a new top level node.
//
// Functions which operate directly with nodes do have a notation in form of
// `foo(arg1:+n1, arg2:+n2) (ret1:+n3)`.
// Each argument is followed by a delta change to its reference counter.
// In case if no change is expected, the delta will be `-0`.

// Map is an associative mapping from keys to values, both represented as
// interface{}. Key comparison and iteration order is defined by a
// client-provided function that implements a strict weak order.
//
// Maps can be Cloned in constant time.
// Get, Store, and Delete operations are done on average in logarithmic time.
// Maps can be Updated in O(m log(n/m)) time for maps of size n and m, where m < n.
//
// Values are reference counted, and a client-supplied release function
// is called when a value is no longer referenced by a map or any clone.
//
// Internally the implementation is based on a randomized persistent treap:
// https://en.wikipedia.org/wiki/Treap.
type Map struct {
	less func(a, b interface{}) bool
	root *mapNode
}

func (m *Map) String() string {
	var buf strings.Builder
	buf.WriteByte('{')
	var sep string
	m.Range(func(k, v interface{}) {
		fmt.Fprintf(&buf, "%s%v: %v", sep, k, v)
		sep = ", "
	})
	buf.WriteByte('}')
	return buf.String()
}

type mapNode struct {
	key         interface{}
	value       *refValue
	weight      uint64
	refCount    int32
	left, right *mapNode
}

type refValue struct {
	refCount int32
	value    interface{}
	release  func(key, value interface{})
}

func newNodeWithRef(key, value interface{}, release func(key, value interface{})) *mapNode {
	return &mapNode{
		key: key,
		value: &refValue{
			value:    value,
			release:  release,
			refCount: 1,
		},
		refCount: 1,
		weight:   rand.Uint64(),
	}
}

func (node *mapNode) shallowCloneWithRef() *mapNode {
	atomic.AddInt32(&node.value.refCount, 1)
	return &mapNode{
		key:      node.key,
		value:    node.value,
		weight:   node.weight,
		refCount: 1,
	}
}

func (node *mapNode) incref() *mapNode {
	if node != nil {
		atomic.AddInt32(&node.refCount, 1)
	}
	return node
}

func (node *mapNode) decref() {
	if node == nil {
		return
	}
	if atomic.AddInt32(&node.refCount, -1) == 0 {
		if atomic.AddInt32(&node.value.refCount, -1) == 0 {
			if node.value.release != nil {
				node.value.release(node.key, node.value.value)
			}
			node.value.value = nil
			node.value.release = nil
		}
		node.left.decref()
		node.right.decref()
	}
}

// NewMap returns a new map whose keys are ordered by the given comparison
// function (a strict weak order). It is the responsibility of the caller to
// Destroy it at later time.
func NewMap(less func(a, b interface{}) bool) *Map {
	return &Map{
		less: less,
	}
}

// Clone returns a copy of the given map. It is a responsibility of the caller
// to Destroy it at later time.
func (pm *Map) Clone() *Map {
	return &Map{
		less: pm.less,
		root: pm.root.incref(),
	}
}

// Destroy destroys the map.
//
// After Destroy, the Map should not be used again.
func (pm *Map) Destroy() {
	// The implementation of these two functions is the same,
	// but their intent is different.
	pm.Clear()
}

// Clear removes all entries from the map.
func (pm *Map) Clear() {
	pm.root.decref()
	pm.root = nil
}

// Range calls f sequentially in ascending key order for all entries in the map.
func (pm *Map) Range(f func(key, value interface{})) {
	pm.root.forEach(f)
}

func (node *mapNode) forEach(f func(key, value interface{})) {
	if node == nil {
		return
	}
	node.left.forEach(f)
	f(node.key, node.value.value)
	node.right.forEach(f)
}

// Get returns the map value associated with the specified key, or nil if no entry
// is present. The ok result indicates whether an entry was found in the map.
func (pm *Map) Get(key interface{}) (interface{}, bool) {
	node := pm.root
	for node != nil {
		if pm.less(key, node.key) {
			node = node.left
		} else if pm.less(node.key, key) {
			node = node.right
		} else {
			return node.value.value, true
		}
	}
	return nil, false
}

// SetAll updates the map with key/value pairs from the other map, overwriting existing keys.
// It is equivalent to calling Set for each entry in the other map but is more efficient.
// Both maps must have the same comparison function, otherwise behavior is undefined.
func (pm *Map) SetAll(other *Map) {
	root := pm.root
	pm.root = union(root, other.root, pm.less, true)
	root.decref()
}

// Set updates the value associated with the specified key.
// If release is non-nil, it will be called with entry's key and value once the
// key is no longer contained in the map or any clone.
func (pm *Map) Set(key, value interface{}, release func(key, value interface{})) {
	first := pm.root
	second := newNodeWithRef(key, value, release)
	pm.root = union(first, second, pm.less, true)
	first.decref()
	second.decref()
}

// union returns a new tree which is a union of first and second one.
// If overwrite is set to true, second one would override a value for any duplicate keys.
//
// union(first:-0, second:-0) (result:+1)
// Union borrows both subtrees without affecting their refcount and returns a
// new reference that the caller is expected to call decref.
func union(first, second *mapNode, less func(a, b interface{}) bool, overwrite bool) *mapNode {
	if first == nil {
		return second.incref()
	}
	if second == nil {
		return first.incref()
	}

	if first.weight < second.weight {
		second, first, overwrite = first, second, !overwrite
	}

	left, mid, right := split(second, first.key, less, false)
	var result *mapNode
	if overwrite && mid != nil {
		result = mid.shallowCloneWithRef()
	} else {
		result = first.shallowCloneWithRef()
	}
	result.weight = first.weight
	result.left = union(first.left, left, less, overwrite)
	result.right = union(first.right, right, less, overwrite)
	left.decref()
	mid.decref()
	right.decref()
	return result
}

// split the tree midway by the key into three different ones.
// Return three new trees: left with all nodes with smaller than key, mid with
// the node matching the key, right with all nodes larger than key.
// If there are no nodes in one of trees, return nil instead of it.
// If requireMid is set (such as during deletion), then all return arguments
// are nil if mid is not found.
//
// split(n:-0) (left:+1, mid:+1, right:+1)
// Split borrows n without affecting its refcount, and returns three
// new references that that caller is expected to call decref.
func split(n *mapNode, key interface{}, less func(a, b interface{}) bool, requireMid bool) (left, mid, right *mapNode) {
	if n == nil {
		return nil, nil, nil
	}

	if less(n.key, key) {
		left, mid, right := split(n.right, key, less, requireMid)
		if requireMid && mid == nil {
			return nil, nil, nil
		}
		newN := n.shallowCloneWithRef()
		newN.left = n.left.incref()
		newN.right = left
		return newN, mid, right
	} else if less(key, n.key) {
		left, mid, right := split(n.left, key, less, requireMid)
		if requireMid && mid == nil {
			return nil, nil, nil
		}
		newN := n.shallowCloneWithRef()
		newN.left = right
		newN.right = n.right.incref()
		return left, mid, newN
	}
	mid = n.shallowCloneWithRef()
	return n.left.incref(), mid, n.right.incref()
}

// Delete deletes the value for a key.
func (pm *Map) Delete(key interface{}) {
	root := pm.root
	left, mid, right := split(root, key, pm.less, true)
	if mid == nil {
		return
	}
	pm.root = merge(left, right)
	left.decref()
	mid.decref()
	right.decref()
	root.decref()
}

// merge two trees while preserving the weight invariant.
// All nodes in left must have smaller keys than any node in right.
//
// merge(left:-0, right:-0) (result:+1)
// Merge borrows its arguments without affecting their refcount
// and returns a new reference that the caller is expected to call decref.
func merge(left, right *mapNode) *mapNode {
	switch {
	case left == nil:
		return right.incref()
	case right == nil:
		return left.incref()
	case left.weight > right.weight:
		root := left.shallowCloneWithRef()
		root.left = left.left.incref()
		root.right = merge(left.right, right)
		return root
	default:
		root := right.shallowCloneWithRef()
		root.left = merge(left, right.left)
		root.right = right.right.incref()
		return root
	}
}