fx/core/bst.c

/*
  The Clear BSD License

  Copyright (c) 2023 Max Wash
  All rights reserved.

  Redistribution and use in source and binary forms, with or without
  modification, are permitted (subject to the limitations in the disclaimer
  below) provided that the following conditions are met:

  - Redistributions of source code must retain the above copyright notice,
    this list of conditions and the following disclaimer.

  - Redistributions in binary form must reproduce the above copyright
    notice, this list of conditions and the following disclaimer in the
    documentation and/or other materials provided with the distribution.

  - Neither the name of the copyright holder nor the names of its
    contributors may be used to endorse or promote products derived from this
    software without specific prior written permission.
 */

/* templated AVL binary tree implementation

   this file implements an extensible AVL binary tree data structure.

   the primary rule of an AVL binary tree is that for a given node N,
   the heights of N's left and right substs can differ by at most 1.

   the height of a subst is the length of the longest path between
   the root of the subst and a leaf node, including the root node itself.

   the height of a leaf node is 1.

   when a node is inserted into or deleted from the tree, this rule may
   be broken, in which the tree must be rotated to restore the balance.

   no more than one rotation is required for any insert operations,
   while multiple rotations may be required for a delete operation.

   there are four types of rotations that can be applied to a tree:
     - left rotation
     - right rotation
     - double left rotations
     - double right rotations

   by enforcing the balance rule, for a tree with n nodes, the worst-case
   performance for insert, delete, and search operations is guaranteed
   to be O(log n).

   this file intentionally excludes any kind of search function implementation.
   it is up to the programmer to implement their own tree node type
   using fx_bst_node, and their own search function using fx_bst.
   this allows the programmer to define their own node types with complex
   non-integer key types. bst.h contains a number of macros to help
   define these functions. the macros do all the work, you just have to
   provide a comparator function.
*/

#include <fx/core/bst.h>
#include <stddef.h>

#define MAX(a, b)            ((a) > (b) ? (a) : (b))
#define MIN(a, b)            ((a) < (b) ? (a) : (b))

#define IS_LEFT_CHILD(p, c)  ((p) && (c) && ((p)->n_left == (c)))
#define IS_RIGHT_CHILD(p, c) ((p) && (c) && ((p)->n_right == (c)))

#define HAS_LEFT_CHILD(x)    ((x) && ((x)->n_left))
#define HAS_RIGHT_CHILD(x)   ((x) && ((x)->n_right))

#define HAS_NO_CHILDREN(x)   ((x) && (!(x)->n_left) && (!(x)->n_right))
#define HAS_ONE_CHILD(x)                                                       \
	((HAS_LEFT_CHILD(x) && !HAS_RIGHT_CHILD(x))                            \
	 || (!HAS_LEFT_CHILD(x) && HAS_RIGHT_CHILD(x)))
#define HAS_TWO_CHILDREN(x) (HAS_LEFT_CHILD(x) && HAS_RIGHT_CHILD(x))

#define HEIGHT(x)           ((x) ? (x)->n_height : 0)

struct fx_bst_iterator_p {
	size_t i, depth;
	fx_bst_node *node;
	fx_bst *_b;
};

static inline void update_height(struct fx_bst_node *x)
{
	x->n_height = MAX(HEIGHT(x->n_left), HEIGHT((x->n_right))) + 1;
}

static inline int bf(struct fx_bst_node *x)
{
	int bf = 0;

	if (!x) {
		return bf;
	}

	if (x->n_right) {
		bf += x->n_right->n_height;
	}

	if (x->n_left) {
		bf -= x->n_left->n_height;
	}

	return bf;
}

