c语言map的实现-使用内核红黑树

一、前言

内核数据结构中，上次讲了一次内核链表的使用《c语言双向循环链表实现-使用内核链表》，这次再接着讲一下内核中常用的<key,value>结构，红黑树（rbtree）。

二、相关知识

2.1 红黑树特点

红黑树为自平衡二叉查找树，具备了以下几种特征[1]：

1. 节点是红色或黑色。
2. 根是黑色。
3. 所有叶子都是黑色（叶子是NIL节点）。
4. 每个红色节点必须有两个黑色的子节点。（从每个叶子到根的所有路径上不能有两个连续的红色节点。）
5. 从任一节点到其每个叶子的所有简单路径都包含相同数目的黑色节点。

红黑树的应用比较广泛，主要是用它来存储有序的数据，它的时间复杂度是O(lgn)，效率非常之高。
例如，Java集合中的TreeSet和TreeMap，C++ STL中的set、map，以及Linux虚拟内存的管理，都是通过红黑树去实现的。
结构图例如下，详细原理请移步文章[2]：

2.2 内核红黑树的使用

源码位置：linux-2.6.22.1/lib/rbtree.c、linux-2.6.22.1/include/linux/rbtree.h

文档位置：linux-2.6.22.1/Documentation/rbtree.txt 

根节点、子节点定义如下：

struct rb_node {
        unsigned long  __rb_parent_color;
        struct rb_node *rb_right;
        struct rb_node *rb_left;
} __attribute__((aligned(sizeof(long))));
    /* The alignment might seem pointless, but allegedly CRIS needs it */

struct rb_root {
        struct rb_node *rb_node;
};

实用接口：

接口	描述
void rb_link_node(struct rb_node * node,     struct rb_node * parent, struct rb_node ** rb_link);	插入结点
void rb_insert_color(struct rb_node *, struct rb_root *);	调整红黑树
void rb_erase(struct rb_node *, struct rb_root *);	删除节点
struct rb_node *rb_first(struct rb_root *);	获取头节点
struct rb_node *rb_next(struct rb_node *);	获取下一个节点

2.3 调整

rbtree.c\rbtree.h 在内核源码中，需要进行调整下头文件定义，提取到工程中：

cp -a /usr/src/linux-2.6.22.1/lib/rbtree.c ./
cp -a /usr/src/linux-2.6.22.1/include/linux/rbtree.h ./

sed -i 's@#include <linux\/rbtree\.h>@#include "rbtree.h"@g' rbtree.c
sed -i 's@#include <linux\/module\.h>@#define EXPORT_SYMBOL(a)@g' rbtree.c

sed -i 's@#include <linux\/kernel\.h>@@g' rbtree.h
sed -i 's@#include <linux\/stddef\.h>@#include <stddef.h>@g' rbtree.h
sed -i '/container_of/ a\#define container_of(ptr, type, member) ({ const typeof( ((type *)0)->member ) *__mptr = (ptr);  (type *)( (char *)__mptr - offsetof(type,member) );})' rbtree.h

rbtree.txt 文档中提到：需要在自定义数据结构中定义，再实现应用的插入、查找、删除接口

三、应用

使用红黑树结构保存目录树信息，key为文件名、value为文件信息，支持文件项新增、按文件名查找、文件项删除三种功能。结构体为 struct memfs_file，其中里面包含了 struct rb_node node作为红黑树的关联结构、path为文件名称：

struct memfs_file {
    char *path;                     /* File path */
    void *data;                     /* File content */
    u8 free_on_delete;

    struct stat vstat;              /* File stat */

    pthread_mutex_t lock;
    pthread_cond_t write_finished;

    struct rb_node node;
};

插入节点，注意插入必须是新的文件名，已有的文件名是无法插入的：

static int __insert(struct rb_root *root, struct memfs_file *pf)
{
    struct rb_node **new = &(root->rb_node), *parent = NULL;

    /* Figure out where to put new node */
    while (*new) {
        struct memfs_file *this = container_of(*new, struct memfs_file, node);
        int result = strcmp(pf->path, this->path);
        parent = *new;
        if (result < 0) {
            new = &((*new)->rb_left);
        }
        else if (result > 0) {
            new = &((*new)->rb_right);
        }
        else {
            printf("%s is exist!!!\n", pf->path);
            return -1;
        }
    }

    /* Add new node and rebalance tree. */
    rb_link_node(&pf->node, parent, new);
    rb_insert_color(&pf->node, root);
    return 0;
}

