gomapllrb.go

// Package gomapllrb implements an in-memory key/value store using LLRB algorithm.
// LLRB (Left-Leaning Red-Black) is a self-balancing binary search tree that
// stores the keys in order which allows ordered iteration and find nearest keys.
//
// Copyright (c) 2023, Seungyoung Kim
// https://github.com/wolkykim/GoMapLLRB
package gomapllrb

import (
	"bytes"
	"fmt"
	"sync"

	"golang.org/x/exp/constraints"
)

const (
	// LLRB234 sets the tree management property and algorithm.
	LLRB234 = true // true: 2-3-4 varian(default), false: 2-3 variant
)

// Tree is the glorious tree struct.
type Tree[K constraints.Ordered] struct {
	isLess Comparator[K] // data comparator (default: string comparator)

	root  *Node[K]     // root node
	len   int          // number of object stored
	mutex sync.RWMutex // reader/writer mutual exclusion lock

	stats Stats // usage and performance metrics
}

// Node is like an apple on the apple trees.
type Node[K constraints.Ordered] struct {
	name K
	data interface{}

	red   bool
	up    *Node[K]
	left  *Node[K]
	right *Node[K]
}

// Stats provides usage statistics accessible via Stats() method.
type Stats struct {
	Put struct {
		Sum    uint64
		New    uint64
		Update uint64
	}
	Delete struct {
		Sum      uint64
		Deleted  uint64
		NotFound uint64
	}
	Get struct {
		Sum      uint64
		Found    uint64
		NotFound uint64
	}
	Perf PerfStats
}

// PerfStats are global stats for debugging purpose.
type PerfStats struct {
	Flip   uint64
	Rotate struct {
		Sum   uint64
		Left  uint64
		Right uint64
	}
}

// New creates a new tree.
func New[K constraints.Ordered]() *Tree[K] {
	return &Tree[K]{
		isLess: IsLess[K],
	}
}

// SetLess sets a user comparator function.
//
//	func myLess[K constraints.Ordered](a, b K) bool {
//	  // return true if a < b, or false
//	}
func (tree *Tree[K]) SetLess(fn Comparator[K]) {
	tree.isLess = fn
}

// Put inserts a new key or replaces old if the same key is found.
func (tree *Tree[K]) Put(name K, data interface{}) {
	tree.mutex.Lock()
	defer tree.mutex.Unlock()
	tree.root = tree.put(tree.root, name, data)
	tree.root.red = false
}

// Delete deletes the key. It returns an error if the key is not found.
func (tree *Tree[K]) Delete(name K) bool {
	tree.mutex.Lock()
	defer tree.mutex.Unlock()
	var deleted bool
	tree.root, deleted = tree.delete(tree.root, name)
	if tree.root != nil {
		tree.root.red = false
	}
	return deleted
}

// Get returns the value of the key. If key is not found, it returns Nil.
// When Nil value is expected as a actual value, use Exist() instead.
func (tree *Tree[K]) Get(name K) interface{} {
	tree.mutex.RLock()
	defer tree.mutex.RUnlock()
	if node := tree.get(tree.root, name); node != nil {
		return node.data
	}
	return nil
}

// Exist checks if the key exists.
func (tree *Tree[K]) Exist(name K) bool {
	tree.mutex.RLock()
	defer tree.mutex.RUnlock()
	if node := tree.get(tree.root, name); node != nil {
		return true
	}
	return false
}

// Min returns a min key and value.
func (tree *Tree[K]) Min() (K, interface{}, bool) {
	tree.mutex.RLock()
	defer tree.mutex.RUnlock()
	if node := findMin(tree.root); node != nil {
		return node.name, node.data, true
	}
	var n K
	return n, nil, false
}

// Max returns a max key and value.
func (tree *Tree[K]) Max() (K, interface{}, bool) {
	tree.mutex.RLock()
	defer tree.mutex.RUnlock()
	if node := findMax(tree.root); node != nil {
		return node.name, node.data, true
	}
	var n K
	return n, nil, false
}

// Bigger finds the next key bigger than given ken.
func (tree *Tree[K]) Bigger(name K) (K, interface{}, bool) {
	tree.mutex.RLock()
	defer tree.mutex.RUnlock()
	if node := tree.bigger(tree.root, name, false); node != nil {
		return node.name, node.data, true
	}
	var n K
	return n, nil, false
}