/* perform a left rotation on a subst

   if you have a tree like this:

        Z
       / \
      X   .
     / \
    .   Y
       / \
      .   .

   and you perform a left rotation on node X,
   you will get the following tree:

        Z
       / \
      Y   .
     / \
    X   .
   / \
  .   .

  note that this function does NOT update fx_height for the rotated
  nodes. it is up to you to call update_height_to_root().
*/
static void rotate_left(struct fx_bst *tree, struct fx_bst_node *x)
{
	struct fx_bst_node *y = x->n_right;
	struct fx_bst_node *p = x->n_parent;

	if (y->n_left) {
		y->n_left->n_parent = x;
	}

	x->n_right = y->n_left;

	if (!p) {
		tree->bst_root = y;
	} else if (x == p->n_left) {
		p->n_left = y;
	} else {
		p->n_right = y;
	}

	x->n_parent = y;
	y->n_left = x;
	y->n_parent = p;
}

static void update_height_to_root(struct fx_bst_node *x)
{
	while (x) {
		update_height(x);
		x = x->n_parent;
	}
}

/* perform a right rotation on a subst

   if you have a tree like this:

        Z
       / \
      .   X
         / \
        Y   .
       / \
      .   .

   and you perform a right rotation on node X,
   you will get the following tree:

        Z
       / \
      .   Y
         / \
        .   X
           / \
          .   .

  note that this function does NOT update fx_height for the rotated
  nodes. it is up to you to call update_height_to_root().
*/
static void rotate_right(struct fx_bst *tree, struct fx_bst_node *y)
{
	struct fx_bst_node *x = y->n_left;
	struct fx_bst_node *p = y->n_parent;

	if (x->n_right) {
		x->n_right->n_parent = y;
	}

	y->n_left = x->n_right;

	if (!p) {
		tree->bst_root = x;
	} else if (y == p->n_left) {
		p->n_left = x;
	} else {
		p->n_right = x;
	}

	y->n_parent = x;
	x->n_right = y;
	x->n_parent = p;
}

/* for a given node Z, perform a right rotation on Z's right child,
   followed by a left rotation on Z itself.

   if you have a tree like this:

        Z
       / \
      .   X
         / \
        Y   .
       / \
      .   .

   and you perform a double-left rotation on node Z,
   you will get the following tree:

         Y
        / \
       /   \
      Z     X
     / \   / \
    .   . .   .

  note that, unlike rotate_left and rotate_right, this function
  DOES update fx_height for the rotated nodes (since it needs to be
  done in a certain order).
*/
static void rotate_double_left(struct fx_bst *tree, struct fx_bst_node *z)
{
	struct fx_bst_node *x = z->n_right;
	struct fx_bst_node *y = x->n_left;

	rotate_right(tree, x);
	rotate_left(tree, z);

	update_height(z);
	update_height(x);

	while (y) {
		update_height(y);
		y = y->n_parent;
	}
}

/* for a given node Z, perform a left rotation on Z's left child,
   followed by a right rotation on Z itself.

   if you have a tree like this:

        Z
       / \
      X   .
     / \
    .   Y
       / \
      .   .

   and you perform a double-right rotation on node Z,
   you will get the following tree:

         Y
        / \
       /   \
      X     Z
     / \   / \
    .   . .   .

  note that, unlike rotate_left and rotate_right, this function
  DOES update fx_height for the rotated nodes (since it needs to be
  done in a certain order).
*/
static void rotate_double_right(struct fx_bst *tree, struct fx_bst_node *z)
{
	struct fx_bst_node *x = z->n_left;
	struct fx_bst_node *y = x->n_right;

	rotate_left(tree, x);
	rotate_right(tree, z);

	update_height(z);
	update_height(x);

	while (y) {
		update_height(y);
		y = y->n_parent;
	}
}

