#include <assert.h>
#include <ofi_tree.h>
#include <ofi_osd.h>
#include <rdma/fi_errno.h>
static struct ofi_rbnode *ofi_rbnode_alloc(struct ofi_rbmap *map)
{
struct ofi_rbnode *node;
if (!map->free_list)
return malloc(sizeof(*node));
node = map->free_list;
map->free_list = node->right;
return node;
}
static void ofi_rbnode_free(struct ofi_rbmap *map, struct ofi_rbnode *node)
{
node->right = map->free_list ? map->free_list : NULL;
map->free_list = node;
}
void ofi_rbmap_init(struct ofi_rbmap *map,
int (*compare)(struct ofi_rbmap *map, void *key, void *data))
{
map->compare = compare;
map->root = &map->sentinel;
map->sentinel.left = &map->sentinel;
map->sentinel.right = &map->sentinel;
map->sentinel.parent = NULL;
map->sentinel.color = BLACK;
map->sentinel.data = NULL;
}
struct ofi_rbmap *
ofi_rbmap_create(int (*compare)(struct ofi_rbmap *map, void *key, void *data))
{
struct ofi_rbmap *map;
map = calloc(1, sizeof *map);
if (map)
ofi_rbmap_init(map, compare);
return map;
}
static void ofi_delete_tree(struct ofi_rbmap *map, struct ofi_rbnode *node)
{
if (node == &map->sentinel)
return;
ofi_delete_tree(map, node->left);
ofi_delete_tree(map, node->right);
free(node);
}
void ofi_rbmap_cleanup(struct ofi_rbmap *map)
{
struct ofi_rbnode *node;
ofi_delete_tree(map, map->root);
while (map->free_list) {
node = map->free_list;
map->free_list = node->right;
free(node);
}
}
void ofi_rbmap_destroy(struct ofi_rbmap *map)
{
ofi_rbmap_cleanup(map);
free(map);
}
int ofi_rbmap_empty(struct ofi_rbmap *map)
{
return map->root == &map->sentinel;
}
static void ofi_rotate_left(struct ofi_rbmap *map, struct ofi_rbnode *node)
{
struct ofi_rbnode *y = node->right;
node->right = y->left;
if (y->left != &map->sentinel)
y->left->parent = node;
if (y != &map->sentinel)
y->parent = node->parent;
if (node->parent) {
if (node== node->parent->left)
node->parent->left = y;
else
node->parent->right = y;
} else {
map->root = y;
}
y->left = node;
if (node != &map->sentinel)
node->parent = y;
}
static void ofi_rotate_right(struct ofi_rbmap *map, struct ofi_rbnode *node)
{
struct ofi_rbnode *y = node->left;
node->left = y->right;
if (y->right != &map->sentinel)
y->right->parent = node;
if (y != &map->sentinel)
y->parent = node->parent;
if (node->parent) {
if (node == node->parent->right)
node->parent->right = y;
else
node->parent->left = y;
} else {
map->root = y;
}
y->right = node;
if (node != &map->sentinel)
node->parent = y;
}
static void
ofi_insert_rebalance(struct ofi_rbmap *map, struct ofi_rbnode *x)
{
struct ofi_rbnode *y;
while (x != map->root && x->parent->color == RED) {
if (x->parent == x->parent->parent->left) {
y = x->parent->parent->right;
if (y->color == RED) {
x->parent->color = BLACK;
y->color = BLACK;
x->parent->parent->color = RED;
x = x->parent->parent;
} else {
if (x == x->parent->right) {
x = x->parent;
ofi_rotate_left(map, x);
}
x->parent->color = BLACK;
x->parent->parent->color = RED;
ofi_rotate_right(map, x->parent->parent);
}
} else {
y = x->parent->parent->left;
if (y->color == RED) {
x->parent->color = BLACK;
y->color = BLACK;
x->parent->parent->color = RED;
x = x->parent->parent;
} else {
if (x == x->parent->left) {
x = x->parent;
ofi_rotate_right(map, x);
}
x->parent->color = BLACK;
x->parent->parent->color = RED;
ofi_rotate_left(map, x->parent->parent);
}
}
}
map->root->color = BLACK;
}
int ofi_rbmap_insert(struct ofi_rbmap *map, void *key, void *data,
struct ofi_rbnode **ret_node)
{
struct ofi_rbnode *current, *parent, *node;
int ret;
current = map->root;
parent = NULL;
while (current != &map->sentinel) {
ret = map->compare(map, key, current->data);
if (ret == 0) {
if (ret_node)
*ret_node = current;
return -FI_EALREADY;
}
parent = current;
current = (ret < 0) ? current->left : current->right;
}
node = ofi_rbnode_alloc(map);
if (!node)
return -FI_ENOMEM;
node->parent = parent;
node->left = &map->sentinel;
node->right = &map->sentinel;
node->color = RED;
node->data = data;
if (parent) {
if (map->compare(map, key, parent->data) < 0)
parent->left = node;
else
parent->right = node;
} else {
map->root = node;
}
ofi_insert_rebalance(map, node);
if (ret_node)
*ret_node = node;
return 0;
}
static void ofi_delete_rebalance(struct ofi_rbmap *map, struct ofi_rbnode *node)
{
struct ofi_rbnode *w;
while (node != map->root && node->color == BLACK) {
if (node == node->parent->left) {
w = node->parent->right;
if (w->color == RED) {
w->color = BLACK;
node->parent->color = RED;
ofi_rotate_left(map, node->parent);
w = node->parent->right;
}
if (w->left->color == BLACK && w->right->color == BLACK) {
w->color = RED;
node = node->parent;
} else {
if (w->right->color == BLACK) {
w->left->color = BLACK;
w->color = RED;
ofi_rotate_right(map, w);
w = node->parent->right;
}
w->color = node->parent->color;
node->parent->color = BLACK;
w->right->color = BLACK;
ofi_rotate_left(map, node->parent);
node = map->root;
}
} else {
w = node->parent->left;
if (w->color == RED) {
w->color = BLACK;
node->parent->color = RED;
ofi_rotate_right(map, node->parent);
w = node->parent->left;
}
if (w->right->color == BLACK && w->left->color == BLACK) {
w->color = RED;
node = node->parent;
} else {
if (w->left->color == BLACK) {
w->right->color = BLACK;
w->color = RED;
ofi_rotate_left(map, w);
w = node->parent->left;
}
w->color = node->parent->color;
node->parent->color = BLACK;
w->left->color = BLACK;
ofi_rotate_right(map, node->parent);
node = map->root;
}
}
}
node->color = BLACK;
}
static void ofi_rbmap_replace_node_ptr(struct ofi_rbmap *map,
struct ofi_rbnode *old_node, struct ofi_rbnode *new_node)
{
if (new_node == old_node)
return;
*new_node = *old_node;
if (!old_node->parent)
map->root = new_node;
else if (old_node == old_node->parent->left)
old_node->parent->left = new_node;
else
old_node->parent->right = new_node;
if (old_node->left != &map->sentinel)
old_node->left->parent = new_node;
if (old_node->right != &map->sentinel)
old_node->right->parent = new_node;
}
void ofi_rbmap_delete(struct ofi_rbmap *map, struct ofi_rbnode *node)
{
struct ofi_rbnode *x, *y;
if (node->left == &map->sentinel) {
y = node;
x = y->right;
} else if (node->right == &map->sentinel) {
y = node;
x = y->left;
} else {
y = node->right;
while (y->left != &map->sentinel)
y = y->left;
x = y->right;
}
x->parent = y->parent;
if (y->parent) {
if (y == y->parent->left)
y->parent->left = x;
else
y->parent->right = x;
} else {
map->root = x;
}
if (y != node)
node->data = y->data;
if (y->color == BLACK)
ofi_delete_rebalance(map, x);
ofi_rbmap_replace_node_ptr(map, node, y);
ofi_rbnode_free(map, node);
}
struct ofi_rbnode *ofi_rbmap_get_root(struct ofi_rbmap *map)
{
if (ofi_rbmap_empty(map))
return NULL;
return map->root;
}
struct ofi_rbnode *ofi_rbmap_find(struct ofi_rbmap *map, void *key)
{
struct ofi_rbnode *node;
int ret;
node = map->root;
while (node != &map->sentinel) {
ret = map->compare(map, key, node->data);
if (ret == 0)
return node;
node = (ret < 0) ? node->left : node->right;
}
return NULL;
}
int ofi_rbmap_find_delete(struct ofi_rbmap *map, void *key)
{
struct ofi_rbnode *node;
node = ofi_rbmap_find(map, key);
if (!node)
return -FI_ENODATA;
ofi_rbmap_delete(map, node);
return 0;
}
struct ofi_rbnode *ofi_rbmap_search(struct ofi_rbmap *map, void *key,
int (*compare)(struct ofi_rbmap *map, void *key, void *data))
{
struct ofi_rbnode *node;
int ret;
node = map->root;
while (node != &map->sentinel) {
ret = compare(map, key, node->data);
if (ret == 0)
return node;
node = (ret < 0) ? node->left : node->right;
}
return NULL;
}