// Smaller finds the next key bigger than given ken.
func (tree *Tree[K]) Smaller(name K) (K, interface{}, bool) {
	tree.mutex.RLock()
	defer tree.mutex.RUnlock()
	if node := tree.smaller(tree.root, name, false); node != nil {
		return node.name, node.data, true
	}
	var n K
	return n, nil, false
}

// EqualOrBigger finds a matching key or the next bigger key.
func (tree *Tree[K]) EqualOrBigger(name K) (K, interface{}, bool) {
	tree.mutex.RLock()
	defer tree.mutex.RUnlock()
	if node := tree.bigger(tree.root, name, true); node != nil {
		return node.name, node.data, true
	}
	var n K
	return n, nil, false
}

// EqualOrSmaller finds a matching key or the next smaller key.
func (tree *Tree[K]) EqualOrSmaller(name K) (K, interface{}, bool) {
	tree.mutex.RLock()
	defer tree.mutex.RUnlock()
	if node := tree.smaller(tree.root, name, true); node != nil {
		return node.name, node.data, true
	}
	var n K
	return n, nil, false
}

// Clear empties the tree without resetting the statistic metrics.
func (tree *Tree[K]) Clear() {
	tree.mutex.Lock()
	defer tree.mutex.Unlock()
	tree.root = nil
	tree.len = 0
}

// Len returns the number of object stored.
func (tree *Tree[K]) Len() int {
	return tree.len
}

// Stats returns a copy of the statistics metrics.
func (tree *Tree[K]) Stats() Stats {
	tree.stats.Put.Sum = tree.stats.Put.New + tree.stats.Put.Update
	tree.stats.Get.Sum = tree.stats.Get.Found + tree.stats.Get.NotFound
	tree.stats.Delete.Sum = tree.stats.Delete.Deleted + tree.stats.Delete.NotFound
	tree.stats.Perf = pstats
	tree.stats.Perf.Rotate.Sum = tree.stats.Perf.Rotate.Left + tree.stats.Perf.Rotate.Right
	return tree.stats
}

// ResetStats resets all the satistics metrics.
func (tree *Tree[K]) ResetStats() {
	tree.stats = Stats{}
	pstats = PerfStats{}
}

// String returns a pretty drawing of the tree structure.
//
//	┌── 6
//	|   └──[5]
//	4
//	│   ┌── 3
//	└──[2]
//	    └── 1
func (tree *Tree[K]) String() string {
	var buf bytes.Buffer
	tree.mutex.RLock()
	defer tree.mutex.RUnlock()
	printNode(tree.root, &buf, nil, false)
	return buf.String()
}

// Map returns the tree in a map
func (tree *Tree[K]) Map() map[K]interface{} {
	m := make(map[K]interface{}, tree.Len())
	for it := tree.Iter(); it.Next(); {
		m[it.Key()] = it.Val()
	}
	return m
}

// String returns a statistics data in a string.
func (s Stats) String() string {
	variant := "234"
	if !LLRB234 {
		variant = "23"
	}
	numUpdate := s.Put.Sum + s.Delete.Sum
	return fmt.Sprintf("Variant:LLRB%s, Put:%d, Delete:%d, Get:%d, Rotate:%0.2f, Flip:%0.2f",
		variant, s.Put.Sum, s.Delete.Sum, s.Get.Sum,
		float64(s.Perf.Rotate.Sum)/float64(numUpdate),
		float64(s.Perf.Flip)/float64(numUpdate))
}

// Check checks that the invariants of the red-black tree are satisfied.
//
//	Root property:  The root is black.
//	Red property:   If a node is red, then both its children are black.
//	Black property: For each node, all simple paths from the node to
//	                descendant leaves contain the same number of black nodes.
//	LLRB property:  3-nodes always lean to the left and 4-nodes are balanced.
func (tree *Tree[K]) Check() error {
	if err := checkRoot(tree.root); err != nil {
		return err
	}
	if err := checkRed(tree.root); err != nil {
		return err
	}
	length := 0
	if err := checkBlack(tree.root, &length); err != nil {
		return err
	}
	return checkLLRB(tree.root)
}

