dtc/src/agent/common/da_rbtree.c

/*
* Copyright [2021] JD.com, Inc.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
*     http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include "da_rbtree.h"
#include "da_log.h"

void rbtree_node_init(struct rbnode *node) {
	node->left = NULL;
	node->right = NULL;
	node->parent = NULL;
	node->key = 0ULL;
	node->data = NULL;
	/* color is left uninitialized */
}

void rbtree_init(struct rbtree *tree, struct rbnode *node) {
	rbtree_node_init(node);
	rbtree_black(node);
	tree->root = node;
	tree->sentinel = node;
}

static struct rbnode *
rbtree_node_min(struct rbnode *node, struct rbnode *sentinel) {
	/* traverse left links */

	while (node->left != sentinel) {
		node = node->left;
	}

	return node;
}

struct rbnode *rbtree_min(struct rbtree *tree) {
	struct rbnode *node = tree->root;
	struct rbnode *sentinel = tree->sentinel;

	/* empty tree */

	if (node == sentinel) {
		return NULL;
	}

	return rbtree_node_min(node, sentinel);
}

struct rbnode *rbtree_search(struct rbtree *tree,struct rbnode *tnode) {
	struct rbnode *node = tree->root;
	struct rbnode *sentinel = tree->sentinel;

	if (node == sentinel) {
		return NULL;
	}
	while (node != sentinel)
	{
		if(!node || !tnode)
			return NULL;

		if(tnode->key == node->key)
		{
			break;
		}
		else if(tnode->key > node->key)
		{
			node = node->right;
		}
		else if(tnode->key < node->key)
		{
			node = node->left;
		}
	}
	if(node==sentinel)
	{
		return NULL;
	}
	else
	{
		return node;
	}
}

static void rbtree_left_rotate(struct rbnode **root, struct rbnode *sentinel,
		struct rbnode *node) {
	struct rbnode *temp;

	temp = node->right;
	node->right = temp->left;

	if (temp->left != sentinel) {
		temp->left->parent = node;
	}

	temp->parent = node->parent;

	if (node == *root) {
		*root = temp;
	} else if (node == node->parent->left) {
		node->parent->left = temp;
	} else {
		node->parent->right = temp;
	}

	temp->left = node;
	node->parent = temp;
}

static void rbtree_right_rotate(struct rbnode **root, struct rbnode *sentinel,
		struct rbnode *node) {
	struct rbnode *temp;

	temp = node->left;
	node->left = temp->right;

	if (temp->right != sentinel) {
		temp->right->parent = node;
	}

	temp->parent = node->parent;

	if (node == *root) {
		*root = temp;
	} else if (node == node->parent->right) {
		node->parent->right = temp;
	} else {
		node->parent->left = temp;
	}

	temp->right = node;
	node->parent = temp;
}

void rbtree_insert(struct rbtree *tree, struct rbnode *node) {
	struct rbnode **root = &tree->root;
	struct rbnode *sentinel = tree->sentinel;
	struct rbnode *temp, **p;

	/* empty tree */

	if (*root == sentinel) {
		node->parent = NULL;
		node->left = sentinel;
		node->right = sentinel;
		rbtree_black(node);
		*root = node;
		return;
	}

	/* a binary tree insert */

	temp = *root;
	for (;;) {

		p = (node->key < temp->key) ? &temp->left : &temp->right;
		if (*p == sentinel) {
			break;
		}
		temp = *p;
	}

	*p = node;
	node->parent = temp;
	node->left = sentinel;
	node->right = sentinel;
	rbtree_red(node);

	/* re-balance tree */

	while (node != *root && rbtree_is_red(node->parent)) {

		if (node->parent == node->parent->parent->left) {
			temp = node->parent->parent->right;

			if (rbtree_is_red(temp)) {
				rbtree_black(node->parent);
				rbtree_black(temp);
				rbtree_red(node->parent->parent);
				node = node->parent->parent;
			} else {
				if (node == node->parent->right) {
					node = node->parent;
					rbtree_left_rotate(root, sentinel, node);
				}

				rbtree_black(node->parent);
				rbtree_red(node->parent->parent);
				rbtree_right_rotate(root, sentinel, node->parent->parent);
			}
		} else {
			temp = node->parent->parent->left;

			if (rbtree_is_red(temp)) {
				rbtree_black(node->parent);
				rbtree_black(temp);
				rbtree_red(node->parent->parent);
				node = node->parent->parent;
			} else {
				if (node == node->parent->left) {
					node = node->parent;
					rbtree_right_rotate(root, sentinel, node);
				}

				rbtree_black(node->parent);
				rbtree_red(node->parent->parent);
				rbtree_left_rotate(root, sentinel, node->parent->parent);
			}
		}
	}

	rbtree_black(*root);
}

