github.com/gnolang/gno@v0.0.0-20240520182011-228e9d0192ce/examples/gno.land/p/demo/avl/node.gno

package avl

//----------------------------------------
// Node

// Node represents a node in an AVL tree.
type Node struct {
	key       string      // key is the unique identifier for the node.
	value     interface{} // value is the data stored in the node.
	height    int8        // height is the height of the node in the tree.
	size      int         // size is the number of nodes in the subtree rooted at this node.
	leftNode  *Node       // leftNode is the left child of the node.
	rightNode *Node       // rightNode is the right child of the node.
}

// NewNode creates a new node with the given key and value.
func NewNode(key string, value interface{}) *Node {
	return &Node{
		key:    key,
		value:  value,
		height: 0,
		size:   1,
	}
}

// Size returns the size of the subtree rooted at the node.
func (node *Node) Size() int {
	if node == nil {
		return 0
	}
	return node.size
}

// IsLeaf checks if the node is a leaf node (has no children).
func (node *Node) IsLeaf() bool {
	return node.height == 0
}

// Key returns the key of the node.
func (node *Node) Key() string {
	return node.key
}

// Value returns the value of the node.
func (node *Node) Value() interface{} {
	return node.value
}

// _copy creates a copy of the node (excluding value).
func (node *Node) _copy() *Node {
	if node.height == 0 {
		panic("Why are you copying a value node?")
	}
	return &Node{
		key:       node.key,
		height:    node.height,
		size:      node.size,
		leftNode:  node.leftNode,
		rightNode: node.rightNode,
	}
}

// Has checks if a node with the given key exists in the subtree rooted at the node.
func (node *Node) Has(key string) (has bool) {
	if node == nil {
		return false
	}
	if node.key == key {
		return true
	}
	if node.height == 0 {
		return false
	}
	if key < node.key {
		return node.getLeftNode().Has(key)
	}
	return node.getRightNode().Has(key)
}

// Get searches for a node with the given key in the subtree rooted at the node
// and returns its index, value, and whether it exists.
func (node *Node) Get(key string) (index int, value interface{}, exists bool) {
	if node == nil {
		return 0, nil, false
	}

	if node.height == 0 {
		if node.key == key {
			return 0, node.value, true
		}
		if node.key < key {
			return 1, nil, false
		}
		return 0, nil, false
	}

	if key < node.key {
		return node.getLeftNode().Get(key)
	}

	rightNode := node.getRightNode()
	index, value, exists = rightNode.Get(key)
	index += node.size - rightNode.size
	return index, value, exists
}

// GetByIndex retrieves the key-value pair of the node at the given index
// in the subtree rooted at the node.
func (node *Node) GetByIndex(index int) (key string, value interface{}) {
	if node.height == 0 {
		if index == 0 {
			return node.key, node.value
		}
		panic("GetByIndex asked for invalid index")
	}
	// TODO: could improve this by storing the sizes
	leftNode := node.getLeftNode()
	if index < leftNode.size {
		return leftNode.GetByIndex(index)
	}
	return node.getRightNode().GetByIndex(index - leftNode.size)
}

// Set inserts a new node with the given key-value pair into the subtree rooted at the node,
// and returns the new root of the subtree and whether an existing node was updated.
//
// XXX consider a better way to do this... perhaps split Node from Node.
func (node *Node) Set(key string, value interface{}) (newSelf *Node, updated bool) {
	if node == nil {
		return NewNode(key, value), false
	}

	if node.height == 0 {
		return node.setLeaf(key, value)
	}

	node = node._copy()
	if key < node.key {
		node.leftNode, updated = node.getLeftNode().Set(key, value)
	} else {
		node.rightNode, updated = node.getRightNode().Set(key, value)
	}

	if updated {
		return node, updated
	}

	node.calcHeightAndSize()
	return node.balance(), updated
}

// setLeaf inserts a new leaf node with the given key-value pair into the subtree rooted at the node,
// and returns the new root of the subtree and whether an existing node was updated.
func (node *Node) setLeaf(key string, value interface{}) (newSelf *Node, updated bool) {
	if key == node.key {
		return NewNode(key, value), true
	}

	if key < node.key {
		return &Node{
			key:       node.key,
			height:    1,
			size:      2,
			leftNode:  NewNode(key, value),
			rightNode: node,
		}, false
	}

	return &Node{
		key:       key,
		height:    1,
		size:      2,
		leftNode:  node,
		rightNode: NewNode(key, value),
	}, false
}

