linux红黑树 移植

10 篇文章 0 订阅

前言

1. linux(4.3)内核的红黑树代码移植修改,使之能够在项目中编译运行
2. 简单封装接口,使之更好用

kernel.h

提供宏函数

#ifndef __MINE_KERNEL_H__
#define __MINE_KERNEL_H__

#include <stddef.h>

#define container_of(ptr, type, member) ({			\
	const typeof( ((type *)0)->member ) *__mptr = (ptr);	\
	(type *)( (char *)__mptr - offsetof(type,member) );})


#define WRITE_ONCE(var, val) \
    (*((volatile typeof(val) *)(&(var))) = (val))

#endif

linux rbtree修改后源码

头文件rbtree.h

/*
  Red Black Trees
  (C) 1999  Andrea Arcangeli <andrea@suse.de>
  
  This program is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; either version 2 of the License, or
  (at your option) any later version.

  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with this program; if not, write to the Free Software
  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

  linux/include/linux/rbtree.h

  To use rbtrees you'll have to implement your own insert and search cores.
  This will avoid us to use callbacks and to drop drammatically performances.
  I know it's not the cleaner way,  but in C (not in C++) to get
  performances and genericity...

  See Documentation/rbtree.txt for documentation and samples.
*/

#ifndef	__LINUX_RBTREE_H__
#define	__LINUX_RBTREE_H__

#include <stdint.h>
#include "kernel.h"

struct rb_node {
	size_t  __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;
};


#define rb_parent(r)   ((struct rb_node *)((r)->__rb_parent_color & ~3))

#define RB_ROOT	(struct rb_root) { NULL, }
#define	rb_entry(ptr, type, member) container_of(ptr, type, member)

#define RB_EMPTY_ROOT(root)  ((root)->rb_node == NULL)

/* 'empty' nodes are nodes that are known not to be inserted in an rbtree */
#define RB_EMPTY_NODE(node)  \
	((node)->__rb_parent_color == (size_t)(node))
#define RB_CLEAR_NODE(node)  \
	((node)->__rb_parent_color = (size_t)(node))


extern void rb_insert_color(struct rb_node *, struct rb_root *);
extern void rb_erase(struct rb_node *, struct rb_root *);


/* Find logical next and previous nodes in a tree */
extern struct rb_node *rb_next(const struct rb_node *);
extern struct rb_node *rb_prev(const struct rb_node *);
extern struct rb_node *rb_first(const struct rb_root *);
extern struct rb_node *rb_last(const struct rb_root *);

/* Postorder iteration - always visit the parent after its children */
extern struct rb_node *rb_first_postorder(const struct rb_root *);
extern struct rb_node *rb_next_postorder(const struct rb_node *);

/* Fast replacement of a single node without remove/rebalance/add/rebalance */
extern void rb_replace_node(struct rb_node *victim, struct rb_node *new,
			    struct rb_root *root);

static inline void rb_link_node(struct rb_node *node, struct rb_node *parent,
				struct rb_node **rb_link)
{
	node->__rb_parent_color = (size_t)parent;
	node->rb_left = node->rb_right = NULL;

	*rb_link = node;
}

#define rb_entry_safe(ptr, type, member) \
	({ typeof(ptr) ____ptr = (ptr); \
	   ____ptr ? rb_entry(____ptr, type, member) : NULL; \
	})

/**
 * rbtree_postorder_for_each_entry_safe - iterate over rb_root in post order of
 * given type safe against removal of rb_node entry
 *
 * @pos:	the 'type *' to use as a loop cursor.
 * @n:		another 'type *' to use as temporary storage
 * @root:	'rb_root *' of the rbtree.
 * @field:	the name of the rb_node field within 'type'.
 */
#define rbtree_postorder_for_each_entry_safe(pos, n, root, field) \
	for (pos = rb_entry_safe(rb_first_postorder(root), typeof(*pos), field); \
	     pos && ({ n = rb_entry_safe(rb_next_postorder(&pos->field), \
			typeof(*pos), field); 1; }); \
	     pos = n)

#endif	/* _LINUX_RBTREE_H */

rbtree_agumented.h

/*
  Red Black Trees
  (C) 1999  Andrea Arcangeli <andrea@suse.de>
  (C) 2002  David Woodhouse <dwmw2@infradead.org>
  (C) 2012  Michel Lespinasse <walken@google.com>

  This program is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; either version 2 of the License, or
  (at your option) any later version.

  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with this program; if not, write to the Free Software
  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

  linux/include/linux/rbtree_augmented.h
*/

#ifndef __LINUX_RBTREE_AUGMENTED_H__
#define __LINUX_RBTREE_AUGMENTED_H__

#include <stdint.h>
#include "rbtree.h"
/*
 * Please note - only struct rb_augment_callbacks and the prototypes for
 * rb_insert_augmented() and rb_erase_augmented() are intended to be public.
 * The rest are implementation details you are not expected to depend on.
 *
 * See Documentation/rbtree.txt for documentation and samples.
 */

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);
};

extern void __rb_insert_augmented(struct rb_node *node, struct rb_root *root,
	void (*augment_rotate)(struct rb_node *old, struct rb_node *new));
/*
 * Fixup the rbtree and update the augmented information when rebalancing.
 *
 * On insertion, the user must update the augmented information on the path
 * leading to the inserted node, then call rb_link_node() as usual and
 * rb_augment_inserted() instead of the usual rb_insert_color() call.
 * If rb_augment_inserted() rebalances the rbtree, it will callback into
 * a user provided function to update the augmented information on the
 * affected subtrees.
 */