/* run after an insert operation. checks that the balance factor
   of the local subst is within the range -1 <= BF <= 1. if it
   is not, rotate the subst to restore balance.

   note that at most one rotation should be required after a node
   is inserted into the tree.

   this function depends on all nodes in the tree having
   correct fx_height values.

   @param w the node that was just inserted into the tree
*/
static void insert_fixup(struct fx_bst *tree, struct fx_bst_node *w)
{
	struct fx_bst_node *z = NULL, *y = NULL, *x = NULL;

	z = w;
	while (z) {
		if (bf(z) >= -1 && bf(z) <= 1) {
			goto next_ancestor;
		}

		if (IS_LEFT_CHILD(z, y)) {
			if (IS_LEFT_CHILD(y, x)) {
				rotate_right(tree, z);
				update_height_to_root(z);
			} else {
				rotate_double_right(tree, z);
			}
		} else {
			if (IS_LEFT_CHILD(y, x)) {
				rotate_double_left(tree, z);
			} else {
				rotate_left(tree, z);
				update_height_to_root(z);
			}
		}

	next_ancestor:
		x = y;
		y = z;
		z = z->n_parent;
	}
}

/* run after a delete operation. checks that the balance factor
   of the local subst is within the range -1 <= BF <= 1. if it
   is not, rotate the subst to restore balance.

   note that, unlike insert_fixup, multiple rotations may be required
   to restore balance after a node is deleted.

   this function depends on all nodes in the tree having
   correct fx_height values.

   @param w one of the following:
            - the parent of the node that was deleted if the node
              had no children.
            - the parent of the node that replaced the deleted node
              if the deleted node had two children.
            - the node that replaced the node that was deleted, if
              the node that was deleted had one child.
*/
static void delete_fixup(struct fx_bst *tree, struct fx_bst_node *w)
{
	struct fx_bst_node *z = w;

	while (z) {
		if (bf(z) > 1) {
			if (bf(z->n_right) >= 0) {
				rotate_left(tree, z);
				update_height_to_root(z);
			} else {
				rotate_double_left(tree, z);
			}
		} else if (bf(z) < -1) {
			if (bf(z->n_left) <= 0) {
				rotate_right(tree, z);
				update_height_to_root(z);
			} else {
				rotate_double_right(tree, z);
			}
		}

		z = z->n_parent;
	}
}

/* updates fx_height for all nodes between the inserted node and the root
   of the tree, and calls insert_fixup.

   @param node the node that was just inserted into the tree.
*/
void fx_bst_insert_fixup(struct fx_bst *tree, struct fx_bst_node *node)
{
	node->n_height = 0;

	struct fx_bst_node *cur = node;
	while (cur) {
		update_height(cur);
		cur = cur->n_parent;
	}

	insert_fixup(tree, node);
}

/* remove a node from a tree.

   this function assumes that `node` has no children, and therefore
   doesn't need to be replaced.

   updates fx_height for all nodes between `node` and the tree root.

   @param node the node to delete.
*/
static struct fx_bst_node *remove_node_with_no_children(
	struct fx_bst *tree, struct fx_bst_node *node)
{
	struct fx_bst_node *w = node->n_parent;
	struct fx_bst_node *p = node->n_parent;
	node->n_parent = NULL;

	if (!p) {
		tree->bst_root = NULL;
	} else if (IS_LEFT_CHILD(p, node)) {
		p->n_left = NULL;
	} else {
		p->n_right = NULL;
	}

	while (p) {
		update_height(p);
		p = p->n_parent;
	}

	return w;
}

/* remove a node from a tree.

   this function assumes that `node` has one child.
   the child of `node` is inherited by `node`'s parent, and `node` is removed.

   updates fx_height for all nodes between the node that replaced
   `node` and the tree root.

   @param node the node to delete.
*/
static struct fx_bst_node *replace_node_with_one_subst(
	struct fx_bst *tree, struct fx_bst_node *node)
{
	struct fx_bst_node *p = node->n_parent;
	struct fx_bst_node *z = NULL;

	if (HAS_LEFT_CHILD(node)) {
		z = node->n_left;
	} else {
		z = node->n_right;
	}

	struct fx_bst_node *w = z;
	if (!p) {
		tree->bst_root = z;
	} else if (IS_LEFT_CHILD(p, node)) {
		p->n_left = z;
	} else if (IS_RIGHT_CHILD(p, node)) {
		p->n_right = z;
	}

	z->n_parent = p;