// Remove deletes the node with the given key from the subtree rooted at the node.
// returns the new root of the subtree, the new leftmost leaf key (if changed),
// the removed value and the removal was successful.
func (node *Node) Remove(key string) (
	newNode *Node, newKey string, value interface{}, removed bool,
) {
	if node == nil {
		return nil, "", nil, false
	}
	if node.height == 0 {
		if key == node.key {
			return nil, "", node.value, true
		}
		return node, "", nil, false
	}
	if key < node.key {
		var newLeftNode *Node
		newLeftNode, newKey, value, removed = node.getLeftNode().Remove(key)
		if !removed {
			return node, "", value, false
		}
		if newLeftNode == nil { // left node held value, was removed
			return node.rightNode, node.key, value, true
		}
		node = node._copy()
		node.leftNode = newLeftNode
		node.calcHeightAndSize()
		node = node.balance()
		return node, newKey, value, true
	}

	var newRightNode *Node
	newRightNode, newKey, value, removed = node.getRightNode().Remove(key)
	if !removed {
		return node, "", value, false
	}
	if newRightNode == nil { // right node held value, was removed
		return node.leftNode, "", value, true
	}
	node = node._copy()
	node.rightNode = newRightNode
	if newKey != "" {
		node.key = newKey
	}
	node.calcHeightAndSize()
	node = node.balance()
	return node, "", value, true
}

// getLeftNode returns the left child of the node.
func (node *Node) getLeftNode() *Node {
	return node.leftNode
}

// getRightNode returns the right child of the node.
func (node *Node) getRightNode() *Node {
	return node.rightNode
}

// rotateRight performs a right rotation on the node and returns the new root.
// NOTE: overwrites node
// TODO: optimize balance & rotate
func (node *Node) rotateRight() *Node {
	node = node._copy()
	l := node.getLeftNode()
	_l := l._copy()

	_lrCached := _l.rightNode
	_l.rightNode = node
	node.leftNode = _lrCached

	node.calcHeightAndSize()
	_l.calcHeightAndSize()

	return _l
}

// rotateLeft performs a left rotation on the node and returns the new root.
// NOTE: overwrites node
// TODO: optimize balance & rotate
func (node *Node) rotateLeft() *Node {
	node = node._copy()
	r := node.getRightNode()
	_r := r._copy()

	_rlCached := _r.leftNode
	_r.leftNode = node
	node.rightNode = _rlCached

	node.calcHeightAndSize()
	_r.calcHeightAndSize()

	return _r
}

// calcHeightAndSize updates the height and size of the node based on its children.
// NOTE: mutates height and size
func (node *Node) calcHeightAndSize() {
	node.height = maxInt8(node.getLeftNode().height, node.getRightNode().height) + 1
	node.size = node.getLeftNode().size + node.getRightNode().size
}

// calcBalance calculates the balance factor of the node.
func (node *Node) calcBalance() int {
	return int(node.getLeftNode().height) - int(node.getRightNode().height)
}

// balance balances the subtree rooted at the node and returns the new root.
// NOTE: assumes that node can be modified
// TODO: optimize balance & rotate
func (node *Node) balance() (newSelf *Node) {
	balance := node.calcBalance()
	if balance >= -1 {
		return node
	}
	if balance > 1 {
		if node.getLeftNode().calcBalance() >= 0 {
			// Left Left Case
			return node.rotateRight()
		}
		// Left Right Case
		left := node.getLeftNode()
		node.leftNode = left.rotateLeft()
		return node.rotateRight()
	}

	if node.getRightNode().calcBalance() <= 0 {
		// Right Right Case
		return node.rotateLeft()
	}

	// Right Left Case
	right := node.getRightNode()
	node.rightNode = right.rotateRight()
	return node.rotateLeft()
}

// Shortcut for TraverseInRange.
func (node *Node) Iterate(start, end string, cb func(*Node) bool) bool {
	return node.TraverseInRange(start, end, true, true, cb)
}

// Shortcut for TraverseInRange.
func (node *Node) ReverseIterate(start, end string, cb func(*Node) bool) bool {
	return node.TraverseInRange(start, end, false, true, cb)
}

