Linux内核中使用红黑树的扩展特性实现区间树（Interval tree）

原文地址：http://galex.cn/linux%E5%86%85%E6%A0%B8%E4%B8%AD%E4%BD%BF%E7%94%A8%E7%BA%A2%E9%BB%91%E6%A0%91%E7%9A%84%E6%89%A9%E5%B1%95%E7%89%B9%E6%80%A7%E5%AE%9E%E7%8E%B0%E5%8C%BA%E9%97%B4%E6%A0%91%EF%BC%88interval-tree%EF%BC%89/

阅读3.0版本以后的内核源代码就会发现，新版本内核中的红黑树的实现与之前有所不同。新的实现提供了增强特性，用户可以在红黑树结点中符加一些数据结构来提供更多的功能。区间树就是一种增强型的红黑树。《算法导论》中有关于区间树详细的介绍。下图是取自《算法导论》的一棵区间树：

区间树

在Linux内核中红黑树的实现中提到了这种特性，但没有详细介绍。下面就利用这种特性实现区间树。

按照《算法导论》中的讲解，区间树需要提供三种基本操作：查找、插入和删除。这里定义区间树的结构是：
1
2
3
4
5
6
    
struct interval_tree_node {
    struct rb_node rb;    /* 红黑树结点 */
    unsigned low;         /* 区间左边界 */
    unsigned high;        /* 区间右边界 */
    unsigned max;         /* 子树中的最大边界值 */
};

先实现查找操作，因为它不需要修改结点信息：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
    
struct interval_tree_node* interval_tree_search(struct rb_root *root, unsigned low, unsigned high)
{
    struct interval_tree_node *node;
 
    if (!root->rb_node)
        return NULL;
    node = rb_entry(root->rb_node, struct interval_tree_node, rb);
 
    while (true) {
        /* 当前区间与目标区间有重合，查找成功，返回 */
        if (node->low <= high && node->high >= low)
                return node;
 
        if (node->rb.rb_left) {
            struct interval_tree_node *left =
                        rb_entry(node->rb.rb_left, struct interval_tree_node, rb);
            /* 左子树的最大值大于要查找的左边界，从左子树查找 */
            if (left->max >= low) {
                node = left;
                continue;
            }
        }
        if (node->rb.rb_right) {
            node = rb_entry(node->rb.rb_right, struct interval_tree_node, rb);
            /* 右子树的最大值大于要查找的左边界，从右子树查找 */
            if (node->max >= low)
                continue;
        }
 
        return NULL;
    }
}

插入删除操作需要修改结点的信息，可以使用三个回调函数实现，它们通常被封装在struct rb_augment_callbacks 结构中：
0
1
2
3
4
5
    
<rbtree_augmented.h>
struct rb_augment_callbacks {
    void (*propagate)(struct rb_node *node, struct rb_node *stop);
    void (*copy)(struct rb_node *old, struct rb_node *new);
    void (*rotate)(struct rb_node *old, struct rb_node *new);
};

propagate 是传播回调函数。它的工作就是更新node到它的祖先stop结点之间的每一个结点的附加信息。

copy 是拷贝回调函数。它的工作是将new结点的附加信息设置成old的附加信息。

rotate 是旋转回调函数。它的工作是将new结点的附加信息设置成old的附加信息，并重新计算old的附加信息。

在Linux内核中红黑树的实现中分析普通红黑树的实现时已经看到，红黑树的实现操作会在合适时间会调用这三个回调函数更新结点信息，现在要做的只是实现这三个回调函数。

对区间树执行插入删除操作时需要修改结点的信息只有子树的最大边界值，所以先创建一个辅助函数来获得子树的最大边界值。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
    
static inline unsigned interval_tree_get_subtree_max(struct interval_tree_node *node)
{
    unsigned max = node->high, subtree_max;
    if (node->rb.rb_left) {
        subtree_max = rb_entry(node->rb.rb_left, struct interval_tree_node, rb)->max;
        if (max < subtree_max)
            max = subtree_max;
    }
    if (node->rb.rb_right) {
        subtree_max = rb_entry(node->rb.rb_right, struct interval_tree_node, rb)->max;
        if (max < subtree_max)
            max = subtree_max;
    }
    return max;
}