	node->n_parent = NULL;
	node->n_left = node->n_right = NULL;

	while (z) {
		update_height(z);
		z = z->n_parent;
	}

	return w;
}

/* remove a node from a tree.

   this function assumes that `node` has two children.
   find the in-order successor Y of `node` (the largest node in `node`'s left
   sub-tree), removes `node` from the tree and moves Y to where `node` used to
   be.

   if Y has a child (it will never have more than one), have Y's parent inherit
   Y's child.

   updates fx_height for all nodes between the deepest node that was modified
   and the tree root.

   @param z the node to delete.
*/
static struct fx_bst_node *replace_node_with_two_substs(
	struct fx_bst *tree, struct fx_bst_node *z)
{
	/* x will replace z */
	struct fx_bst_node *x = z->n_left;

	while (x->n_right) {
		x = x->n_right;
	}

	/* y is the node that will replace x (if x has a left child) */
	struct fx_bst_node *y = x->n_left;

	/* w is the starting point for the height update and fixup */
	struct fx_bst_node *w = x;
	if (w->n_parent != z) {
		w = w->n_parent;
	}

	if (y) {
		w = y;
	}

	if (IS_LEFT_CHILD(x->n_parent, x)) {
		x->n_parent->n_left = y;
	} else if (IS_RIGHT_CHILD(x->n_parent, x)) {
		x->n_parent->n_right = y;
	}

	if (y) {
		y->n_parent = x->n_parent;
	}

	if (IS_LEFT_CHILD(z->n_parent, z)) {
		z->n_parent->n_left = x;
	} else if (IS_RIGHT_CHILD(z->n_parent, z)) {
		z->n_parent->n_right = x;
	}

	x->n_parent = z->n_parent;
	x->n_left = z->n_left;
	x->n_right = z->n_right;

	if (x->n_left) {
		x->n_left->n_parent = x;
	}

	if (x->n_right) {
		x->n_right->n_parent = x;
	}

	if (!x->n_parent) {
		tree->bst_root = x;
	}

	struct fx_bst_node *cur = w;
	while (cur) {
		update_height(cur);
		cur = cur->n_parent;
	}

	return w;
}

/* delete a node from the tree and re-balance it afterwards */
void fx_bst_delete(struct fx_bst *tree, struct fx_bst_node *node)
{
	struct fx_bst_node *w = NULL;

	if (HAS_NO_CHILDREN(node)) {
		w = remove_node_with_no_children(tree, node);
	} else if (HAS_ONE_CHILD(node)) {
		w = replace_node_with_one_subst(tree, node);
	} else if (HAS_TWO_CHILDREN(node)) {
		w = replace_node_with_two_substs(tree, node);
	}

	if (w) {
		delete_fixup(tree, w);
	}

	node->n_left = node->n_right = node->n_parent = NULL;
}

static struct fx_bst_node *first_node(const struct fx_bst *tree, int *depth)
{
	/* the first node in the tree is the node with the smallest key.
	   we keep moving left until we can't go any further */
	struct fx_bst_node *cur = tree->bst_root;
	int d = 0;

	if (!cur) {
		*depth = 0;
		return NULL;
	}

	while (cur->n_left) {
		d++;
		cur = cur->n_left;
	}

	*depth = d;
	return cur;
}

struct fx_bst_node *fx_bst_first(const struct fx_bst *tree)
{
	int d;
	return first_node(tree, &d);
}

static struct fx_bst_node *last_node(const struct fx_bst *tree, int *depth)
{
	/* the first node in the tree is the node with the largest key.
	   we keep moving right until we can't go any further */
	struct fx_bst_node *cur = tree->bst_root;
	int d = 0;
	if (!cur) {
		return NULL;
	}

	while (cur->n_right) {
		d++;
		cur = cur->n_right;
	}

	*depth = d;
	return cur;
}

fx_bst_node *fx_bst_last(const struct fx_bst *tree)
{
	int d;
	return last_node(tree, &d);
}