// TraverseInRange traverses all nodes, including inner nodes.
// Start is inclusive and end is exclusive when ascending,
// Start and end are inclusive when descending.
// Empty start and empty end denote no start and no end.
// If leavesOnly is true, only visit leaf nodes.
// NOTE: To simulate an exclusive reverse traversal,
// just append 0x00 to start.
func (node *Node) TraverseInRange(start, end string, ascending bool, leavesOnly bool, cb func(*Node) bool) bool {
	if node == nil {
		return false
	}
	afterStart := (start == "" || start < node.key)
	startOrAfter := (start == "" || start <= node.key)
	beforeEnd := false
	if ascending {
		beforeEnd = (end == "" || node.key < end)
	} else {
		beforeEnd = (end == "" || node.key <= end)
	}

	// Run callback per inner/leaf node.
	stop := false
	if (!node.IsLeaf() && !leavesOnly) ||
		(node.IsLeaf() && startOrAfter && beforeEnd) {
		stop = cb(node)
		if stop {
			return stop
		}
	}
	if node.IsLeaf() {
		return stop
	}

	if ascending {
		// check lower nodes, then higher
		if afterStart {
			stop = node.getLeftNode().TraverseInRange(start, end, ascending, leavesOnly, cb)
		}
		if stop {
			return stop
		}
		if beforeEnd {
			stop = node.getRightNode().TraverseInRange(start, end, ascending, leavesOnly, cb)
		}
	} else {
		// check the higher nodes first
		if beforeEnd {
			stop = node.getRightNode().TraverseInRange(start, end, ascending, leavesOnly, cb)
		}
		if stop {
			return stop
		}
		if afterStart {
			stop = node.getLeftNode().TraverseInRange(start, end, ascending, leavesOnly, cb)
		}
	}

	return stop
}

// TraverseByOffset traverses all nodes, including inner nodes.
// A limit of math.MaxInt means no limit.
func (node *Node) TraverseByOffset(offset, limit int, descending bool, leavesOnly bool, cb func(*Node) bool) bool {
	if node == nil {
		return false
	}

	// fast paths. these happen only if TraverseByOffset is called directly on a leaf.
	if limit <= 0 || offset >= node.size {
		return false
	}
	if node.IsLeaf() {
		if offset > 0 {
			return false
		}
		return cb(node)
	}

	// go to the actual recursive function.
	return node.traverseByOffset(offset, limit, descending, leavesOnly, cb)
}

// TraverseByOffset traverses the subtree rooted at the node by offset and limit,
// in either ascending or descending order, and applies the callback function to each traversed node.
// If leavesOnly is true, only leaf nodes are visited.
func (node *Node) traverseByOffset(offset, limit int, descending bool, leavesOnly bool, cb func(*Node) bool) bool {
	// caller guarantees: offset < node.size; limit > 0.
	if !leavesOnly {
		if cb(node) {
			return true
		}
	}
	first, second := node.getLeftNode(), node.getRightNode()
	if descending {
		first, second = second, first
	}
	if first.IsLeaf() {
		// either run or skip, based on offset
		if offset > 0 {
			offset--
		} else {
			cb(first)
			limit--
			if limit <= 0 {
				return false
			}
		}
	} else {
		// possible cases:
		// 1 the offset given skips the first node entirely
		// 2 the offset skips none or part of the first node, but the limit requires some of the second node.
		// 3 the offset skips none or part of the first node, and the limit stops our search on the first node.
		if offset >= first.size {
			offset -= first.size // 1
		} else {
			if first.traverseByOffset(offset, limit, descending, leavesOnly, cb) {
				return true
			}
			// number of leaves which could actually be called from inside
			delta := first.size - offset
			offset = 0
			if delta >= limit {
				return true // 3
			}
			limit -= delta // 2
		}
	}

	// because of the caller guarantees and the way we handle the first node,
	// at this point we know that limit > 0 and there must be some values in
	// this second node that we include.

	// => if the second node is a leaf, it has to be included.
	if second.IsLeaf() {
		return cb(second)
	}
	// => if it is not a leaf, it will still be enough to recursively call this
	// function with the updated offset and limit
	return second.traverseByOffset(offset, limit, descending, leavesOnly, cb)
}

// Only used in testing...
func (node *Node) lmd() *Node {
	if node.height == 0 {
		return node
	}
	return node.getLeftNode().lmd()
}

// Only used in testing...
func (node *Node) rmd() *Node {
	if node.height == 0 {
		return node
	}
	return node.getRightNode().rmd()
}

func maxInt8(a, b int8) int8 {
	if a > b {
		return a
	}
	return b
}