mango/include/socks/btree.h

/*
  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.
 */

#ifndef SOCKS_BTREE_H_
#define SOCKS_BTREE_H_

#include <stdint.h>

#ifdef __cplusplus
extern "C" {
#endif

/* if your custom structure contains a btree_node_t (i.e. it can be part of a btree),
   you can use this macro to convert a btree_node_t* to a your_type*

   @param t the name of your custom type (something that can be passed to offsetof)
   @param m the name of the btree_node_t member variable within your custom type.
   @param v the btree_node_t pointer that you wish to convert. if this is NULL, NULL will be returned.
*/
#define BTREE_CONTAINER(t, m, v) ((void *)((v) ? (uintptr_t)(v) - (offsetof(t, m)) : 0))

/* defines a simple node insertion function.
   this function assumes that your nodes have simple integer keys that can be compared with the usual operators.

   EXAMPLE:
   if you have a tree node type like this:

     struct my_tree_node {
             int key;
	     btree_node_t base;
     }

   You would use the following call to generate an insert function for a tree with this node type:

     BTREE_DEFINE_SIMPLE_INSERT(struct my_tree_node, base, key, my_tree_node_insert);

   Which would emit a function defined like:

     static void my_tree_node_insert(btree_t *tree, struct my_tree_node *node);

   @param node_type your custom tree node type. usually a structure that contains a btree_node_t member.
   @param container_node_member the name of the btree_node_t member variable within your custom type.
   @param container_key_member the name of the key member variable within your custom type.
   @param function_name the name of the function to generate.
*/
#define BTREE_DEFINE_SIMPLE_INSERT(node_type, container_node_member, container_key_member, function_name) \
	static void function_name(btree_t *tree, node_type *node)                                         \
	{                                                                                                 \
		if (!tree->b_root) {                                                                      \
			tree->b_root = &node->container_node_member;                                      \
			btree_insert_fixup(tree, &node->container_node_member);                           \
			return;                                                                           \
		}                                                                                         \
                                                                                                          \
		btree_node_t *cur = tree->b_root;                                                         \
		while (1) {                                                                               \
			node_type *cur_node = BTREE_CONTAINER(node_type, container_node_member, cur);     \
			btree_node_t *next = NULL;                                                        \
                                                                                                          \
			if (node->container_key_member > cur_node->container_key_member) {                \
				next = btree_right(cur);                                                  \
                                                                                                          \
				if (!next) {                                                              \
					btree_put_right(cur, &node->container_node_member);               \
					break;                                                            \
				}                                                                         \
			} else if (node->container_key_member < cur_node->container_key_member) {         \
				next = btree_left(cur);                                                   \
                                                                                                          \
				if (!next) {                                                              \
					btree_put_left(cur, &node->container_node_member);                \
					break;                                                            \
				}                                                                         \
			} else {                                                                          \
				return;                                                                   \
			}                                                                                 \
                                                                                                          \
			cur = next;                                                                       \
		}                                                                                         \
                                                                                                          \
		btree_insert_fixup(tree, &node->container_node_member);                                   \
	}

/* defines a node insertion function.
   this function should be used for trees with complex node keys that cannot be directly compared.
   a comparator for your keys must be supplied.

   EXAMPLE:
   if you have a tree node type like this:

     struct my_tree_node {
             complex_key_t key;
	     btree_node_t base;
     }

   You would need to define a comparator function or macro with the following signature:

     int my_comparator(struct my_tree_node *a, struct my_tree_node *b);

   Which implements the following:

     return -1   if a < b
     return  0   if a == b
     return  1   if a > b

   You would use the following call to generate an insert function for a tree with this node type:

     BTREE_DEFINE_INSERT(struct my_tree_node, base, key, my_tree_node_insert, my_comparator);

   Which would emit a function defined like:

     static void my_tree_node_insert(btree_t *tree, struct my_tree_node *node);

   @param node_type your custom tree node type. usually a structure that contains a btree_node_t member.
   @param container_node_member the name of the btree_node_t member variable within your custom type.
   @param container_key_member the name of the key member variable within your custom type.
   @param function_name the name of the function to generate.
   @param comparator the name of a comparator function or functional-macro that conforms to the
                     requirements listed above.
*/
#define BTREE_DEFINE_INSERT(node_type, container_node_member, container_key_member, function_name, comparator) \
	static void function_name(btree_t *tree, node_type *node)                                              \
	{                                                                                                      \
		if (!tree->b_root) {                                                                           \
			tree->b_root = &node->container_node_member;                                           \
			btree_insert_fixup(tree, &node->container_node_member);                                \
			return;                                                                                \
		}                                                                                              \
                                                                                                               \
		btree_node_t *cur = tree->b_root;                                                              \
		while (1) {                                                                                    \
			node_type *cur_node = BTREE_CONTAINER(node_type, container_node_member, cur);          \
			btree_node_t *next = NULL;                                                             \
			int cmp = comparator(node, cur_node);                                                  \
                                                                                                               \
			if (cmp == 1) {                                                                        \
				next = btree_right(cur);                                                       \
                                                                                                               \
				if (!next) {                                                                   \
					btree_put_right(cur, &node->container_node_member);                    \
					break;                                                                 \
				}                                                                              \
			} else if (cmp == -1) {                                                                \
				next = btree_left(cur);                                                        \
                                                                                                               \
				if (!next) {                                                                   \
					btree_put_left(cur, &node->container_node_member);                     \
					break;                                                                 \
				}                                                                              \
			} else {                                                                               \
				return;                                                                        \
			}                                                                                      \
                                                                                                               \
			cur = next;                                                                            \
		}                                                                                              \
                                                                                                               \
		btree_insert_fixup(tree, &node->container_node_member);                                        \
	}

/* defines a simple tree search function.
   this function assumes that your nodes have simple integer keys that can be compared with the usual operators.

   EXAMPLE:
   if you have a tree node type like this:

     struct my_tree_node {
             int key;
	     btree_node_t base;
     }

   You would use the following call to generate a search function for a tree with this node type:

     BTREE_DEFINE_SIMPLE_GET(struct my_tree_node, int, base, key, my_tree_node_get);

   Which would emit a function defined like:

     static struct my_tree_node *my_tree_node_get(btree_t *tree, int key);

   @param node_type your custom tree node type. usually a structure that contains a btree_node_t member.
   @param key_type the type name of the key embedded in your custom tree node type. this type must be
                   compatible with the builtin comparison operators.
   @param container_node_member the name of the btree_node_t member variable within your custom type.
   @param container_key_member the name of the key member variable within your custom type.
   @param function_name the name of the function to generate.
*/
#define BTREE_DEFINE_SIMPLE_GET(node_type, key_type, container_node_member, container_key_member, function_name) \
node_type *function_name(btree_t *tree, key_type key) \
{ \
	btree_node_t *cur = tree->b_root; \
	while (cur) { \
		node_type *cur_node = BTREE_CONTAINER(node_type, container_node_member, cur); \
		if (key > cur_node->container_key_member) { \
			cur = btree_right(cur); \
		} else if (key < cur_node->container_key_member) { \
			cur = btree_left(cur); \
		} else { \
			return cur_node; \
		} \
	} \
 \
	return NULL; \
}

/* perform an in-order traversal of a binary tree

   If you have a tree defined like:

     btree_t my_tree;

   with nodes defined like:

     struct my_tree_node {
             int key;
	     btree_node_t base;
     }

   and you want to do something like:

     foreach (struct my_tree_node *node : my_tree) { ... }

   you should use this:

     btree_foreach (struct my_tree_node, node, &my_tree, base) { ... }

   @param iter_type the type name of the iterator variable. this should be the tree's node type, and shouldn't be a pointer.
   @param iter_name the name of the iterator variable.
   @param tree_name a pointer to the tree to traverse.
   @param node_member the name of the btree_node_t member variable within the tree node type.
*/
#define btree_foreach(iter_type, iter_name, tree_name, node_member) \
	for (iter_type *iter_name = BTREE_CONTAINER(iter_type, node_member, btree_first(tree_name)); \
	     iter_name; \
	     iter_name = BTREE_CONTAINER(iter_type, node_member, btree_next(&((iter_name)->node_member))))

/* perform an reverse in-order traversal of a binary tree

   If you have a tree defined like:

     btree_t my_tree;

   with nodes defined like:

     struct my_tree_node {
             int key;
	     btree_node_t base;
     }

   and you want to do something like:

     foreach (struct my_tree_node *node : reverse(my_tree)) { ... }

   you should use this:

     btree_foreach_r (struct my_tree_node, node, &my_tree, base) { ... }

   @param iter_type the type name of the iterator variable. this should be the tree's node type, and shouldn't be a pointer.
   @param iter_name the name of the iterator variable.
   @param tree_name a pointer to the tree to traverse.
   @param node_member the name of the btree_node_t member variable within the tree node type.
*/
#define btree_foreach_r(iter_type, iter_name, tree_name, node_member) \
	for (iter_type *iter_name = BTREE_CONTAINER(iter_type, node_member, btree_last(tree_name)); \
	     iter_name; \
	     iter_name = BTREE_CONTAINER(iter_type, node_member, btree_prev(&((iter_name)->node_member))))

/* binary tree nodes. this *cannot* be used directly. you need to define a custom node type
   that contains a member variable of type btree_node_t.

   you would then use the supplied macros to define functions to manipulate your custom binary tree.
*/
typedef struct btree_node {
	struct btree_node *b_parent, *b_left, *b_right;
	unsigned short b_height;
} btree_node_t;

/* binary tree. unlike btree_node_t, you can define variables of type btree_t. */
typedef struct btree {
	struct btree_node *b_root;
} btree_t;

/* re-balance a binary tree after an insertion operation.

   NOTE that, if you define an insertion function using BTREE_DEFINE_INSERT or similar,
   this function will automatically called for you.

   @param tree the tree to re-balance.
   @param node the node that was just inserted into the tree.
*/
extern void btree_insert_fixup(btree_t *tree, btree_node_t *node);

/* delete a node from a binary tree and re-balance the tree afterwards.

   @param tree the tree to delete from
   @param node the node to delete.
*/
extern void btree_delete(btree_t *tree, btree_node_t *node);

/* get the first node in a binary tree.

   this will be the node with the smallest key (i.e. the node that is furthest-left from the root)
*/
extern btree_node_t *btree_first(btree_t *tree);

/* get the last node in a binary tree.

   this will be the node with the largest key (i.e. the node that is furthest-right from the root)
*/
extern btree_node_t *btree_last(btree_t *tree);
/* for any binary tree node, this function returns the node with the next-largest key value */
extern btree_node_t *btree_next(btree_node_t *node);
/* for any binary tree node, this function returns the node with the next-smallest key value */
extern btree_node_t *btree_prev(btree_node_t *node);

/* sets `child` as the immediate left-child of `parent` */
static inline void btree_put_left(btree_node_t *parent, btree_node_t *child)
{
	parent->b_left = child;
	child->b_parent = parent;
}

/* sets `child` as the immediate right-child of `parent` */
static inline void btree_put_right(btree_node_t *parent, btree_node_t *child)
{
	parent->b_right = child;
	child->b_parent = parent;
}

/* get the immediate left-child of `node` */
static inline btree_node_t *btree_left(btree_node_t *node)
{
	return node->b_left;
}

/* get the immediate right-child of `node` */
static inline btree_node_t *btree_right(btree_node_t *node)
{
	return node->b_right;
}

/* get the immediate parent of `node` */
static inline btree_node_t *btree_parent(btree_node_t *node)
{
	return node->b_parent;
}

/* get the height of `node`.

   the height of a node is defined as the length of the longest path
   between the node and a leaf node.

   this count includes the node itself, so the height of a leaf node will be 1.
*/
static inline unsigned short btree_height(btree_node_t *node)
{
	return node->b_height;
}

#ifdef __cplusplus
}
#endif

#endif