static fx_bst_node *next_node(const struct fx_bst_node *node, int *depth_diff)
{
	if (!node) {
		return NULL;
	}

	int depth = 0;

	/* there are two possibilities for the next node:

	   1. if `node` has a right sub-tree, every node in this sub-tree is
	   bigger than node. the in-order successor of `node` is the smallest
	   node in this subst.
	   2. if `node` has no right sub-tree, we've reached the largest node in
	      the sub-tree rooted at `node`. we need to go back to our parent
	      and continue the search elsewhere.
	*/
	if (node->n_right) {
		/* case 1: step into `node`'s right sub-tree and keep going
		   left to find the smallest node */
		struct fx_bst_node *cur = node->n_right;
		depth++;
		while (cur->n_left) {
			cur = cur->n_left;
			depth++;
		}

		*depth_diff = depth;
		return cur;
	}

	/* case 2: keep stepping back up towards the root of the tree.
	   if we encounter a step where we are our parent's left child,
	   we've found a parent with a value larger than us. this parent
	   is the in-order successor of `node` */
	while (node->n_parent && node->n_parent->n_left != node) {
		node = node->n_parent;
		depth--;
	}

	*depth_diff = depth - 1;
	return node->n_parent;
}

static fx_bst_node *prev_node(const struct fx_bst_node *node, int *depth_diff)
{
	if (!node) {
		return NULL;
	}

	int depth = 0;

	/* there are two possibilities for the previous node:

	   1. if `node` has a left sub-tree, every node in this sub-tree is
	   smaller than `node`. the in-order predecessor of `node` is the
	   largest node in this subst.
	   2. if `node` has no left sub-tree, we've reached the smallest node in
	      the sub-tree rooted at `node`. we need to go back to our parent
	      and continue the search elsewhere.
	*/
	if (node->n_left) {
		/* case 1: step into `node`'s left sub-tree and keep going
		   right to find the largest node */
		fx_bst_node *cur = node->n_left;
		depth++;
		while (cur->n_right) {
			cur = cur->n_right;
			depth++;
		}

		*depth_diff = depth;
		return cur;
	}

	/* case 2: keep stepping back up towards the root of the tree.
	   if we encounter a step where we are our parent's right child,
	   we've found a parent with a value smaller than us. this parent
	   is the in-order predecessor of `node`. */
	while (node->n_parent && node->n_parent->n_right != node) {
		node = node->n_parent;
		depth--;
	}

	*depth_diff = depth - 1;
	return node->n_parent;
}

fx_bst_node *fx_bst_next(const struct fx_bst_node *node)
{
	int d;
	return next_node(node, &d);
}

fx_bst_node *fx_bst_prev(const struct fx_bst_node *node)
{
	int d;
	return prev_node(node, &d);
}

void fx_bst_move(
	struct fx_bst *tree, struct fx_bst_node *dest, struct fx_bst_node *src)
{
	if (src->n_parent) {
		if (src->n_parent->n_left == src) {
			src->n_parent->n_left = dest;
		} else {
			src->n_parent->n_right = dest;
		}
	}

	if (src->n_left) {
		src->n_left->n_parent = dest;
	}

	if (src->n_right) {
		src->n_right->n_parent = dest;
	}

	if (tree->bst_root == src) {
		tree->bst_root = dest;
	}

	memmove(dest, src, sizeof *src);
}

fx_iterator *fx_bst_begin(struct fx_bst *tree)
{
	fx_iterator *it_obj = fx_object_create(FX_TYPE_BST_ITERATOR);
	if (!it_obj) {
		return NULL;
	}

	struct fx_bst_iterator_p *it
		= fx_object_get_private(it_obj, FX_TYPE_BST_ITERATOR);
	int depth = 0;

	it->_b = (struct fx_bst *)tree;
	it->i = 0;
	it->node = first_node(tree, &depth);
	it->depth = depth;

	return it_obj;
}