根据文件名查找节点，可以看出复杂度为O(logN)，注意通过rb_node再获取父结构体指针需要用到 container_of 宏： 

static struct memfs_file *__search(struct rb_root *root, const char *path)
{
    struct memfs_file *pf = NULL;
    struct rb_node *node = root->rb_node;

    while (node) {
        pf = container_of(node, struct memfs_file, node);

        int result = strcmp(path, pf->path);
        if (result < 0) {
            node = node->rb_left;
        }
        else if (result > 0) {
            node = node->rb_right;
        }
        else {
            return pf;
        }
    }

//  printf("%s not found!!!\n", path);
    return NULL;
}

删除节点，为了配合应用自定义了__free释放函数，__delete返回-1表示未找到，>=0表示成功：

static inline void __free(struct memfs_file *pf)
{
    if (pf->free_on_delete) {
        if (pf->data) {
            free(pf->data);
        }
        if (pf->path) {
            free(pf->path);
        }
        free(pf);
    }
}

static int __delete(struct rb_root *root, const char *path)
{
    struct memfs_file *pf = __search(root, path);
    if (!pf) {
        return -1;
    }   
    
    int blocks = pf->vstat.st_blocks;
    rb_erase(&pf->node, root);
    __free(pf);

    return blocks;
} 

简单的测试函数如下： 

int main(int argc, char *argv[])
{
    struct rb_root root = RB_ROOT;

    struct memfs_file files[] = {
        [0] = { "/tmp/a.jpg",       "0x1234", },
        [1] = { "/tmp/c.txt",       "0x1234", },
        [2] = { "/tmp/d.exe",       "0x1234", },
        [3] = { "/tmp/b.dir/b.txt", "0x1234", },
    };

    __insert(&root, &files[0]);
    __insert(&root, &files[1]);
    __insert(&root, &files[2]);
    __insert(&root, &files[3]);


#define TEST_SEARCH(root, name) do { \
    const struct memfs_file *pf = __search(&root, name); \
    assert(strcmp(pf->path, name) == 0); \
    printf("search: %s\t\t\t\t[OK]\n", name); \
} while(0)

    TEST_SEARCH(root, files[0].path);
    TEST_SEARCH(root, files[1].path);
    TEST_SEARCH(root, files[2].path);
    TEST_SEARCH(root, files[3].path);

    /* Trasver all */
    if (1) {
        struct rb_node *node = NULL;
        for (node = rb_first(&root); node; node = rb_next(node)) {
            const struct memfs_file *pf = rb_entry(node, struct memfs_file, node);
            printf("iterator: %s\t\t\t\t[OK]\n", pf->path);
        }
    }

#define TEST_DELETE(root, name) do { \
    __delete(&root, name); \
    assert(__search(&root, name) == NULL); \
    printf("delete: %s\t\t\t\t[OK]\n", name); \
} while(0)

    TEST_DELETE(root, files[0].path);
    exit(EXIT_SUCCESS);
}

执行结果如下： 

search: /tmp/a.jpg                              [OK]
search: /tmp/c.txt                              [OK]
search: /tmp/d.exe                              [OK]
search: /tmp/b.dir/b.txt                        [OK]
iterator: /tmp/a.jpg                            [OK]
iterator: /tmp/b.dir/b.txt                      [OK]
iterator: /tmp/c.txt                            [OK]
iterator: /tmp/d.exe                            [OK]
delete: /tmp/a.jpg                              [OK]

四、结论

内核的红黑树实现提供的是教科书式的使用，为c语言的首选，学会了内核链表的用法，再使用红黑树上手难度不大；跟c++相比可相当于map的地位，其rb_first、rb_next用法跟迭代器思路非常贴合，跟hashmap的区别是红黑树插入、查找、删除效率均衡，并且可以实现根据key顺序遍历、范围查找的功能；

2.6内核的红黑树算是最好移植到的，3.0之后红黑树开始考虑了c++的兼容，文件变多反而不太好移植了；本文作为铺垫，后续将再提供一个基于红黑树结构实现一个文件系统的例子；

 

参考文章：

[1] wiki，https://zh.wikipedia.org/wiki/%E7%BA%A2%E9%BB%91%E6%A0%91

[2] 原理详解，https://www.cnblogs.com/skywang12345/p/3245399.html