Linux kernel rb-tree (2)

寫了個簡單的Demo，使用內(nèi)核提供的接口創(chuàng)建了一個紅黑樹。

API

// include/linux/rbtree_augmented.h
#define RB_RED          0
#define RB_BLACK        1

// linux-5.4/include/linux/rbtree.h

struct rb_root {
        struct rb_node *rb_node;
};

struct rb_node {
        unsigned long  __rb_parent_color;
        struct rb_node *rb_right;
        struct rb_node *rb_left;
} __attribute__((aligned(sizeof(long))));

#define RB_ROOT (struct rb_root) { NULL, }

#define rb_parent(r)   ((struct rb_node *)((r)->__rb_parent_color & ~3))
#define rb_color(rb)       __rb_color((rb)->__rb_parent_color)
#define __rb_color(pc)     ((pc) & 1)

#define rb_entry(ptr, type, member) container_of(ptr, type, member)

static inline void rb_link_node(struct rb_node * node, struct rb_node * parent, struct rb_node ** rb_link);

void rb_insert_color(struct rb_node *node, struct rb_root *root);

Demo

code

/*
 * Linux kernel rb-tree test
 *
 * Tested on: Linux 5.4.0-58-generic #64-Ubuntu SMP Wed Dec 9 08:16:25 UTC 2020 x86_64 x86_64 x86_64 GNU/Linux
 *
 * */
#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/rbtree.h>
#include <linux/rbtree_augmented.h>

#define BUFF_SIZE 24

struct planet {
        char name[BUFF_SIZE];
        int mass;
        struct rb_node rb_node;
};

struct rb_root planet_rb_root = RB_ROOT;

void insert_planet( struct planet * planet);
void print_planet_tree(struct rb_node *n);
struct planet *create_planet(char * name, int mass);

// 創(chuàng)建一個節(jié)點(planet)
struct planet *create_planet(char * name, int mass){
        struct planet *new_plt;

        if(strlen(name) >= BUFF_SIZE){
                pr_info("[-] name too long\n");
                return NULL;
        }

        new_plt = vmalloc(sizeof(struct planet));

        strcpy(new_plt->name, name);

        new_plt->mass = mass;

        pr_info("%s created: %px\n", name, new_plt);

        return new_plt;
}

// 插入到紅黑樹(planet_rb_root)中
void insert_planet( struct planet * planet){
        struct rb_node **p = &planet_rb_root.rb_node;
        struct rb_node *parent = NULL;

        while(*p){
                struct planet *tmp_planet;
                parent = *p;
                tmp_planet = rb_entry(parent, struct planet, rb_node);
                if(planet->mass < tmp_planet->mass){
                        p = &(*p)->rb_left;
                }else{
                        p = &(*p)->rb_right;
                }
        }

        rb_link_node(&planet->rb_node, parent, p);
        rb_insert_color(&planet->rb_node, &planet_rb_root);
}


// 遍歷紅黑樹并打印節(jié)點信息
void print_planet_tree(struct rb_node *n){
        struct planet *planet;
        char *p_buff =  vmalloc(4096);
        char *pos;

        if(!n){
                return;
        }

        planet = rb_entry(n, struct planet, rb_node);

        pos = p_buff;
        pos += sprintf(pos, "<%px> planet start ****************************************\n", planet);
        pos += sprintf(pos, "<%px> planet->name: %s(%px)\n", &planet->name, planet->name, planet->name);
        pos += sprintf(pos, "<%px> planet->mass: %d\n", &planet->mass, planet->mass);
        pos += sprintf(pos, "<%px> planet->rb_node: \n", &planet->rb_node);

        pos += sprintf(pos, "<%px> \tplanet->rb_node.__rb_parent_color %lx, rb_parent: %px, rb_color: %s\n",
                                &planet->rb_node.__rb_parent_color,
                                planet->rb_node.__rb_parent_color,
                                rb_parent(&planet->rb_node),
                                rb_color(&planet->rb_node)? "black": "red");

        pos += sprintf(pos, "<%px> \tplanet->rb_node.rb_right %px\n", &planet->rb_node.rb_right, planet->rb_node.rb_right);
        pos += sprintf(pos, "<%px> \tplanet->rb_node.rb_left %px\n", &planet->rb_node.rb_left, planet->rb_node.rb_left);
        pos += sprintf(pos, "<%px> planet end  *****************************************\n\n", planet+1);

        pr_info("%s", p_buff);

        vfree(p_buff);

        if(n->rb_left){
                print_planet_tree(n->rb_left);
        }
        if(n->rb_right){
                print_planet_tree(n->rb_right);
        }

        return;
}


void rbt(void){
        struct planet *earth, *mars, *saturn, *venus, *jupiter;

        earth = create_planet("Earth", 0);
        venus = create_planet("Venus", 1);
        saturn = create_planet("Saturn", 3);
        mars = create_planet("Mars", 4);
        jupiter = create_planet("Jupiter", 7);

        insert_planet(earth);
        insert_planet(venus);
        insert_planet(saturn);
        insert_planet(mars);
        insert_planet(jupiter);

        pr_info("\n<%px> planet_rb_root.rb_node: %px\n\n", &planet_rb_root.rb_node, planet_rb_root.rb_node);

        print_planet_tree(planet_rb_root.rb_node);
}


static int __init rbt_init(void)
{
    printk(KERN_INFO "Hello rbt\n");
    rbt();
    return 0;
}

static void __exit rbt_exit(void)
{
    printk(KERN_INFO "Goodbye rbt\n");
}


module_init(rbt_init);
module_exit(rbt_exit);