static inline void
rb_insert_augmented(struct rb_node *node, struct rb_root *root,
		    const struct rb_augment_callbacks *augment)
{
	__rb_insert_augmented(node, root, augment->rotate);
}

#define RB_DECLARE_CALLBACKS(rbstatic, rbname, rbstruct, rbfield,	\
			     rbtype, rbaugmented, rbcompute)		\
static inline void							\
rbname ## _propagate(struct rb_node *rb, struct rb_node *stop)		\
{									\
	while (rb != stop) {						\
		rbstruct *node = rb_entry(rb, rbstruct, rbfield);	\
		rbtype augmented = rbcompute(node);			\
		if (node->rbaugmented == augmented)			\
			break;						\
		node->rbaugmented = augmented;				\
		rb = rb_parent(&node->rbfield);				\
	}								\
}									\
static inline void							\
rbname ## _copy(struct rb_node *rb_old, struct rb_node *rb_new)		\
{									\
	rbstruct *old = rb_entry(rb_old, rbstruct, rbfield);		\
	rbstruct *new = rb_entry(rb_new, rbstruct, rbfield);		\
	new->rbaugmented = old->rbaugmented;				\
}									\
static void								\
rbname ## _rotate(struct rb_node *rb_old, struct rb_node *rb_new)	\
{									\
	rbstruct *old = rb_entry(rb_old, rbstruct, rbfield);		\
	rbstruct *new = rb_entry(rb_new, rbstruct, rbfield);		\
	new->rbaugmented = old->rbaugmented;				\
	old->rbaugmented = rbcompute(old);				\
}									\
rbstatic const struct rb_augment_callbacks rbname = {			\
	rbname ## _propagate, rbname ## _copy, rbname ## _rotate	\
};


#define	RB_RED		0
#define	RB_BLACK	1

#define __rb_parent(pc)    ((struct rb_node *)(pc & ~3))

#define __rb_color(pc)     ((pc) & 1)
#define __rb_is_black(pc)  __rb_color(pc)
#define __rb_is_red(pc)    (!__rb_color(pc))
#define rb_color(rb)       __rb_color((rb)->__rb_parent_color)
#define rb_is_red(rb)      __rb_is_red((rb)->__rb_parent_color)
#define rb_is_black(rb)    __rb_is_black((rb)->__rb_parent_color)

static inline void rb_set_parent(struct rb_node *rb, struct rb_node *p)
{
	rb->__rb_parent_color = rb_color(rb) | (size_t)p;
}

static inline void rb_set_parent_color(struct rb_node *rb,
				       struct rb_node *p, int color)
{
	rb->__rb_parent_color = (size_t)p | color;
}

static inline void
__rb_change_child(struct rb_node *old, struct rb_node *new,
		  struct rb_node *parent, struct rb_root *root)
{
	if (parent) {
		if (parent->rb_left == old)
			WRITE_ONCE(parent->rb_left, new);
		else
			WRITE_ONCE(parent->rb_right, new);
	} else
		WRITE_ONCE(root->rb_node, new);
}

extern void __rb_erase_color(struct rb_node *parent, struct rb_root *root,
	void (*augment_rotate)(struct rb_node *old, struct rb_node *new));

static inline struct rb_node *
__rb_erase_augmented(struct rb_node *node, struct rb_root *root,
		     const struct rb_augment_callbacks *augment)
{
	struct rb_node *child = node->rb_right;
	struct rb_node *tmp = node->rb_left;
	struct rb_node *parent, *rebalance;
	size_t pc;

	if (!tmp) {
		/*
		 * Case 1: node to erase has no more than 1 child (easy!)
		 *
		 * Note that if there is one child it must be red due to 5)
		 * and node must be black due to 4). We adjust colors locally
		 * so as to bypass __rb_erase_color() later on.
		 */
		pc = node->__rb_parent_color;
		parent = __rb_parent(pc);
		__rb_change_child(node, child, parent, root);
		if (child) {
			child->__rb_parent_color = pc;
			rebalance = NULL;
		} else
			rebalance = __rb_is_black(pc) ? parent : NULL;
		tmp = parent;
	} else if (!child) {
		/* Still case 1, but this time the child is node->rb_left */
		tmp->__rb_parent_color = pc = node->__rb_parent_color;
		parent = __rb_parent(pc);
		__rb_change_child(node, tmp, parent, root);
		rebalance = NULL;
		tmp = parent;
	} else {
		struct rb_node *successor = child, *child2;

		tmp = child->rb_left;
		if (!tmp) {
			/*
			 * Case 2: node's successor is its right child
			 *
			 *    (n)          (s)
			 *    / \          / \
			 *  (x) (s)  ->  (x) (c)
			 *        \
			 *        (c)
			 */
			parent = successor;
			child2 = successor->rb_right;

			augment->copy(node, successor);
		} else {
			/*
			 * Case 3: node's successor is leftmost under
			 * node's right child subtree
			 *
			 *    (n)          (s)
			 *    / \          / \
			 *  (x) (y)  ->  (x) (y)
			 *      /            /
			 *    (p)          (p)
			 *    /            /
			 *  (s)          (c)
			 *    \
			 *    (c)
			 */
			do {
				parent = successor;
				successor = tmp;
				tmp = tmp->rb_left;
			} while (tmp);
			child2 = successor->rb_right;
			WRITE_ONCE(parent->rb_left, child2);
			WRITE_ONCE(successor->rb_right, child);
			rb_set_parent(child, successor);

			augment->copy(node, successor);
			augment->propagate(parent, successor);
		}

		tmp = node->rb_left;
		WRITE_ONCE(successor->rb_left, tmp);
		rb_set_parent(tmp, successor);

		pc = node->__rb_parent_color;
		tmp = __rb_parent(pc);
		__rb_change_child(node, successor, tmp, root);

		if (child2) {
			successor->__rb_parent_color = pc;
			rb_set_parent_color(child2, parent, RB_BLACK);
			rebalance = NULL;
		} else {
			unsigned long pc2 = successor->__rb_parent_color;
			successor->__rb_parent_color = pc;
			rebalance = __rb_is_black(pc2) ? parent : NULL;
		}
		tmp = successor;
	}

	augment->propagate(tmp, NULL);
	return rebalance;
}

