#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_entry(node, type, member) ((type*)((char*)(node) - offsetof (type, member)))

#define rbtree_node_color(x) ((x) ? (x)->color : RBTREE_BLACK)

#define rbtree_rotate_left(root_ptr, x_node)                                                       \
  do {                                                                                             \
    struct rb_node_link* __x = (x_node);                                                           \
    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_node)                                                      \
  do {                                                                                             \
    struct rb_node_link* __y = (y_node);                                                           \
    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 (rbtree_node_color (__y) == 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 (rbtree_node_color (__y) == 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 = __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;                                                                                  \
    struct rb_node_link* __n = (node);                                                             \
    while (__n && __n->left)                                                                       \
      __n = __n->left;                                                                             \
    (out) = __n;                                                                                   \
  } while (0)

#define rbtree_max(node, out)                                                                      \
  do {                                                                                             \
    (out) = NULL;                                                                                  \
    struct rb_node_link* __n = (node);                                                             \
    while (__n && __n->right)                                                                      \
      __n = __n->right;                                                                            \
    (out) = __n;                                                                                   \
  } while (0)

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

#define rbtree_transplant(root_ptr, u_node, v_node)                                                \
  do {                                                                                             \
    struct rb_node_link* __u = (u_node);                                                           \
    struct rb_node_link* __v = (v_node);                                                           \
    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* __rdf_x = (x_node);                                                       \
    struct rb_node_link* __rdf_xp = (xparent_node);                                                \
    while (__rdf_xp && (__rdf_x == NULL || __rdf_x->color == RBTREE_BLACK)) {                      \
      if (__rdf_x == __rdf_xp->left) {                                                             \
        struct rb_node_link* __w = __rdf_xp->right;                                                \
        if (rbtree_node_color (__w) == RBTREE_RED) {                                               \
          __w->color = RBTREE_BLACK;                                                               \
          __rdf_xp->color = RBTREE_RED;                                                            \
          rbtree_rotate_left (root_ptr, __rdf_xp);                                                 \
          __w = __rdf_xp->right;                                                                   \
        }                                                                                          \
        if (rbtree_node_color (__w->left) == RBTREE_BLACK &&                                       \
            rbtree_node_color (__w->right) == RBTREE_BLACK) {                                      \
          if (__w)                                                                                 \
            __w->color = RBTREE_RED;                                                               \
          __rdf_x = __rdf_xp;                                                                      \
          __rdf_xp = __rdf_x->parent;                                                              \
        } else {                                                                                   \
          if (rbtree_node_color (__w->right) == RBTREE_BLACK) {                                    \
            if (__w->left)                                                                         \
              __w->left->color = RBTREE_BLACK;                                                     \
            __w->color = RBTREE_RED;                                                               \
            rbtree_rotate_right (root_ptr, __w);                                                   \
            __w = __rdf_xp->right;                                                                 \
          }                                                                                        \
          __w->color = __rdf_xp->color;                                                            \
          __rdf_xp->color = RBTREE_BLACK;                                                          \
          if (__w->right)                                                                          \
            __w->right->color = RBTREE_BLACK;                                                      \
          rbtree_rotate_left (root_ptr, __rdf_xp);                                                 \
          __rdf_x = *(root_ptr);                                                                   \
          break;                                                                                   \
        }                                                                                          \
      } else {                                                                                     \
        struct rb_node_link* __w = __rdf_xp->left;                                                 \
        if (rbtree_node_color (__w) == RBTREE_RED) {                                               \
          __w->color = RBTREE_BLACK;                                                               \
          __rdf_xp->color = RBTREE_RED;                                                            \
          rbtree_rotate_right (root_ptr, __rdf_xp);                                                \
          __w = __rdf_xp->left;                                                                    \
        }                                                                                          \
        if (rbtree_node_color (__w->right) == RBTREE_BLACK &&                                      \
            rbtree_node_color (__w->left) == RBTREE_BLACK) {                                       \
          if (__w)                                                                                 \
            __w->color = RBTREE_RED;                                                               \
          __rdf_x = __rdf_xp;                                                                      \
          __rdf_xp = __rdf_x->parent;                                                              \
        } else {                                                                                   \
          if (rbtree_node_color (__w->left) == RBTREE_BLACK) {                                     \
            if (__w->right)                                                                        \
              __w->right->color = RBTREE_BLACK;                                                    \
            __w->color = RBTREE_RED;                                                               \
            rbtree_rotate_left (root_ptr, __w);                                                    \
            __w = __rdf_xp->left;                                                                  \
          }                                                                                        \
          __w->color = __rdf_xp->color;                                                            \
          __rdf_xp->color = RBTREE_BLACK;                                                          \
          if (__w->left)                                                                           \
            __w->left->color = RBTREE_BLACK;                                                       \
          rbtree_rotate_right (root_ptr, __rdf_xp);                                                \
          __rdf_x = *(root_ptr);                                                                   \
          break;                                                                                   \
        }                                                                                          \
      }                                                                                            \
    }                                                                                              \
    if (__rdf_x)                                                                                   \
      __rdf_x->color = RBTREE_BLACK;                                                               \
  } while (0)

#define rbtree_delete(root_ptr, z_node)                                                            \
  do {                                                                                             \
    struct rb_node_link* __rd_z = (z_node);                                                        \
    struct rb_node_link* __rd_y = __rd_z;                                                          \
    struct rb_node_link* __rd_x = NULL;                                                            \
    struct rb_node_link* __rd_xp = NULL;                                                           \
    int __rd_y_orig_color = __rd_y->color;                                                         \
    if (!__rd_z->left) {                                                                           \
      __rd_x = __rd_z->right;                                                                      \
      __rd_xp = __rd_z->parent;                                                                    \
      rbtree_transplant (root_ptr, __rd_z, __rd_z->right);                                         \
    } else if (!__rd_z->right) {                                                                   \
      __rd_x = __rd_z->left;                                                                       \
      __rd_xp = __rd_z->parent;                                                                    \
      rbtree_transplant (root_ptr, __rd_z, __rd_z->left);                                          \
    } else {                                                                                       \
      rbtree_min (__rd_z->right, __rd_y);                                                          \
      __rd_y_orig_color = __rd_y->color;                                                           \
      __rd_x = __rd_y->right;                                                                      \
      if (__rd_y->parent == __rd_z) {                                                              \
        __rd_xp = __rd_y;                                                                          \
        if (__rd_x)                                                                                \
          __rd_x->parent = __rd_y;                                                                 \
      } else {                                                                                     \
        __rd_xp = __rd_y->parent;                                                                  \
        rbtree_transplant (root_ptr, __rd_y, __rd_y->right);                                       \
        __rd_y->right = __rd_z->right;                                                             \
        __rd_y->right->parent = __rd_y;                                                            \
      }                                                                                            \
      rbtree_transplant (root_ptr, __rd_z, __rd_y);                                                \
      __rd_y->left = __rd_z->left;                                                                 \
      __rd_y->left->parent = __rd_y;                                                               \
      __rd_y->color = __rd_z->color;                                                               \
    }                                                                                              \
    if (__rd_y_orig_color == RBTREE_BLACK)                                                         \
      rbtree_delete_fixup (root_ptr, __rd_x, __rd_xp);                                             \
  } 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)

#endif // _KERNEL_LIBK_RBTREE_H