#ifndef _KERNEL_LIBK_RBTREE_H
#define _KERNEL_LIBK_RBTREE_H

struct rb_node_link {
  struct rb_node_link* left;
  struct rb_node_link* right;
  struct rb_node_link* parent;
  int color;
};

#define RBTREE_RED   0
#define RBTREE_BLACK 1

#define rbtree_parent(x) ((x)->parent)
#define rbtree_left(x)   ((x)->left)
#define rbtree_right(x)  ((x)->right)
#define rbtree_color(x)  ((x)->color)

#define rbtree_container_of(ptr, type, member) ((type*)((char*)(ptr) - offsetof (type, member)))

#define rbtree_entry(node, type, member) rbtree_container_of (node, type, member)

#define rbtree_rotate_left(root_ptr, x)                                                            \
  do {                                                                                             \
    struct rb_node_link* __y = (x)->right;                                                         \
    (x)->right = __y->left;                                                                        \
    if (__y->left)                                                                                 \
      __y->left->parent = (x);                                                                     \
    __y->parent = (x)->parent;                                                                     \
    if (!(x)->parent)                                                                              \
      *(root_ptr) = __y;                                                                           \
    else if ((x) == (x)->parent->left)                                                             \
      (x)->parent->left = __y;                                                                     \
    else                                                                                           \
      (x)->parent->right = __y;                                                                    \
    __y->left = (x);                                                                               \
    (x)->parent = __y;                                                                             \
  } while (0)

#define rbtree_rotate_right(root_ptr, y)                                                           \
  do {                                                                                             \
    struct rb_node_link* __x = (y)->left;                                                          \
    (y)->left = __x->right;                                                                        \
    if (__x->right)                                                                                \
      __x->right->parent = (y);                                                                    \
    __x->parent = (y)->parent;                                                                     \
    if (!(y)->parent)                                                                              \
      *(root_ptr) = __x;                                                                           \
    else if ((y) == (y)->parent->right)                                                            \
      (y)->parent->right = __x;                                                                    \
    else                                                                                           \
      (y)->parent->left = __x;                                                                     \
    __x->right = (y);                                                                              \
    (y)->parent = __x;                                                                             \
  } while (0)

#define rbtree_insert_fixup(root_ptr, z_node)                                                      \
  do {                                                                                             \
    struct rb_node_link* __z = (z_node);                                                           \
    while (__z->parent && __z->parent->color == RBTREE_RED) {                                      \
      if (__z->parent == __z->parent->parent->left) {                                              \
        struct rb_node_link* __y = __z->parent->parent->right;                                     \
        if (__y && __y->color == RBTREE_RED) {                                                     \
          __z->parent->color = RBTREE_BLACK;                                                       \
          __y->color = RBTREE_BLACK;                                                               \
          __z->parent->parent->color = RBTREE_RED;                                                 \
          __z = __z->parent->parent;                                                               \
        } else {                                                                                   \
          if (__z == __z->parent->right) {                                                         \
            __z = __z->parent;                                                                     \
            rbtree_rotate_left ((root_ptr), __z);                                                  \
          }                                                                                        \
          __z->parent->color = RBTREE_BLACK;                                                       \
          __z->parent->parent->color = RBTREE_RED;                                                 \
          rbtree_rotate_right ((root_ptr), __z->parent->parent);                                   \
        }                                                                                          \
      } else {                                                                                     \
        struct rb_node_link* __y = __z->parent->parent->left;                                      \
        if (__y && __y->color == RBTREE_RED) {                                                     \
          __z->parent->color = RBTREE_BLACK;                                                       \
          __y->color = RBTREE_BLACK;                                                               \
          __z->parent->parent->color = RBTREE_RED;                                                 \
          __z = __z->parent->parent;                                                               \
        } else {                                                                                   \
          if (__z == __z->parent->left) {                                                          \
            __z = __z->parent;                                                                     \
            rbtree_rotate_right ((root_ptr), __z);                                                 \
          }                                                                                        \
          __z->parent->color = RBTREE_BLACK;                                                       \
          __z->parent->parent->color = RBTREE_RED;                                                 \
          rbtree_rotate_left ((root_ptr), __z->parent->parent);                                    \
        }                                                                                          \
      }                                                                                            \
    }                                                                                              \
    (*(root_ptr))->color = RBTREE_BLACK;                                                           \
  } while (0)