void rbtree_delete(struct rbtree *tree, struct rbnode *node) {
	struct rbnode **root = &tree->root;
	struct rbnode *sentinel = tree->sentinel;
	struct rbnode *subst, *temp, *w;
	uint8_t red;

	/* a binary tree delete */

	if (node->left == sentinel) {
		temp = node->right;
		subst = node;
	} else if (node->right == sentinel) {
		temp = node->left;
		subst = node;
	} else {
		subst = rbtree_node_min(node->right, sentinel);
		if (subst->left != sentinel) {
			temp = subst->left;
		} else {
			temp = subst->right;
		}
	}

	if (subst == *root) {
		*root = temp;
		rbtree_black(temp);

		rbtree_node_init(node);

		return;
	}

	red = rbtree_is_red(subst);

	if (subst == subst->parent->left) {
		subst->parent->left = temp;
	} else {
		subst->parent->right = temp;
	}

	if (subst == node) {
		temp->parent = subst->parent;
	} else {

		if (subst->parent == node) {
			temp->parent = subst;
		} else {
			temp->parent = subst->parent;
		}

		subst->left = node->left;
		subst->right = node->right;
		subst->parent = node->parent;
		rbtree_copy_color(subst, node);

		if (node == *root) {
			*root = subst;
		} else {
			if (node == node->parent->left) {
				node->parent->left = subst;
			} else {
				node->parent->right = subst;
			}
		}

		if (subst->left != sentinel) {
			subst->left->parent = subst;
		}

		if (subst->right != sentinel) {
			subst->right->parent = subst;
		}
	}

	rbtree_node_init(node);

	if (red) {
		return;
	}

	/* a delete fixup */

	while (temp != *root && rbtree_is_black(temp)) {

		if (temp == temp->parent->left) {
			w = temp->parent->right;

			if (rbtree_is_red(w)) {
				rbtree_black(w);
				rbtree_red(temp->parent);
				rbtree_left_rotate(root, sentinel, temp->parent);
				w = temp->parent->right;
			}

			if (rbtree_is_black(w->left) && rbtree_is_black(w->right)) {
				rbtree_red(w);
				temp = temp->parent;
			} else {
				if (rbtree_is_black(w->right)) {
					rbtree_black(w->left);
					rbtree_red(w);
					rbtree_right_rotate(root, sentinel, w);
					w = temp->parent->right;
				}

				rbtree_copy_color(w, temp->parent);
				rbtree_black(temp->parent);
				rbtree_black(w->right);
				rbtree_left_rotate(root, sentinel, temp->parent);
				temp = *root;
			}

		} else {
			w = temp->parent->left;

			if (rbtree_is_red(w)) {
				rbtree_black(w);
				rbtree_red(temp->parent);
				rbtree_right_rotate(root, sentinel, temp->parent);
				w = temp->parent->left;
			}

			if (rbtree_is_black(w->left) && rbtree_is_black(w->right)) {
				rbtree_red(w);
				temp = temp->parent;
			} else {
				if (rbtree_is_black(w->left)) {
					rbtree_black(w->right);
					rbtree_red(w);
					rbtree_left_rotate(root, sentinel, w);
					w = temp->parent->left;
				}

				rbtree_copy_color(w, temp->parent);
				rbtree_black(temp->parent);
				rbtree_black(w->left);
				rbtree_right_rotate(root, sentinel, temp->parent);
				temp = *root;
			}
		}
	}

	rbtree_black(temp);
}