MODULE_LICENSE("GPL");
MODULE_AUTHOR("X++D");
MODULE_DESCRIPTION("Kernel xxx Module.");
MODULE_VERSION("0.1");

output

$ insmod rbt-t.ko
$ rmmod rbt-t.ko
$ dmesg

[1954676.388309] Earth created: ffff9d49c00b5000
[1954676.388311] Venus created: ffff9d49c0213000
[1954676.388314] Saturn created: ffff9d49c0247000
[1954676.388316] Mars created: ffff9d49c0277000
[1954676.388318] Jupiter created: ffff9d49c0279000
[1954676.388319]
                 <ffffffffc0ed5380> planet_rb_root.rb_node: ffff9d49c0213020

[1954676.388325] <ffff9d49c0213000> planet start ****************************************
                 <ffff9d49c0213000> planet->name: Venus(ffff9d49c0213000)
                 <ffff9d49c0213018> planet->mass: 1
                 <ffff9d49c0213020> planet->rb_node:
                 <ffff9d49c0213020>     planet->rb_node.__rb_parent_color 1, rb_parent: 0000000000000000, rb_color: black
                 <ffff9d49c0213028>     planet->rb_node.rb_right ffff9d49c0277020
                 <ffff9d49c0213030>     planet->rb_node.rb_left ffff9d49c00b5020
                 <ffff9d49c0213038> planet end  *****************************************

[1954676.388330] <ffff9d49c00b5000> planet start ****************************************
                 <ffff9d49c00b5000> planet->name: Earth(ffff9d49c00b5000)
                 <ffff9d49c00b5018> planet->mass: 0
                 <ffff9d49c00b5020> planet->rb_node:
                 <ffff9d49c00b5020>     planet->rb_node.__rb_parent_color ffff9d49c0213021, rb_parent: ffff9d49c0213020, rb_color: black
                 <ffff9d49c00b5028>     planet->rb_node.rb_right 0000000000000000
                 <ffff9d49c00b5030>     planet->rb_node.rb_left 0000000000000000
                 <ffff9d49c00b5038> planet end  *****************************************

[1954676.388335] <ffff9d49c0277000> planet start ****************************************
                 <ffff9d49c0277000> planet->name: Mars(ffff9d49c0277000)
                 <ffff9d49c0277018> planet->mass: 4
                 <ffff9d49c0277020> planet->rb_node:
                 <ffff9d49c0277020>     planet->rb_node.__rb_parent_color ffff9d49c0213021, rb_parent: ffff9d49c0213020, rb_color: black
                 <ffff9d49c0277028>     planet->rb_node.rb_right ffff9d49c0279020
                 <ffff9d49c0277030>     planet->rb_node.rb_left ffff9d49c0247020
                 <ffff9d49c0277038> planet end  *****************************************

[1954676.388340] <ffff9d49c0247000> planet start ****************************************
                 <ffff9d49c0247000> planet->name: Saturn(ffff9d49c0247000)
                 <ffff9d49c0247018> planet->mass: 3
                 <ffff9d49c0247020> planet->rb_node:
                 <ffff9d49c0247020>     planet->rb_node.__rb_parent_color ffff9d49c0277020, rb_parent: ffff9d49c0277020, rb_color: red
                 <ffff9d49c0247028>     planet->rb_node.rb_right 0000000000000000
                 <ffff9d49c0247030>     planet->rb_node.rb_left 0000000000000000
                 <ffff9d49c0247038> planet end  *****************************************

[1954676.388344] <ffff9d49c0279000> planet start ****************************************
                 <ffff9d49c0279000> planet->name: Jupiter(ffff9d49c0279000)
                 <ffff9d49c0279018> planet->mass: 7
                 <ffff9d49c0279020> planet->rb_node:
                 <ffff9d49c0279020>     planet->rb_node.__rb_parent_color ffff9d49c0277020, rb_parent: ffff9d49c0277020, rb_color: red
                 <ffff9d49c0279028>     planet->rb_node.rb_right 0000000000000000
                 <ffff9d49c0279030>     planet->rb_node.rb_left 0000000000000000
                 <ffff9d49c0279038> planet end  *****************************************

visualization

rbt-demo.png

畫圖工具：https://www.cs.usfca.edu/~galles/visualization/RedBlack.html

說明

比較迷惑的是__rb_parent_color這個點，它是unsigned long類型，同時包含了parent的指針和當前節(jié)點的color。

由于struct rb_node是long長度對齊的，也就是說在64位機器上rb_node的內(nèi)存起始地址都是8的倍數(shù)，所以地址的最低4個bit要么是0000要么是 1000，最低3bit是用不上的，而color只有0(red)和1(black) 剛好用parent地址的最低1個bit來保存就行了。

這樣就能理解rb_parent和rb_parent這兩個宏了。

至于為什么rb_parent用的 __rb_parent_color & ~3而不是~7應該是為了兼容32位系統(tǒng)。但是為什么不是~1呢？

還有，為什么8字節(jié)對齊，內(nèi)存起始地址就是8的倍數(shù)，??？

?著作權歸作者所有,轉(zhuǎn)載或內(nèi)容合作請聯(lián)系作者
【社區(qū)內(nèi)容提示】社區(qū)部分內(nèi)容疑似由AI輔助生成，瀏覽時請結(jié)合常識與多方信息審慎甄別。
平臺聲明：文章內(nèi)容（如有圖片或視頻亦包括在內(nèi)）由作者上傳并發(fā)布，文章內(nèi)容僅代表作者本人觀點，簡書系信息發(fā)布平臺，僅提供信息存儲服務。

相關閱讀更多精彩內(nèi)容

友情鏈接更多精彩內(nèi)容