static inline void
rb_erase_augmented(struct rb_node *node, struct rb_root *root,
		   const struct rb_augment_callbacks *augment)
{
	struct rb_node *rebalance = __rb_erase_augmented(node, root, augment);
	if (rebalance)
		__rb_erase_color(rebalance, root, augment->rotate);
}

#endif	/* _LINUX_RBTREE_AUGMENTED_H */

rbtree.c

/*
  Red Black Trees
  (C) 1999  Andrea Arcangeli <andrea@suse.de>
  (C) 2002  David Woodhouse <dwmw2@infradead.org>
  (C) 2012  Michel Lespinasse <walken@google.com>

  This program is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; either version 2 of the License, or
  (at your option) any later version.

  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with this program; if not, write to the Free Software
  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

  linux/lib/rbtree.c
*/
#include <stddef.h>
#include <stdint.h>
#include "rbtree_augmented.h"

/*
 * red-black trees properties:  http://en.wikipedia.org/wiki/Rbtree
 *
 *  1) A node is either red or black
 *  2) The root is black
 *  3) All leaves (NULL) are black
 *  4) Both children of every red node are black
 *  5) Every simple path from root to leaves contains the same number
 *     of black nodes.
 *
 *  4 and 5 give the O(log n) guarantee, since 4 implies you cannot have two
 *  consecutive red nodes in a path and every red node is therefore followed by
 *  a black. So if B is the number of black nodes on every simple path (as per
 *  5), then the longest possible path due to 4 is 2B.
 *
 *  We shall indicate color with case, where black nodes are uppercase and red
 *  nodes will be lowercase. Unknown color nodes shall be drawn as red within
 *  parentheses and have some accompanying text comment.
 */

/*
 * Notes on lockless lookups:
 *
 * All stores to the tree structure (rb_left and rb_right) must be done using
 * WRITE_ONCE(). And we must not inadvertently cause (temporary) loops in the
 * tree structure as seen in program order.
 *
 * These two requirements will allow lockless iteration of the tree -- not
 * correct iteration mind you, tree rotations are not atomic so a lookup might
 * miss entire subtrees.
 *
 * But they do guarantee that any such traversal will only see valid elements
 * and that it will indeed complete -- does not get stuck in a loop.
 *
 * It also guarantees that if the lookup returns an element it is the 'correct'
 * one. But not returning an element does _NOT_ mean it's not present.
 *
 * NOTE:
 *
 * Stores to __rb_parent_color are not important for simple lookups so those
 * are left undone as of now. Nor did I check for loops involving parent
 * pointers.
 */

static inline void rb_set_black(struct rb_node *rb)
{
	rb->__rb_parent_color |= RB_BLACK;
}

static inline struct rb_node *rb_red_parent(struct rb_node *red)
{
	return (struct rb_node *)red->__rb_parent_color;
}

/*
 * Helper function for rotations:
 * - old's parent and color get assigned to new
 * - old gets assigned new as a parent and 'color' as a color.
 */
static inline void
__rb_rotate_set_parents(struct rb_node *old, struct rb_node *new,
			struct rb_root *root, int color)
{
	struct rb_node *parent = rb_parent(old);
	new->__rb_parent_color = old->__rb_parent_color;
	rb_set_parent_color(old, new, color);
	__rb_change_child(old, new, parent, root);
}