/*************************************************************************
 * Iterator
 ************************************************************************/

// Iter is a iterator object.
type Iter[K constraints.Ordered] struct {
	tree *Tree[K]
	cur  *Node[K] // cursor, start from
	last *Node[K] // last node pointer after next()
	end  K        // end boundary is span is set
	span bool     // indicates the end boundary is set
	done bool     // indicates the iteration is complete
}

// Iter returns an iterator.
func (tree *Tree[K]) Iter() *Iter[K] {
	tree.mutex.RLock()
	defer tree.mutex.RUnlock()
	it := &Iter[K]{
		tree: tree,
		cur:  findMin(tree.root),
	}
	if it.cur == nil {
		it.done = true
	}
	return it
}

// Range returns a ranged iterator.
func (tree *Tree[K]) Range(start, end K) *Iter[K] {
	tree.mutex.RLock()
	defer tree.mutex.RUnlock()
	it := &Iter[K]{
		tree: tree,
		cur:  tree.bigger(tree.root, start, true),
		end:  end,
		span: true,
	}
	if it.cur == nil {
		it.done = true
	}
	return it
}

// Next travels the keys in the tree.
func (it *Iter[K]) Next() bool {
	if it.done {
		return false
	}
	it.last = it.cur
	it.tree.mutex.RLock()
	defer it.tree.mutex.RUnlock()
	if it.cur = it.tree.bigger(it.cur.right, it.cur.name, false); it.cur == nil {
		it.cur = it.last.up
		// go up until bigger value found
		for it.cur != nil && it.tree.isLess(it.cur.name, it.last.name) {
			it.cur = it.cur.up
		}
		// go down again
		if it.cur != nil {
			it.cur = it.tree.bigger(it.cur, it.last.name, false)
		}
		if it.cur == nil {
			it.done = true
		}
	}
	if !it.done && it.span && it.tree.isLess(it.end, it.cur.name) {
		it.done = true
	}
	return true
}

// Key returns the key name.
func (it *Iter[K]) Key() K {
	if it.last == nil {
		var k K
		return k
	}
	return it.last.name
}

// Val returns the value data.
func (it *Iter[K]) Val() interface{} {
	if it.last == nil {
		return nil
	}
	return it.last.data
}

/*************************************************************************
 * Default comparators
 ************************************************************************/

// Comparator is the type.
type Comparator[K constraints.Ordered] func(a, b K) bool

// IsLess is the default comparator.
func IsLess[K constraints.Ordered](a, b K) bool {
	return a < b
}

/*************************************************************************
 * User data manipulation functions
 ************************************************************************/
func (tree *Tree[K]) put(node *Node[K], name K, data interface{}) *Node[K] {
	if node == nil {
		tree.len++
		tree.stats.Put.New++
		return newNode[K](name, data)
	}

	if LLRB234 {
		// split 4-nodes on the way down
		if isRed(node.left) && isRed(node.right) {
			flipColor(node)
		}
	}

	if tree.isLess(name, node.name) {
		node.left = tree.put(node.left, name, data)
		node.left.up = node
	} else if tree.isLess(node.name, name) {
		node.right = tree.put(node.right, name, data)
		node.right.up = node
	} else { // existing key found
		node.data = data
		tree.stats.Put.Update++
	}

	// fix right-leaning reds on the way up
	if isRed(node.right) && !isRed(node.left) {
		node = rotateLeft(node)
	}

	// fix two reds in a row on the way up
	if isRed(node.left) && isRed(node.left.left) {
		node = rotateRight(node)
	}

	if !LLRB234 {
		// split 4-nodes on the way up
		if isRed(node.left) && isRed(node.right) {
			flipColor(node)
		}
	}

	// return new root
	return node
}