propagate 函数的两个参数分别是起始结点和终止结点，在它们之间的所有结点都会被重新计算子树最大边界值：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
    
static void interval_tree_augment_propagate(struct rb_node *rb, struct rb_node *stop)
{
    while (rb != stop) {
        struct interval_tree_node *node = rb_entry(rb, struct interval_tree_node, rb);
        unsigned subtree_max = interval_tree_get_subtree_max(node);
 
        if (node->max == subtree_max)
            break;
 
        node->max = subtree_max;
 
        rb = rb_parent(&node->rb);
    }
}

copy函数只是简单地复制：
1
2
3
4
5
6
7
    
static void interval_tree_augment_copy(struct rb_node *rb_old, struct rb_node *rb_new)
{
    struct interval_tree_node *old = rb_entry(rb_old, struct interval_tree_node, rb);
    struct interval_tree_node *new = rb_entry(rb_new, struct interval_tree_node, rb);
 
    new->max = old->max;
}

类似的写出 rotate 函数：
1
2
3
4
5
6
7
8
    
static void interval_tree_augment_rotate(struct rb_node *rb_old, struct rb_node *rb_new)
{
    struct interval_tree_node *old = rb_entry(rb_old, struct interval_tree_node, rb);
    struct interval_tree_node *new = rb_entry(rb_new, struct interval_tree_node, rb);
 
    new->max = old->max;
    old->max = interval_tree_get_subtree_max(old);
}

将这三个回调函数封装到 rb_augment_callbacks 结构中：
1
2
3
4
5
    
static const struct rb_augment_callbacks interval_tree_augment_callbacks = {
    interval_tree_augment_propagate,
    interval_tree_augment_copy,
    interval_tree_augment_rotate
};

接下来就可以实现插入删除操作了：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
    
void interval_tree_insert(struct rb_root *root, struct interval_tree_node *node)
{
    unsigned low = node->low, high = node->high;
    struct rb_node **link = &root->rb_node, *rb_parent;
    struct interval_tree_node *parent;
 
    while (*link) {
        rb_parent = *link;
        parent = rb_entry(rb_parent, struct interval_tree_node, rb);
 
        if (parent->max < high)
            parent->max = high;
 
        if (low < parent->low)
            link = &parent->rb.rb_left;
        else
            link = &parent->rb.rb_right;
    }
 
    node->max = high;
    rb_link_node(&node->rb, rb_parent, link);
    rb_insert_augmented(&node->rb, root, &interval_tree_augment_callbacks);
}
 
void interval_tree_delete(struct rb_root *root, struct interval_tree_node *node)
{
    rb_erase_augmented(&node->rb, root, &interval_tree_augment_callbacks);
}

删除操作很简单，只需要提供回调函数，红黑树的删除操作就可以做所有工作。

插入结点时，需要找到待插入结点的位置，Linux内核中红黑树的实现中说过，新插入的结点总是叶子结点。所以在while循环中不断向下遍历，同时更新经过的结点的子树最大边界值。

rb_link_node 函数的作用是将结点连接到红黑树中，它有三个参数，第一个是待插入的结点，第二个是待插入结点的父结点，第三个参数是一个二维指针，它指向的是父结点的左孩子或者右孩子，也就是在while循环中找到的位置。该函数定义在<rbtree.h>中：
0
1
2
3
4
5
6
7
8
    
<rbtree.h>
static inline void rb_link_node(struct rb_node * node, struct rb_node * parent,
                struct rb_node ** rb_link)
{
    node->__rb_parent_color = (unsigned long)parent;
    node->rb_left = node->rb_right = NULL;
 
    *rb_link = node;
}

最后调用 rb_insert_augmented 函数调整红黑树使之符合红黑树的特性。