tree.h - Add RB_AUGMENT support

Submitted-by: Brills Peng

diff --git a/sys/sys/tree.h b/sys/sys/tree.h
index f4d7807..5d35db3 100644 (file)
--- a/sys/sys/tree.h
+++ b/sys/sys/tree.h
@@ -340,8 +340,8 @@ struct {                                                            \
        RB_COLOR(red, field) = RB_RED;                                  \
 } while (/*CONSTCOND*/ 0)
 
-#ifdef RB_AUGMENT
-#error "RB_AUGMENT not supported by DragonFly"
+#ifndef RB_AUGMENT
+#define RB_AUGMENT(x)  do {} while (0)
 #endif
 
 #define RB_ROTATE_LEFT(head, elm, tmp, field) do {                     \
@@ -349,6 +349,7 @@ struct {                                                            \
        if ((RB_RIGHT(elm, field) = RB_LEFT(tmp, field)) != NULL) {     \
                RB_PARENT(RB_LEFT(tmp, field), field) = (elm);          \
        }                                                               \
+       RB_AUGMENT(elm);                                                \
        if ((RB_PARENT(tmp, field) = RB_PARENT(elm, field)) != NULL) {  \
                if ((elm) == RB_LEFT(RB_PARENT(elm, field), field))     \
                        RB_LEFT(RB_PARENT(elm, field), field) = (tmp);  \
@@ -358,6 +359,9 @@ struct {                                                            \
                (head)->rbh_root = (tmp);                               \
        RB_LEFT(tmp, field) = (elm);                                    \
        RB_PARENT(elm, field) = (tmp);                                  \
+       RB_AUGMENT(tmp);                                                \
+       if ((RB_PARENT(tmp, field)))                                    \
+               RB_AUGMENT(RB_PARENT(tmp, field));                      \
 } while (/*CONSTCOND*/ 0)
 
 #define RB_ROTATE_RIGHT(head, elm, tmp, field) do {                    \
@@ -365,6 +369,7 @@ struct {                                                            \
        if ((RB_LEFT(elm, field) = RB_RIGHT(tmp, field)) != NULL) {     \
                RB_PARENT(RB_RIGHT(tmp, field), field) = (elm);         \
        }                                                               \
+       RB_AUGMENT(elm);                                                \
        if ((RB_PARENT(tmp, field) = RB_PARENT(elm, field)) != NULL) {  \
                if ((elm) == RB_LEFT(RB_PARENT(elm, field), field))     \
                        RB_LEFT(RB_PARENT(elm, field), field) = (tmp);  \
@@ -374,6 +379,9 @@ struct {                                                            \
                (head)->rbh_root = (tmp);                               \
        RB_RIGHT(tmp, field) = (elm);                                   \
        RB_PARENT(elm, field) = (tmp);                                  \
+       RB_AUGMENT(tmp);                                                \
+       if ((RB_PARENT(tmp, field)))                                    \
+               RB_AUGMENT(RB_PARENT(tmp, field));                      \
 } while (/*CONSTCOND*/ 0)
 
 /* Generates prototypes and inline functions */
@@ -583,6 +591,7 @@ name##_RB_REMOVE(struct name *head, struct type *elm)                       \
                                RB_LEFT(parent, field) = child;         \
                        else                                            \
                                RB_RIGHT(parent, field) = child;        \
+                       RB_AUGMENT(parent);                             \
                } else                                                  \
                        RB_ROOT(head) = child;                          \
                if (RB_PARENT(elm, field) == old)                       \
@@ -593,11 +602,18 @@ name##_RB_REMOVE(struct name *head, struct type *elm)                     \
                                RB_LEFT(RB_PARENT(old, field), field) = elm;\
                        else                                            \
                                RB_RIGHT(RB_PARENT(old, field), field) = elm;\
+                       RB_AUGMENT(RB_PARENT(old, field));              \
                } else                                                  \
                        RB_ROOT(head) = elm;                            \
                RB_PARENT(RB_LEFT(old, field), field) = elm;            \
                if (RB_RIGHT(old, field))                               \
                        RB_PARENT(RB_RIGHT(old, field), field) = elm;   \
+               if (parent) {                                           \
+                       left = parent;                                  \
+                       do {                                            \
+                               RB_AUGMENT(left);                       \
+                       } while ((left = RB_PARENT(left, field)) != NULL); \
+               }                                                       \
                goto color;                                             \
        }                                                               \
        parent = RB_PARENT(elm, field);                                 \
@@ -609,6 +625,7 @@ name##_RB_REMOVE(struct name *head, struct type *elm)                       \
                        RB_LEFT(parent, field) = child;                 \
                else                                                    \
                        RB_RIGHT(parent, field) = child;                \
+               RB_AUGMENT(parent);                                     \
        } else                                                          \
                RB_ROOT(head) = child;                                  \
 color:                                                                 \
@@ -641,6 +658,7 @@ name##_RB_INSERT(struct name *head, struct type *elm)                       \
                        RB_LEFT(parent, field) = elm;                   \
                else                                                    \
                        RB_RIGHT(parent, field) = elm;                  \
+               RB_AUGMENT(parent);                                     \
        } else                                                          \
                RB_ROOT(head) = elm;                                    \
        name##_RB_INSERT_COLOR(head, elm);                              \