func (tree *Tree[K]) delete(node *Node[K], name K) (*Node[K], bool) {
	if node == nil {
		tree.stats.Delete.NotFound++
		return nil, false
	}

	deleted := false
	if tree.isLess(name, node.name) {
		// move red left
		if node.left != nil && (!isRed(node.left) && !isRed(node.left.left)) {
			node = moveRedLeft(node)
		}
		// keep going down to the left
		node.left, deleted = tree.delete(node.left, name)
	} else { // right or equal
		if isRed(node.left) {
			node = rotateRight(node)
		}
		// remove if equal at the bottom
		if node.right == nil && !tree.isLess(node.name, name) {
			tree.len--
			tree.stats.Delete.Deleted++
			return nil, true
		}
		// move red right
		if node.right != nil && (!isRed(node.right) && !isRed(node.right.left)) {
			node = moveRedRight(node)
		}
		// found in the middle
		if !tree.isLess(node.name, name) {
			// we delete the min node from the right instead
			var min *Node[K]
			node.right, min = deleteMin(node.right)
			// then copy the min node to this
			node.name = min.name
			node.data = min.data
			tree.len--
			tree.stats.Delete.Deleted++
		} else {
			// keep going down to the right
			node.right, deleted = tree.delete(node.right, name)
		}
	}
	// fix right-leaning red nodes on the way up
	return fixNode(node), deleted
}

func (tree *Tree[K]) get(node *Node[K], name K) *Node[K] {
	// do linear search for performance
	for node != nil {
		if tree.isLess(name, node.name) {
			node = node.left
		} else if tree.isLess(node.name, name) {
			node = node.right
		} else {
			tree.stats.Get.Found++
			return node
		}
	}
	tree.stats.Get.NotFound++
	return nil
}

func (tree *Tree[K]) bigger(node *Node[K], name K, equal bool) *Node[K] {
	if node == nil {
		return nil
	}
	this := node
	if tree.isLess(name, node.name) {
		if node = tree.bigger(node.left, name, equal); node == nil {
			node = this
		}
	} else if tree.isLess(node.name, name) {
		node = tree.bigger(node.right, name, equal)
	} else if !equal {
		// match found, continue to the right
		node = tree.bigger(node.right, name, equal)
	}
	return node
}

func (tree *Tree[K]) smaller(node *Node[K], name K, equal bool) *Node[K] {
	if node == nil {
		return nil
	}
	this := node
	if tree.isLess(name, node.name) {
		node = tree.smaller(node.left, name, equal)
	} else if tree.isLess(node.name, name) {
		if node = tree.smaller(node.right, name, equal); node == nil {
			node = this
		}
	} else if !equal {
		// match found, continue to the left
		node = tree.smaller(node.left, name, equal)
	}
	return node
}

/*************************************************************************
 * Tree property management functions
 ************************************************************************/

var pstats PerfStats

func newNode[K constraints.Ordered](name K, data interface{}) *Node[K] {
	return &Node[K]{
		name: name,
		data: data,
		red:  true,
	}
}

func isRed[K constraints.Ordered](node *Node[K]) bool {
	if node == nil {
		return false
	}
	return node.red
}

func flipColor[K constraints.Ordered](node *Node[K]) {
	node.red = !node.red
	node.left.red = !node.left.red
	node.right.red = !node.right.red
	pstats.Flip++
}

func rotateLeft[K constraints.Ordered](node *Node[K]) *Node[K] {
	n := node.right
	n.up = node.up
	node.up = n
	node.right = n.left
	n.left = node
	n.red = n.left.red
	n.left.red = true
	pstats.Rotate.Left++
	return n
}

func rotateRight[K constraints.Ordered](node *Node[K]) *Node[K] {
	n := node.left
	n.up = node.up
	node.up = n
	node.left = n.right
	n.right = node
	n.red = n.right.red
	n.right.red = true
	pstats.Rotate.Right++
	return n
}

func moveRedLeft[K constraints.Ordered](node *Node[K]) *Node[K] {
	flipColor(node)
	if isRed(node.right.left) {
		node.right = rotateRight(node.right)
		node = rotateLeft(node)
		flipColor(node)
		if LLRB234 {
			// 2-3-4 exclusive
			if isRed(node.right.right) {
				node.right = rotateLeft(node.right)
			}
		}
	}
	return node
}

func moveRedRight[K constraints.Ordered](node *Node[K]) *Node[K] {
	flipColor(node)
	if isRed(node.left.left) {
		node = rotateRight(node)
		flipColor(node)
	}
	return node
}

func findMin[K constraints.Ordered](node *Node[K]) *Node[K] {
	if node == nil {
		return nil
	}

	for node.left != nil {
		node = node.left
	}
	return node
}