static inline void
__rb_insert(struct rb_node *node, struct rb_root *root,
	    void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
{
	struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;

	while (1) {
		/*
		 * Loop invariant: node is red
		 *
		 * If there is a black parent, we are done.
		 * Otherwise, take some corrective action as we don't
		 * want a red root or two consecutive red nodes.
		 */
		if (!parent) {
			rb_set_parent_color(node, NULL, RB_BLACK);
			break;
		} else if (rb_is_black(parent))
			break;

		gparent = rb_red_parent(parent);

		tmp = gparent->rb_right;
		if (parent != tmp) {	/* parent == gparent->rb_left */
			if (tmp && rb_is_red(tmp)) {
				/*
				 * Case 1 - color flips
				 *
				 *       G            g
				 *      / \          / \
				 *     p   u  -->   P   U
				 *    /            /
				 *   n            n
				 *
				 * However, since g's parent might be red, and
				 * 4) does not allow this, we need to recurse
				 * at g.
				 */
				rb_set_parent_color(tmp, gparent, RB_BLACK);
				rb_set_parent_color(parent, gparent, RB_BLACK);
				node = gparent;
				parent = rb_parent(node);
				rb_set_parent_color(node, parent, RB_RED);
				continue;
			}

			tmp = parent->rb_right;
			if (node == tmp) {
				/*
				 * Case 2 - left rotate at parent
				 *
				 *      G             G
				 *     / \           / \
				 *    p   U  -->    n   U
				 *     \           /
				 *      n         p
				 *
				 * This still leaves us in violation of 4), the
				 * continuation into Case 3 will fix that.
				 */
				tmp = node->rb_left;
				WRITE_ONCE(parent->rb_right, tmp);
				WRITE_ONCE(node->rb_left, parent);
				if (tmp)
					rb_set_parent_color(tmp, parent,
							    RB_BLACK);
				rb_set_parent_color(parent, node, RB_RED);
				augment_rotate(parent, node);
				parent = node;
				tmp = node->rb_right;
			}

			/*
			 * Case 3 - right rotate at gparent
			 *
			 *        G           P
			 *       / \         / \
			 *      p   U  -->  n   g
			 *     /                 \
			 *    n                   U
			 */
			WRITE_ONCE(gparent->rb_left, tmp); /* == parent->rb_right */
			WRITE_ONCE(parent->rb_right, gparent);
			if (tmp)
				rb_set_parent_color(tmp, gparent, RB_BLACK);
			__rb_rotate_set_parents(gparent, parent, root, RB_RED);
			augment_rotate(gparent, parent);
			break;
		} else {
			tmp = gparent->rb_left;
			if (tmp && rb_is_red(tmp)) {
				/* Case 1 - color flips */
				rb_set_parent_color(tmp, gparent, RB_BLACK);
				rb_set_parent_color(parent, gparent, RB_BLACK);
				node = gparent;
				parent = rb_parent(node);
				rb_set_parent_color(node, parent, RB_RED);
				continue;
			}

			tmp = parent->rb_left;
			if (node == tmp) {
				/* Case 2 - right rotate at parent */
				tmp = node->rb_right;
				WRITE_ONCE(parent->rb_left, tmp);
				WRITE_ONCE(node->rb_right, parent);
				if (tmp)
					rb_set_parent_color(tmp, parent,
							    RB_BLACK);
				rb_set_parent_color(parent, node, RB_RED);
				augment_rotate(parent, node);
				parent = node;
				tmp = node->rb_left;
			}

			/* Case 3 - left rotate at gparent */
			WRITE_ONCE(gparent->rb_right, tmp); /* == parent->rb_left */
			WRITE_ONCE(parent->rb_left, gparent);
			if (tmp)
				rb_set_parent_color(tmp, gparent, RB_BLACK);
			__rb_rotate_set_parents(gparent, parent, root, RB_RED);
			augment_rotate(gparent, parent);
			break;
		}
	}
}

/*
 * Inline version for rb_erase() use - we want to be able to inline
 * and eliminate the dummy_rotate callback there
 */
static inline void
____rb_erase_color(struct rb_node *parent, struct rb_root *root,
	void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
{
	struct rb_node *node = NULL, *sibling, *tmp1, *tmp2;

	while (1) {
		/*
		 * Loop invariants:
		 * - node is black (or NULL on first iteration)
		 * - node is not the root (parent is not NULL)
		 * - All leaf paths going through parent and node have a
		 *   black node count that is 1 lower than other leaf paths.
		 */
		sibling = parent->rb_right;
		if (node != sibling) {	/* node == parent->rb_left */
			if (rb_is_red(sibling)) {
				/*
				 * Case 1 - left rotate at parent
				 *
				 *     P               S
				 *    / \             / \
				 *   N   s    -->    p   Sr
				 *      / \         / \
				 *     Sl  Sr      N   Sl
				 */
				tmp1 = sibling->rb_left;
				WRITE_ONCE(parent->rb_right, tmp1);
				WRITE_ONCE(sibling->rb_left, parent);
				rb_set_parent_color(tmp1, parent, RB_BLACK);
				__rb_rotate_set_parents(parent, sibling, root,
							RB_RED);
				augment_rotate(parent, sibling);
				sibling = tmp1;
			}
			tmp1 = sibling->rb_right;
			if (!tmp1 || rb_is_black(tmp1)) {
				tmp2 = sibling->rb_left;
				if (!tmp2 || rb_is_black(tmp2)) {
					/*
					 * Case 2 - sibling color flip
					 * (p could be either color here)
					 *
					 *    (p)           (p)
					 *    / \           / \
					 *   N   S    -->  N   s
					 *      / \           / \
					 *     Sl  Sr        Sl  Sr
					 *
					 * This leaves us violating 5) which
					 * can be fixed by flipping p to black
					 * if it was red, or by recursing at p.
					 * p is red when coming from Case 1.
					 */
					rb_set_parent_color(sibling, parent,
							    RB_RED);
					if (rb_is_red(parent))
						rb_set_black(parent);
					else {
						node = parent;
						parent = rb_parent(node);
						if (parent)
							continue;
					}
					break;
				}
				/*
				 * Case 3 - right rotate at sibling
				 * (p could be either color here)
				 *
				 *   (p)           (p)
				 *   / \           / \
				 *  N   S    -->  N   Sl
				 *     / \             \
				 *    sl  Sr            s
				 *                       \
				 *                        Sr
				 */
				tmp1 = tmp2->rb_right;
				WRITE_ONCE(sibling->rb_left, tmp1);
				WRITE_ONCE(tmp2->rb_right, sibling);
				WRITE_ONCE(parent->rb_right, tmp2);
				if (tmp1)
					rb_set_parent_color(tmp1, sibling,
							    RB_BLACK);
				augment_rotate(sibling, tmp2);
				tmp1 = sibling;
				sibling = tmp2;
			}
			/*
			 * Case 4 - left rotate at parent + color flips
			 * (p and sl could be either color here.
			 *  After rotation, p becomes black, s acquires
			 *  p's color, and sl keeps its color)
			 *
			 *      (p)             (s)
			 *      / \             / \
			 *     N   S     -->   P   Sr
			 *        / \         / \
			 *      (sl) sr      N  (sl)
			 */
			tmp2 = sibling->rb_left;
			WRITE_ONCE(parent->rb_right, tmp2);
			WRITE_ONCE(sibling->rb_left, parent);
			rb_set_parent_color(tmp1, sibling, RB_BLACK);
			if (tmp2)
				rb_set_parent(tmp2, parent);
			__rb_rotate_set_parents(parent, sibling, root,
						RB_BLACK);
			augment_rotate(parent, sibling);
			break;
		} else {
			sibling = parent->rb_left;
			if (rb_is_red(sibling)) {
				/* Case 1 - right rotate at parent */
				tmp1 = sibling->rb_right;
				WRITE_ONCE(parent->rb_left, tmp1);
				WRITE_ONCE(sibling->rb_right, parent);
				rb_set_parent_color(tmp1, parent, RB_BLACK);
				__rb_rotate_set_parents(parent, sibling, root,
							RB_RED);
				augment_rotate(parent, sibling);
				sibling = tmp1;
			}
			tmp1 = sibling->rb_left;
			if (!tmp1 || rb_is_black(tmp1)) {
				tmp2 = sibling->rb_right;
				if (!tmp2 || rb_is_black(tmp2)) {
					/* Case 2 - sibling color flip */
					rb_set_parent_color(sibling, parent,
							    RB_RED);
					if (rb_is_red(parent))
						rb_set_black(parent);
					else {
						node = parent;
						parent = rb_parent(node);
						if (parent)
							continue;
					}
					break;
				}
				/* Case 3 - right rotate at sibling */
				tmp1 = tmp2->rb_left;
				WRITE_ONCE(sibling->rb_right, tmp1);
				WRITE_ONCE(tmp2->rb_left, sibling);
				WRITE_ONCE(parent->rb_left, tmp2);
				if (tmp1)
					rb_set_parent_color(tmp1, sibling,
							    RB_BLACK);
				augment_rotate(sibling, tmp2);
				tmp1 = sibling;
				sibling = tmp2;
			}
			/* Case 4 - left rotate at parent + color flips */
			tmp2 = sibling->rb_right;
			WRITE_ONCE(parent->rb_left, tmp2);
			WRITE_ONCE(sibling->rb_right, parent);
			rb_set_parent_color(tmp1, sibling, RB_BLACK);
			if (tmp2)
				rb_set_parent(tmp2, parent);
			__rb_rotate_set_parents(parent, sibling, root,
						RB_BLACK);
			augment_rotate(parent, sibling);
			break;
		}
	}
}

/* Non-inline version for rb_erase_augmented() use */
void __rb_erase_color(struct rb_node *parent, struct rb_root *root,
	void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
{
	____rb_erase_color(parent, root, augment_rotate);
}

/*
 * Non-augmented rbtree manipulation functions.
 *
 * We use dummy augmented callbacks here, and have the compiler optimize them
 * out of the rb_insert_color() and rb_erase() function definitions.
 */

static inline void dummy_propagate(struct rb_node *node, struct rb_node *stop) {}
static inline void dummy_copy(struct rb_node *old, struct rb_node *new) {}
static inline void dummy_rotate(struct rb_node *old, struct rb_node *new) {}

static const struct rb_augment_callbacks dummy_callbacks = {
	dummy_propagate, dummy_copy, dummy_rotate
};

void rb_insert_color(struct rb_node *node, struct rb_root *root)
{
	__rb_insert(node, root, dummy_rotate);
}

void rb_erase(struct rb_node *node, struct rb_root *root)
{
	struct rb_node *rebalance;
	rebalance = __rb_erase_augmented(node, root, &dummy_callbacks);
	if (rebalance)
		____rb_erase_color(rebalance, root, dummy_rotate);
}

/*
 * Augmented rbtree manipulation functions.
 *
 * This instantiates the same inline functions as in the non-augmented
 * case, but this time with user-defined callbacks.
 */

void __rb_insert_augmented(struct rb_node *node, struct rb_root *root,
	void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
{
	__rb_insert(node, root, augment_rotate);
}

/*
 * This function returns the first node (in sort order) of the tree.
 */
struct rb_node *rb_first(const struct rb_root *root)
{
	struct rb_node	*n;

	n = root->rb_node;
	if (!n)
		return NULL;
	while (n->rb_left)
		n = n->rb_left;
	return n;
}

struct rb_node *rb_last(const struct rb_root *root)
{
	struct rb_node	*n;

	n = root->rb_node;
	if (!n)
		return NULL;
	while (n->rb_right)
		n = n->rb_right;
	return n;
}

struct rb_node *rb_next(const struct rb_node *node)
{
	struct rb_node *parent;

	if (RB_EMPTY_NODE(node))
		return NULL;

	/*
	 * If we have a right-hand child, go down and then left as far
	 * as we can.
	 */
	if (node->rb_right) {
		node = node->rb_right; 
		while (node->rb_left)
			node=node->rb_left;
		return (struct rb_node *)node;
	}

	/*
	 * No right-hand children. Everything down and left is smaller than us,
	 * so any 'next' node must be in the general direction of our parent.
	 * Go up the tree; any time the ancestor is a right-hand child of its
	 * parent, keep going up. First time it's a left-hand child of its
	 * parent, said parent is our 'next' node.
	 */
	while ((parent = rb_parent(node)) && node == parent->rb_right)
		node = parent;

	return parent;
}

struct rb_node *rb_prev(const struct rb_node *node)
{
	struct rb_node *parent;

	if (RB_EMPTY_NODE(node))
		return NULL;

	/*
	 * If we have a left-hand child, go down and then right as far
	 * as we can.
	 */
	if (node->rb_left) {
		node = node->rb_left; 
		while (node->rb_right)
			node=node->rb_right;
		return (struct rb_node *)node;
	}

	/*
	 * No left-hand children. Go up till we find an ancestor which
	 * is a right-hand child of its parent.
	 */
	while ((parent = rb_parent(node)) && node == parent->rb_left)
		node = parent;

	return parent;
}

void rb_replace_node(struct rb_node *victim, struct rb_node *new,
		     struct rb_root *root)
{
	struct rb_node *parent = rb_parent(victim);

	/* Set the surrounding nodes to point to the replacement */
	__rb_change_child(victim, new, parent, root);
	if (victim->rb_left)
		rb_set_parent(victim->rb_left, new);
	if (victim->rb_right)
		rb_set_parent(victim->rb_right, new);

	/* Copy the pointers/colour from the victim to the replacement */
	*new = *victim;
}

static struct rb_node *rb_left_deepest_node(const struct rb_node *node)
{
	for (;;) {
		if (node->rb_left)
			node = node->rb_left;
		else if (node->rb_right)
			node = node->rb_right;
		else
			return (struct rb_node *)node;
	}
}

struct rb_node *rb_next_postorder(const struct rb_node *node)
{
	const struct rb_node *parent;
	if (!node)
		return NULL;
	parent = rb_parent(node);

	/* If we're sitting on node, we've already seen our children */
	if (parent && node == parent->rb_left && parent->rb_right) {
		/* If we are the parent's left node, go to the parent's right
		 * node then all the way down to the left */
		return rb_left_deepest_node(parent->rb_right);
	} else
		/* Otherwise we are the parent's right node, and the parent
		 * should be next */
		return (struct rb_node *)parent;
}

struct rb_node *rb_first_postorder(const struct rb_root *root)
{
	if (!root->rb_node)
		return NULL;

	return rb_left_deepest_node(root->rb_node);
}

封装

rbtree_api.h

#ifndef __RBTREE_API_H__
#define __RBTREE_API_H__

#include "rbtree.h"
#include "rbtree_augmented.h"

typedef struct _rbtree_node {
    struct rb_node node;
    void *key;
    void* val;
} rbtree_node;

struct _rbtree_root;
typedef struct _rbtree_root rbtree_root;

typedef size_t(*rbtree_hash_func)(const void *key);
typedef int(*rbtree_cmp_func)(const void *my_key, const void *other_key);
typedef void*(*rbtree_malloc_func)(size_t len);
typedef void(*rbtree_free_func)(void *ptr);
typedef void (*rbtree_release_func)(void *key, void *val, rbtree_root *root);

typedef struct _rbtree_root {
    struct rb_root root;
    rbtree_hash_func hash;
    rbtree_cmp_func cmp;
    rbtree_release_func release;
    rbtree_malloc_func my_malloc;
    rbtree_free_func my_free;
    int count;
} rbtree_root;

/*******************  以下接口不负责malloc/free rbtree_node  ****************/
// 字符串hash
size_t rbtree_string_hash(const void *key);
// 默认hash
size_t rbtree_hash_default(const void *key);
// 默认释放节点资源的函数
void rbtree_release_default(void *key, void *val, rbtree_root *root);
// (https://github.com/gcc-mirror/gcc/blob/master/libstdc++-v3/libsupc++/hash_bytes.cc)
size_t hash_api(const void* ptr, size_t len, size_t seed);

// 初始化根节点
int32_t rbtree_init(rbtree_root *root,
                            rbtree_hash_func hash_func,
                            rbtree_cmp_func cmp_func,
                            rbtree_release_func release_func);
// 销毁跟根节点
void rbtree_destroy(rbtree_root *root);
// 查找节点
rbtree_node* rbtree_find(const void *key, rbtree_root *root);
// 插入节点
int32_t rbtree_insert(rbtree_node *entry, rbtree_root *root);
// 去掉节点
int32_t rbtree_remove(rbtree_node *entry, rbtree_root *root);
// 去掉节点
void rbtree_free(rbtree_node *entry, rbtree_root *root);
// 去掉节点并释放节点
int32_t rbtree_erase(rbtree_node *node, rbtree_root *root);

/*******************  以下接口负责malloc/free rbtree_node  ****************/
// 初始化根节点
int32_t rbapi_init(rbtree_root *root,
                            rbtree_hash_func hash_func,
                            rbtree_cmp_func cmp_func,
                            rbtree_release_func release_func,
                            rbtree_malloc_func malloc_func,
                            rbtree_free_func free_func);
// 销毁键值对
void rbapi_destroy(rbtree_root *root);
// 清空节点
void rbapi_clear(rbtree_root *root);
// 遍历节点
int32_t rbapi_travel(void *data, rbtree_root *root, 
                            int32_t (*travel)(void *data, rbtree_node *node, rbtree_root *root));
// 获取键值
int32_t rbapi_find(const void *key, void *val, rbtree_root *root);
// 插入键值对
int32_t rbapi_insert(void *key, void *val, rbtree_root *root);
// 移除键值对
rbtree_node* rbapi_remove(const void *key, rbtree_root *root);
// 释放键值对, 和remove配套使用
void rbapi_free(rbtree_node *entry, rbtree_root *root);
// 移除且释放键值对
int32_t rbapi_erase(const void *key, rbtree_root *root);
// 替换键值对
int32_t rbapi_replace(void *key, void *val, rbtree_root *root);


#endif


rbtree_api.c

#include <stdlib.h>
#include <string.h>

#include "rbtree_api.h"

uint32_t buf2uint32(const uint8_t *ptr)
{
    uint32_t val = 0;
    for (int i = 0; i < 4; i++) {
        val |= ptr[i];
        val = val << 8;
    }
    return val;
}

// ptr必须4字节对齐
size_t hash_api(const void* ptr, size_t len, size_t seed)
{
    const size_t m = 0x5bd1e995;
    size_t hash = seed ^ len;
    const char* buf = (const char*)(ptr);

    // Mix 4 bytes at a time into the hash.
    while (len >= 4)
    {
        // buf地址非4字节对齐将异常
        uint32_t k = buf2uint32((const uint8_t *)buf);
        k *= m;
        k ^= k >> 24;
        k *= m;
        hash *= m;
        hash ^= k;
        buf += 4;
        len -= 4;
    }

    // Handle the last few bytes of the input array.
    switch (len)
    {
        case 3:
        hash ^= (unsigned char)(buf[2]) << 16;

        case 2:
        hash ^= (unsigned char)(buf[1]) << 8;

        case 1:
        hash ^= (unsigned char)(buf[0]);
        hash *= m;
    };

    // Do a few final mixes of the hash.
    hash ^= hash >> 13;
    hash *= m;
    hash ^= hash >> 15;
    return hash;
}

size_t rbtree_hash_default(const void *key)
{
    size_t val = (size_t) key;
    return (size_t)val;
}

size_t rbtree_string_hash(const void *key)
{
    return hash_api(key, strlen((char*)key), 0xaa55aa55);
}

void rbtree_release_default(void *key, void *val, rbtree_root *root)
{
    return;
}

static int rbtree_cmp(const void *my_key, const void *other_key, rbtree_root *root)
{
    if (!root->cmp) {
        return 0;
    }

    return root->cmp(my_key, other_key);
}

int32_t rbtree_travel(void *data, rbtree_root *root, 
                            int32_t (*travel)(void *data, rbtree_node *node, rbtree_root *root))
{
	struct rb_node *iter;
    struct rb_node *next = NULL;
    rbtree_node *node = NULL;
	for (iter = rb_first(&root->root); iter; iter = next) {
        next = rb_next(iter);
        node = container_of(iter, rbtree_node, node);
        if (travel(data, node, root) < 0) {
            return -1;
        }
	}
    return 0;
}

int32_t rbtree_clear(rbtree_root *root)
{
    rbtree_node *entry = NULL;
    for (struct rb_node *node = rb_first(&root->root); node; node = rb_first(&root->root)) {
        entry = container_of(node, rbtree_node, node);
        rbtree_erase(entry, root);
	}
    return 0;
}

rbtree_node* rbtree_find(const void *key, rbtree_root *root)
{
    int cmp_ret = 0;
    size_t hash_cmp = 0;
    rbtree_node *node = NULL;
    
    rbtree_hash_func hash = root->hash;
    struct rb_node **new_node = &root->root.rb_node, *parent = NULL;

    size_t hash_mine = hash(key);

	while (*new_node) {
		parent = *new_node;
        node = container_of(parent, rbtree_node, node);

        hash_cmp = hash(node->key);

        if (hash_mine < hash_cmp) {
			new_node = &parent->rb_left;

        } else if (hash_mine > hash_cmp) {
			new_node = &parent->rb_right;

        } else {
            cmp_ret = rbtree_cmp(key, node->key, root);
            if (!cmp_ret) {
                return node;
            }
			// 将导致多计算一次hash
            if (cmp_ret < 0) {
                new_node = &parent->rb_left;
            } else {
                new_node = &parent->rb_right;
            }
        }
	}
    return NULL;
}

int32_t rbtree_insert(rbtree_node *entry, rbtree_root *root)
{
    int cmp_ret = 0;
    size_t hash_cmp = 0;
    rbtree_node *node = NULL;
    
    rbtree_hash_func hash = root->hash;
    struct rb_node **new_node = &root->root.rb_node, *parent = NULL;

    size_t hash_mine = hash(entry->key);

	while (*new_node) {
		parent = *new_node;
        node = container_of(*new_node, rbtree_node, node);

        hash_cmp = hash(node->key);

        if (hash_mine < hash_cmp) {
			new_node = &parent->rb_left;

        } else if (hash_mine > hash_cmp) {
			new_node = &parent->rb_right;

        } else {
            cmp_ret = rbtree_cmp(entry->key, node->key, root);
            if (!cmp_ret) {
                return -1;
            }
			// 将导致多计算一次hash
            if (cmp_ret < 0) {
                new_node = &parent->rb_left;
            } else {
                new_node = &parent->rb_right;
            }
        }
	}

	rb_link_node(&entry->node, parent, new_node);
	rb_insert_color(&entry->node, &root->root);

    root->count ++;
    return 0;
}

int32_t rbtree_remove(rbtree_node *entry, rbtree_root *root)
{
    rb_erase(&entry->node, &root->root);
    root->count --;
    return 0;
}

void rbtree_free(rbtree_node *node, rbtree_root *root)
{
    if (root->release) {
        root->release(node->key, node->val, root);
    }
    node->key = NULL;
    node->val = NULL;
}

int32_t rbtree_erase(rbtree_node *node, rbtree_root *root)
{
    rbtree_remove(node, root);
    rbtree_free(node, root);
    return 0;
}

int32_t rbtree_init(rbtree_root *root,
                        rbtree_hash_func hash_func,
                        rbtree_cmp_func cmp_func,
                        rbtree_release_func release_func)
{
    if (!root) {
        return -1;
    }

    memset(root, 0, sizeof(*root));
    root->hash = hash_func ? hash_func : rbtree_hash_default;
    root->cmp = cmp_func ? cmp_func : NULL;
    root->release = release_func ? release_func : rbtree_release_default;
    root->count = 0;
    return 0;
}

void rbtree_destroy(rbtree_root *root)
{
    if (!root) {
        return ;
    }

    rbtree_clear(root);
}


int32_t rbapi_travel(void *data, rbtree_root *root, 
                            int32_t (*travel)(void *data, rbtree_node *node, rbtree_root *root))
{
	struct rb_node *iter;
    struct rb_node *next = NULL;
    rbtree_node *node = NULL;
	for (iter = rb_first(&root->root); iter; iter = next) {
        next = rb_next(iter);
        node = container_of(iter, rbtree_node, node);
        if (travel(data, node, root) < 0) {
            return -1;
        }
	}
    return 0;
}

int32_t rbapi_find(const void *key, void *val, rbtree_root *root)
{
    rbtree_node *node = rbtree_find(key, root);
    if (!node) {
        return -1;
    }
	if (val) {
	    void **pval = (void **)val;
	    //*pval = node->val;
	    memcpy(pval, &node->val, sizeof(*pval));
    }
    return 0;
}

rbtree_node* rbapi_remove(const void *key, rbtree_root *root)
{   
    rbtree_node * node = rbtree_find(key, root);
    if (!node) {
        return NULL;
    }
    
    rbtree_remove(node, root);
    return node;
}

void rbtapi_free(rbtree_node *node, rbtree_root *root)
{
    rbtree_free(node, root);
    root->my_free(node);
    // free(node);
}

int32_t rbapi_insert(void *key, void *val, rbtree_root *root)
{
    // rbtree_node *entry = (rbtree_node *)malloc(sizeof(rbtree_node));
    rbtree_node *entry = (rbtree_node *)root->my_malloc(sizeof(rbtree_node));
    if (!entry) {
        return -1;
    }
    memset(entry, 0, sizeof(*entry));

    entry->key = key;
    entry->val = val;
    
    if (rbtree_insert(entry, root) < 0) {
    	root->myfree(entry);
    	return -1;
	}
	return 0;
}

int32_t rbapi_erase(const void *key, rbtree_root *root)
{
    rbtree_node * node = rbtree_find(key, root);
    if (!node) {
        return -1;
    }
    rbtree_erase(node, root);
    root->my_free(node);
    // free(node);
    return 0;
}

int32_t rbapi_replace(void *key, void *val, rbtree_root *root)
{
    rbapi_erase(key, root);
    return rbapi_insert(key, val, root);
}

void rbapi_clear(rbtree_root *root)
{
    rbtree_node *entry = NULL;
    for (struct rb_node *node = rb_first(&root->root); node; node = rb_first(&root->root)) {
        entry = container_of(node, rbtree_node, node);
        rbtree_erase(entry, root);
        root->my_free(entry);
	}
}

int32_t rbapi_init(rbtree_root *root,
                            rbtree_hash_func hash_func,
                            rbtree_cmp_func cmp_func,
                            rbtree_release_func release_func,
                            rbtree_malloc_func my_malloc,
                            rbtree_free_func my_free)
{
    if (!root) {
        return -1;
    }

    // memset清空
    rbtree_init(root, hash_func, cmp_func, release_func);

    root->my_malloc = my_malloc ? my_malloc : malloc;
    root->my_free = my_free ? my_free : free;
    return 0;
}

void rbapi_destroy(rbtree_root *root)
{
    rbapi_clear(root);
}

用例

#include <stdio.h>
#include "rbtree_api.h"
int main() {
    rbtree_root root;
    /**
    * 第二参数:计算字符串的hash值,rbtree_string_hash 
    * 第三参数:比较key值的函数,当hash重复时,会调用该函数,传NULL时hash重复认为相等
    * 第四参数:释放节点函数,当调用erase或者clear释放节点时,调用该函数,释放内存
    * 第五/第六参数:malloc和free,如果不传使用默认的malloc和free
    */
    rbapi_init(&root, rbtree_string_hash, NULL, NULL, NULL, NULL);

    rbapi_insert("1", "11", &root);
    rbapi_insert("2", "22", &root);
    rbapi_insert("3", "33", &root);
    rbapi_insert("4", "44", &root);
    rbapi_insert("5", "55", &root);

    char* val = NULL;
    rbapi_find("1", &val, &root);
    printf("val = %s\n", val);

    rbapi_find("2", &val, &root);
    printf("val = %s\n", val);

    rbapi_find("3", &val, &root);
    printf("val = %s\n", val);

    rbapi_find("4", &val, &root);
    printf("val = %s\n", val);

    rbapi_find("5", &val, &root);
    printf("val = %s\n", val);

    rbapi_clear(&root);


    int ret = rbapi_find("1", &val, &root);
    printf("ret = %d\n", ret);

    return 0;
}

  • 1
    点赞
  • 0
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值