#define rbtree_insert(type, root_ptr, node, member, keyfield)                                      \
  do {                                                                                             \
    struct rb_node_link** __link = (root_ptr);                                                     \
    struct rb_node_link* __parent = NULL;                                                          \
    struct rb_node_link* __new = (node);                                                           \
    type* __nobj = rbtree_entry (__new, type, member);                                             \
    while (*__link) {                                                                              \
      __parent = *__link;                                                                          \
      type* __xobj = rbtree_entry (*__link, type, member);                                         \
      if (__nobj->keyfield < __xobj->keyfield)                                                     \
        __link = &((*__link)->left);                                                               \
      else                                                                                         \
        __link = &((*__link)->right);                                                              \
    }                                                                                              \
    __new->parent = __parent;                                                                      \
    __new->left = NULL;                                                                            \
    __new->right = NULL;                                                                           \
    __new->color = RBTREE_RED;                                                                     \
    *__link = __new;                                                                               \
    rbtree_insert_fixup (root_ptr, __new);                                                         \
  } while (0)

#define rbtree_find(type, root_ptr, keyval, out, member, keyfield)                                 \
  do {                                                                                             \
    (out) = NULL;                                                                                  \
    struct rb_node_link* __cur = *(root_ptr);                                                      \
    while (__cur) {                                                                                \
      type* __obj = rbtree_entry (__cur, type, member);                                            \
      if ((keyval) == __obj->keyfield) {                                                           \
        (out) = __cur;                                                                             \
        break;                                                                                     \
      } else if ((keyval) < __obj->keyfield)                                                       \
        __cur = __cur->left;                                                                       \
      else                                                                                         \
        __cur = __cur->right;                                                                      \
    }                                                                                              \
  } while (0)

#define rbtree_min(node, out)                                                                      \
  do {                                                                                             \
    (out) = NULL;                                                                                  \
    if ((node)) {                                                                                  \
      struct rb_node_link* __n = (node);                                                           \
      while (__n->left)                                                                            \
        __n = __n->left;                                                                           \
      (out) = __n;                                                                                 \
    }                                                                                              \
  } while (0)

#define rbtree_transplant(root_ptr, u, v)                                                          \
  do {                                                                                             \
    if (!(u)->parent)                                                                              \
      *(root_ptr) = (v);                                                                           \
    else if ((u) == (u)->parent->left)                                                             \
      (u)->parent->left = (v);                                                                     \
    else                                                                                           \
      (u)->parent->right = (v);                                                                    \
    if (v)                                                                                         \
      (v)->parent = (u)->parent;                                                                   \
  } while (0)