const fx_iterator *fx_bst_cbegin(const struct fx_bst *tree)
{
	fx_iterator *it_obj = fx_object_create(FX_TYPE_BST_ITERATOR);
	if (!it_obj) {
		return NULL;
	}

	struct fx_bst_iterator_p *it
		= fx_object_get_private(it_obj, FX_TYPE_BST_ITERATOR);
	int depth = 0;

	it->_b = (struct fx_bst *)tree;
	it->i = 0;
	it->node = first_node(tree, &depth);
	it->depth = depth;

	return it_obj;
}

static enum fx_status iterator_move_next(const fx_iterator *obj)
{
	struct fx_bst_iterator_p *it
		= fx_object_get_private(obj, FX_TYPE_BST_ITERATOR);

	int depth_diff = 0;
	struct fx_bst_node *next = next_node(it->node, &depth_diff);

	if (!next) {
		it->node = NULL;
		it->depth = 0;
		it->i++;
		return false;
	}

	it->node = next;
	it->i++;
	it->depth += depth_diff;
	return true;
}

static enum fx_status iterator_erase(fx_iterator *obj)
{
	struct fx_bst_iterator_p *it
		= fx_object_get_private(obj, FX_TYPE_BST_ITERATOR);

	if (!it->node) {
		return FX_ERR_OUT_OF_BOUNDS;
	}

	int depth_diff = 0;
	struct fx_bst_node *next = next_node(it->node, &depth_diff);

	fx_bst_delete(it->_b, it->node);
	if (!next) {
		it->node = NULL;
		it->depth = 0;
	} else {
		it->node = next;
		it->depth = 0;

		struct fx_bst_node *cur = next->n_parent;
		while (cur) {
			it->depth++;
			cur = cur->n_parent;
		}
	}

	return FX_SUCCESS;
}

static fx_iterator_value iterator_get_value(fx_iterator *obj)
{
	struct fx_bst_iterator_p *it
		= fx_object_get_private(obj, FX_TYPE_BST_ITERATOR);

	return FX_ITERATOR_VALUE_PTR(it->node);
}

static const fx_iterator_value iterator_get_cvalue(const fx_iterator *obj)
{
	struct fx_bst_iterator_p *it
		= fx_object_get_private(obj, FX_TYPE_BST_ITERATOR);

	return FX_ITERATOR_VALUE_CPTR(it->node);
}

/*** CLASS DEFINITION *********************************************************/

// ---- fx_bst_iterator DEFINITION
FX_TYPE_CLASS_DEFINITION_BEGIN(fx_bst_iterator)
	FX_TYPE_CLASS_INTERFACE_BEGIN(fx_object, FX_TYPE_OBJECT)
		FX_INTERFACE_ENTRY(to_string) = NULL;
	FX_TYPE_CLASS_INTERFACE_END(fx_object, FX_TYPE_OBJECT)

	FX_TYPE_CLASS_INTERFACE_BEGIN(fx_iterator, FX_TYPE_ITERATOR)
		FX_INTERFACE_ENTRY(it_move_next) = iterator_move_next;
		FX_INTERFACE_ENTRY(it_erase) = iterator_erase;
		FX_INTERFACE_ENTRY(it_get_value) = iterator_get_value;
		FX_INTERFACE_ENTRY(it_get_cvalue) = iterator_get_cvalue;
	FX_TYPE_CLASS_INTERFACE_END(fx_iterator, FX_TYPE_ITERATOR)
FX_TYPE_CLASS_DEFINITION_END(fx_bst_iterator)

FX_TYPE_DEFINITION_BEGIN(fx_bst_iterator)
	FX_TYPE_ID(0x432779d7, 0xc03a, 0x48ea, 0xae8f, 0x12c666c767ae);
	FX_TYPE_EXTENDS(FX_TYPE_ITERATOR);
	FX_TYPE_CLASS(fx_bst_iterator_class);
	FX_TYPE_INSTANCE_PRIVATE(struct fx_bst_iterator_p);
FX_TYPE_DEFINITION_END(fx_bst_iterator)