func findMax[K constraints.Ordered](node *Node[K]) *Node[K] {
	if node == nil {
		return nil
	}

	for node.right != nil {
		node = node.right
	}
	return node
}

func deleteMin[K constraints.Ordered](node *Node[K]) (*Node[K], *Node[K]) {
	if node.left == nil {
		// 3-nodes are left-leaning, so this is a leaf.
		return nil, node
	}
	if !isRed(node.left) && !isRed(node.left.left) {
		node = moveRedLeft(node)
	}
	var min *Node[K]
	node.left, min = deleteMin(node.left)
	return fixNode(node), min
}

func fixNode[K constraints.Ordered](node *Node[K]) *Node[K] {
	// rotate right red to left
	if isRed(node.right) {
		if LLRB234 {
			if isRed(node.right.left) {
				node.right = rotateRight(node.right)
			}
		}
		node = rotateLeft(node)
	}
	// rotate left red-red to right
	if isRed(node.left) && isRed(node.left.left) {
		node = rotateRight(node)
	}

	if !LLRB234 {
		// split 4-nodes
		if isRed(node.left) && isRed(node.right) {
			flipColor(node)
		}
	}
	return node
}

/*************************************************************************
 * Integrity checks
 ************************************************************************/

// checkRoot verifies that root property of the red-black tree is satisfied.
func checkRoot[K constraints.Ordered](root *Node[K]) error {
	if isRed(root) {
		return fmt.Errorf("root property violation found")
	}

	return nil
}

// checkRed verifies that red property of the red-black tree is satisfied.
func checkRed[K constraints.Ordered](node *Node[K]) error {
	if node == nil {
		return nil
	}

	if isRed(node) && (isRed(node.right) || isRed(node.left)) {
		return fmt.Errorf("red property violation found")
	}
	if err := checkRed(node.right); err != nil {
		return err
	}
	return checkRed(node.left)
}

// checkBlack verifies that black property of the red-black tree is satisfied.
func checkBlack[K constraints.Ordered](node *Node[K], length *int) error {
	if node == nil {
		*length = 1
		return nil
	}

	var rightLength int
	if err := checkBlack(node.right, &rightLength); err != nil {
		return err
	}
	var leftLength int
	if err := checkBlack(node.left, &leftLength); err != nil {
		return err
	}

	if rightLength != leftLength {
		return fmt.Errorf("black property violation found")
	}
	if !isRed(node) {
		*length = rightLength + 1
	} else {
		*length = rightLength
	}
	return nil
}

// checkLLRB verifies that LLRB property of the left-leaning red-black tree is satisfied.
func checkLLRB[K constraints.Ordered](node *Node[K]) error {
	if node == nil {
		return nil
	}

	if isRed(node.right) && !isRed(node.left) {
		return fmt.Errorf("LLRB property violation found")
	}
	if err := checkLLRB(node.right); err != nil {
		return err
	}
	return checkLLRB(node.left)
}

/*************************************************************************
 * Tree printing functions
 ************************************************************************/

type branchObj struct {
	prev *branchObj
	str  string
}

func printBranch(branch *branchObj, out *bytes.Buffer) {
	if branch == nil {
		return
	}
	printBranch(branch.prev, out)
	out.WriteString(branch.str)
}

func printNode[K constraints.Ordered](node *Node[K], out *bytes.Buffer, pbranch *branchObj, right bool) {
	if node == nil {
		return
	}

	branch := &branchObj{
		prev: pbranch,
	}

	if pbranch != nil {
		if right {
			branch.str = "    "
		} else {
			branch.str = "│   "
		}
	} else {
		branch.str = ""
	}
	printNode(node.right, out, branch, true)

	printBranch(pbranch, out)
	if pbranch != nil {
		if right {
			out.WriteString("┌──")
		} else {
			out.WriteString("└──")
		}
		if node.red {
			out.WriteString("[")
		} else {
			out.WriteString("─")
		}
	}
	out.WriteString(fmt.Sprintf("%v", node.name))
	if node.red {
		out.WriteString("]\n")
	} else {
		out.WriteString(" \n")
	}

	if pbranch != nil {
		if right {
			branch.str = "│   "
		} else {
			branch.str = "    "
		}
	} else {
		branch.str = ""
	}
	printNode(node.left, out, branch, false)
}