#define rbtree_delete_fixup(root_ptr, x_node, xparent_node)                                        \
  do {                                                                                             \
    struct rb_node_link* __x = (x_node);                                                           \
    struct rb_node_link* __xparent = (xparent_node);                                               \
    while (__x != *(root_ptr) && (__x == NULL || __x->color == RBTREE_BLACK)) {                    \
      if (__x == __xparent->left) {                                                                \
        struct rb_node_link* __w = __xparent->right;                                               \
        if (__w && __w->color == RBTREE_RED) {                                                     \
          __w->color = RBTREE_BLACK;                                                               \
          __xparent->color = RBTREE_RED;                                                           \
          rbtree_rotate_left (root_ptr, __xparent);                                                \
          __w = __xparent->right;                                                                  \
        }                                                                                          \
        if ((!__w->left || __w->left->color == RBTREE_BLACK) &&                                    \
            (!__w->right || __w->right->color == RBTREE_BLACK)) {                                  \
          __w->color = RBTREE_RED;                                                                 \
          __x = __xparent;                                                                         \
          __xparent = __x->parent;                                                                 \
        } else {                                                                                   \
          if (!__w->right || __w->right->color == RBTREE_BLACK) {                                  \
            if (__w->left)                                                                         \
              __w->left->color = RBTREE_BLACK;                                                     \
            __w->color = RBTREE_RED;                                                               \
            rbtree_rotate_right (root_ptr, __w);                                                   \
            __w = __xparent->right;                                                                \
          }                                                                                        \
          __w->color = __xparent->color;                                                           \
          __xparent->color = RBTREE_BLACK;                                                         \
          if (__w->right)                                                                          \
            __w->right->color = RBTREE_BLACK;                                                      \
          rbtree_rotate_left (root_ptr, __xparent);                                                \
          __x = *(root_ptr);                                                                       \
        }                                                                                          \
      } else {                                                                                     \
        struct rb_node_link* __w = __xparent->left;                                                \
        if (__w && __w->color == RBTREE_RED) {                                                     \
          __w->color = RBTREE_BLACK;                                                               \
          __xparent->color = RBTREE_RED;                                                           \
          rbtree_rotate_right (root_ptr, __xparent);                                               \
          __w = __xparent->left;                                                                   \
        }                                                                                          \
        if ((!__w->right || __w->right->color == RBTREE_BLACK) &&                                  \
            (!__w->left || __w->left->color == RBTREE_BLACK)) {                                    \
          __w->color = RBTREE_RED;                                                                 \
          __x = __xparent;                                                                         \
          __xparent = __x->parent;                                                                 \
        } else {                                                                                   \
          if (!__w->left || __w->left->color == RBTREE_BLACK) {                                    \
            if (__w->right)                                                                        \
              __w->right->color = RBTREE_BLACK;                                                    \
            __w->color = RBTREE_RED;                                                               \
            rbtree_rotate_left (root_ptr, __w);                                                    \
            __w = __xparent->left;                                                                 \
          }                                                                                        \
          __w->color = __xparent->color;                                                           \
          __xparent->color = RBTREE_BLACK;                                                         \
          if (__w->left)                                                                           \
            __w->left->color = RBTREE_BLACK;                                                       \
          rbtree_rotate_right (root_ptr, __xparent);                                               \
          __x = *(root_ptr);                                                                       \
        }                                                                                          \
      }                                                                                            \
    }                                                                                              \
    if (__x)                                                                                       \
      __x->color = RBTREE_BLACK;                                                                   \
  } while (0)

#define rbtree_next(node, out)                                                                     \
  do {                                                                                             \
    (out) = NULL;                                                                                  \
    if (node) {                                                                                    \
      if ((node)->right) {                                                                         \
        struct rb_node_link* __n = (node)->right;                                                  \
        while (__n->left)                                                                          \
          __n = __n->left;                                                                         \
        (out) = __n;                                                                               \
      } else {                                                                                     \
        struct rb_node_link* __n = (node);                                                         \
        struct rb_node_link* __p = (node)->parent;                                                 \
        while (__p && __n == __p->right) {                                                         \
          __n = __p;                                                                               \
          __p = __p->parent;                                                                       \
        }                                                                                          \
        (out) = __p;                                                                               \
      }                                                                                            \
    }                                                                                              \
  } while (0)

#define rbtree_prev(node, out)                                                                     \
  do {                                                                                             \
    (out) = NULL;                                                                                  \
    if (node) {                                                                                    \
      if ((node)->left) {                                                                          \
        struct rb_node_link* __n = (node)->left;                                                   \
        while (__n->right)                                                                         \
          __n = __n->right;                                                                        \
        (out) = __n;                                                                               \
      } else {                                                                                     \
        struct rb_node_link* __n = (node);                                                         \
        struct rb_node_link* __p = (node)->parent;                                                 \
        while (__p && __n == __p->left) {                                                          \
          __n = __p;                                                                               \
          __p = __p->parent;                                                                       \
        }                                                                                          \
        (out) = __p;                                                                               \
      }                                                                                            \
    }                                                                                              \
  } while (0)

#define rbtree_first(root_ptr, out) rbtree_min (*(root_ptr), out)

#define rbtree_last(root_ptr, out)                                                                 \
  do {                                                                                             \
    (out) = NULL;                                                                                  \
    struct rb_node_link* __n = *(root_ptr);                                                        \
    if (__n) {                                                                                     \
      while (__n->right)                                                                           \
        __n = __n->right;                                                                          \
      (out) = __n;                                                                                 \
    }                                                                                              \
  } while (0)

#endif // _KERNEL_LIBK_RBTREE_H