/*
* Copyright (c) 2017, NVIDIA CORPORATION. All rights reserved.
*
* This program is free software; you can redistribute it and/or modify it
* under the terms and conditions of the GNU General Public License,
* version 2, as published by the Free Software Foundation.
*
* This program is distributed in the hope it will be useful, but WITHOUT
* ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
* FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
* more details.
*
* You should have received a copy of the GNU General Public License
* along with this program. If not, see <http://www.gnu.org/licenses/>.
*/
#include <nvgpu/rbtree.h>
/*
* rotate node x to left
*/
static void rotate_left(struct nvgpu_rbtree_node **root,
struct nvgpu_rbtree_node *x)
{
struct nvgpu_rbtree_node *y = x->right;
/* establish x->right link */
x->right = y->left;
if (y->left)
y->left->parent = x;
/* establish y->parent link */
y->parent = x->parent;
if (x->parent) {
if (x == x->parent->left)
x->parent->left = y;
else
x->parent->right = y;
} else {
*root = y;
}
/* link x and y */
y->left = x;
x->parent = y;
}
/*
* rotate node x to right
*/
static void rotate_right(struct nvgpu_rbtree_node **root,
struct nvgpu_rbtree_node *x)
{
struct nvgpu_rbtree_node *y = x->left;
/* establish x->left link */
x->left = y->right;
if (y->right)
y->right->parent = x;
/* establish y->parent link */
y->parent = x->parent;
if (x->parent) {
if (x == x->parent->right)
x->parent->right = y;
else
x->parent->left = y;
} else {
*root = y;
}
/* link x and y */
y->right = x;
x->parent = y;
}
/*
* maintain red-black tree balance after inserting node x
*/
static void insert_fixup(struct nvgpu_rbtree_node **root,
struct nvgpu_rbtree_node *x)
{
/* check red-black properties */
while ((x != *root) && x->parent->is_red) {
/* we have a violation */
if (x->parent == x->parent->parent->left) {
struct nvgpu_rbtree_node *y = x->parent->parent->right;
if (y && y->is_red) {
/* uncle is RED */
x->parent->is_red = false;
y->is_red = false;
x->parent->parent->is_red = true;
x = x->parent->parent;
} else {
/* uncle is BLACK */
if (x == x->parent->right) {
/* make x a left child */
x = x->parent;
rotate_left(root, x);
}
/* recolor and rotate */
x->parent->is_red = false;
x->parent->parent->is_red = true;
rotate_right(root, x->parent->parent);
}
} else {
/* mirror image of above code */
struct nvgpu_rbtree_node *y = x->parent->parent->left;
if (y && y->is_red) {
/* uncle is RED */
x->parent->is_red = false;
y->is_red = false;
x->parent->parent->is_red = true;
x = x->parent->parent;
} else {
/* uncle is BLACK */
if (x == x->parent->left) {
x = x->parent;
rotate_right(root, x);
}
x->parent->is_red = false;
x->parent->parent->is_red = true;
rotate_left(root, x->parent->parent);
}
}
}
(*root)->is_red = false;
}
void nvgpu_rbtree_insert(struct nvgpu_rbtree_node *new_node,
struct nvgpu_rbtree_node **root)
{
struct nvgpu_rbtree_node *curr;
struct nvgpu_rbtree_node *parent;
/* find future parent */
curr = *root;
parent = NULL;
while (curr) {
parent = curr;
if (new_node->key_start < curr->key_start)
curr = curr->left;
else if (new_node->key_start > curr->key_start)
curr = curr->right;
else
return; /* duplicate entry */
}
/* the caller allocated the node already, just fix the links */
new_node->parent = parent;
new_node->left = NULL;
new_node->right = NULL;
new_node->is_red = true;
/* insert node in tree */
if (parent) {
if (new_node->key_start < parent->key_start)
parent->left = new_node;
else
parent->right = new_node;
} else {
*root = new_node;
}
insert_fixup(root, new_node);
}
/*
* maintain red-black tree balance after deleting node x
*/
static void _delete_fixup(struct nvgpu_rbtree_node **root,
struct nvgpu_rbtree_node *parent_of_x,
struct nvgpu_rbtree_node *x)
{
while ((x != *root) && (!x || !x->is_red)) {
/*
* NULL nodes are sentinel nodes. If we delete a sentinel
* node (x==NULL) it must have a parent node (or be the root).
* Hence, parent_of_x == NULL with
* x==NULL is never possible (tree invariant)
*/
if ((parent_of_x != NULL) && (x == parent_of_x->left)) {
struct nvgpu_rbtree_node *w = parent_of_x->right;
if (w && w->is_red) {
w->is_red = false;
parent_of_x->is_red = true;
rotate_left(root, parent_of_x);
w = parent_of_x->right;
}
if (!w || ((!w->left || !w->left->is_red)
&& (!w->right || !w->right->is_red))) {
if (w)
w->is_red = true;
x = parent_of_x;
} else {
if (!w->right || !w->right->is_red) {
w->left->is_red = false;
w->is_red = true;
rotate_right(root, w);
w = parent_of_x->right;
}
w->is_red = parent_of_x->is_red;
parent_of_x->is_red = false;
w->right->is_red = false;
rotate_left(root, parent_of_x);
x = *root;
}
} else if (parent_of_x != NULL) {
struct nvgpu_rbtree_node *w = parent_of_x->left;
if (w && w->is_red) {
w->is_red = false;
parent_of_x->is_red = true;
rotate_right(root, parent_of_x);
w = parent_of_x->left;
}
if (!w || ((!w->right || !w->right->is_red)
&& (!w->left || !w->left->is_red))) {
if (w)
w->is_red = true;
x = parent_of_x;
} else {
if (!w->left || !w->left->is_red) {
w->right->is_red = false;
w->is_red = true;
rotate_left(root, w);
w = parent_of_x->left;
}
w->is_red = parent_of_x->is_red;
parent_of_x->is_red = false;
w->left->is_red = false;
rotate_right(root, parent_of_x);
x = *root;
}
}
parent_of_x = x->parent;
}
if (x)
x->is_red = false;
}
void nvgpu_rbtree_unlink(struct nvgpu_rbtree_node *node,
struct nvgpu_rbtree_node **root)
{
struct nvgpu_rbtree_node *x;
struct nvgpu_rbtree_node *y;
struct nvgpu_rbtree_node *z;
struct nvgpu_rbtree_node *parent_of_x;
bool y_was_black;
z = node;
/* unlink */
if (!z->left || !z->right) {
/* y has a SENTINEL node as a child */
y = z;
} else {
/* find tree successor */
y = z->right;
while (y->left)
y = y->left;
}
/* x is y's only child */
if (y->left)
x = y->left;
else
x = y->right;
/* remove y from the parent chain */
parent_of_x = y->parent;
if (x)
x->parent = parent_of_x;
if (y->parent) {
if (y == y->parent->left)
y->parent->left = x;
else
y->parent->right = x;
} else {
*root = x;
}
y_was_black = !y->is_red;
if (y != z) {
/* we need to replace z with y so
* the memory for z can be freed
*/
y->parent = z->parent;
if (z->parent) {
if (z == z->parent->left)
z->parent->left = y;
else
z->parent->right = y;
} else {
*root = y;
}
y->is_red = z->is_red;
y->left = z->left;
if (z->left)
z->left->parent = y;
y->right = z->right;
if (z->right)
z->right->parent = y;
if (parent_of_x == z)
parent_of_x = y;
}
if (y_was_black)
_delete_fixup(root, parent_of_x, x);
}
void nvgpu_rbtree_search(u64 key_start, struct nvgpu_rbtree_node **node,
struct nvgpu_rbtree_node *root)
{
struct nvgpu_rbtree_node *curr = root;
while (curr) {
if (key_start < curr->key_start) {
curr = curr->left;
} else if (key_start > curr->key_start) {
curr = curr->right;
} else {
*node = curr;
return;
}
}
*node = NULL;
}
void nvgpu_rbtree_range_search(u64 key,
struct nvgpu_rbtree_node **node,
struct nvgpu_rbtree_node *root)
{
struct nvgpu_rbtree_node *curr = root;
while (curr) {
if (key >= curr->key_start &&
key < curr->key_end) {
*node = curr;
return;
} else if (key < curr->key_start) {
curr = curr->left;
} else {
curr = curr->right;
}
}
*node = NULL;
}
void nvgpu_rbtree_less_than_search(u64 key_start,
struct nvgpu_rbtree_node **node,
struct nvgpu_rbtree_node *root)
{
struct nvgpu_rbtree_node *curr = root;
while (curr) {
if (key_start <= curr->key_start) {
curr = curr->left;
} else {
*node = curr;
curr = curr->right;
}
}
}
void nvgpu_rbtree_enum_start(u64 key_start, struct nvgpu_rbtree_node **node,
struct nvgpu_rbtree_node *root)
{
*node = NULL;
if (root) {
struct nvgpu_rbtree_node *curr = root;
while (curr) {
if (key_start < curr->key_start) {
*node = curr;
curr = curr->left;
} else if (key_start > curr->key_start) {
curr = curr->right;
} else {
*node = curr;
break;
}
}
}
}
void nvgpu_rbtree_enum_next(struct nvgpu_rbtree_node **node,
struct nvgpu_rbtree_node *root)
{
struct nvgpu_rbtree_node *curr = NULL;
if (root && *node) {
/* if we don't have a right subtree return the parent */
curr = *node;
/* pick the leftmost node of the right subtree ? */
if (curr->right) {
curr = curr->right;
for (; curr->left;)
curr = curr->left;
} else {
/* go up until we find the right inorder node */
for (curr = curr->parent; curr; curr = curr->parent) {
if (curr->key_start > (*node)->key_start)
break;
}
}
}
*node = curr;
}
