[linux-block.git] / lib / rbtree_test.c

#include <linux/module.h>
#include <linux/moduleparam.h>
#include <linux/rbtree_augmented.h>
#include <linux/random.h>
#include <linux/slab.h>
#include <asm/timex.h>

#define __param(type, name, init, msg)		\
	static type name = init;		\
	module_param(name, type, 0444);		\
	MODULE_PARM_DESC(name, msg);

__param(int, nnodes, 100, "Number of nodes in the rb-tree");
__param(int, perf_loops, 100000, "Number of iterations modifying the rb-tree");
__param(int, check_loops, 100, "Number of iterations modifying and verifying the rb-tree");

struct test_node {
	u32 key;
	struct rb_node rb;

	/* following fields used for testing augmented rbtree functionality */
	u32 val;
	u32 augmented;
};

static struct rb_root root = RB_ROOT;
static struct test_node *nodes = NULL;

static struct rnd_state rnd;

static void insert(struct test_node *node, struct rb_root *root)
{
	struct rb_node **new = &root->rb_node, *parent = NULL;
	u32 key = node->key;

	while (*new) {
		parent = *new;
		if (key < rb_entry(parent, struct test_node, rb)->key)
			new = &parent->rb_left;
		else
			new = &parent->rb_right;
	}

	rb_link_node(&node->rb, parent, new);
	rb_insert_color(&node->rb, root);
}

static inline void erase(struct test_node *node, struct rb_root *root)
{
	rb_erase(&node->rb, root);
}

static inline u32 augment_recompute(struct test_node *node)
{
	u32 max = node->val, child_augmented;
	if (node->rb.rb_left) {
		child_augmented = rb_entry(node->rb.rb_left, struct test_node,
					   rb)->augmented;
		if (max < child_augmented)
			max = child_augmented;
	}
	if (node->rb.rb_right) {
		child_augmented = rb_entry(node->rb.rb_right, struct test_node,
					   rb)->augmented;
		if (max < child_augmented)
			max = child_augmented;
	}
	return max;
}

RB_DECLARE_CALLBACKS(static, augment_callbacks, struct test_node, rb,
		     u32, augmented, augment_recompute)

static void insert_augmented(struct test_node *node, struct rb_root *root)
{
	struct rb_node **new = &root->rb_node, *rb_parent = NULL;
	u32 key = node->key;
	u32 val = node->val;
	struct test_node *parent;

	while (*new) {
		rb_parent = *new;
		parent = rb_entry(rb_parent, struct test_node, rb);
		if (parent->augmented < val)
			parent->augmented = val;
		if (key < parent->key)
			new = &parent->rb.rb_left;
		else
			new = &parent->rb.rb_right;
	}

	node->augmented = val;
	rb_link_node(&node->rb, rb_parent, new);
	rb_insert_augmented(&node->rb, root, &augment_callbacks);
}

static void erase_augmented(struct test_node *node, struct rb_root *root)
{
	rb_erase_augmented(&node->rb, root, &augment_callbacks);
}

static void init(void)
{
	int i;
	for (i = 0; i < nnodes; i++) {
		nodes[i].key = prandom_u32_state(&rnd);
		nodes[i].val = prandom_u32_state(&rnd);
	}
}

static bool is_red(struct rb_node *rb)
{
	return !(rb->__rb_parent_color & 1);
}

static int black_path_count(struct rb_node *rb)
{
	int count;
	for (count = 0; rb; rb = rb_parent(rb))
		count += !is_red(rb);
	return count;
}

static void check_postorder_foreach(int nr_nodes)
{
	struct test_node *cur, *n;
	int count = 0;
	rbtree_postorder_for_each_entry_safe(cur, n, &root, rb)
		count++;

	WARN_ON_ONCE(count != nr_nodes);
}

static void check_postorder(int nr_nodes)
{
	struct rb_node *rb;
	int count = 0;
	for (rb = rb_first_postorder(&root); rb; rb = rb_next_postorder(rb))
		count++;

	WARN_ON_ONCE(count != nr_nodes);
}

static void check(int nr_nodes)
{
	struct rb_node *rb;
	int count = 0, blacks = 0;
	u32 prev_key = 0;

	for (rb = rb_first(&root); rb; rb = rb_next(rb)) {
		struct test_node *node = rb_entry(rb, struct test_node, rb);
		WARN_ON_ONCE(node->key < prev_key);
		WARN_ON_ONCE(is_red(rb) &&
			     (!rb_parent(rb) || is_red(rb_parent(rb))));
		if (!count)
			blacks = black_path_count(rb);
		else
			WARN_ON_ONCE((!rb->rb_left || !rb->rb_right) &&
				     blacks != black_path_count(rb));
		prev_key = node->key;
		count++;
	}

	WARN_ON_ONCE(count != nr_nodes);
	WARN_ON_ONCE(count < (1 << black_path_count(rb_last(&root))) - 1);

	check_postorder(nr_nodes);
	check_postorder_foreach(nr_nodes);
}

static void check_augmented(int nr_nodes)
{
	struct rb_node *rb;

	check(nr_nodes);
	for (rb = rb_first(&root); rb; rb = rb_next(rb)) {
		struct test_node *node = rb_entry(rb, struct test_node, rb);
		WARN_ON_ONCE(node->augmented != augment_recompute(node));
	}
}

static int __init rbtree_test_init(void)
{
	int i, j;
	cycles_t time1, time2, time;

	nodes = kmalloc(nnodes * sizeof(*nodes), GFP_KERNEL);
	if (!nodes)
		return -ENOMEM;

	printk(KERN_ALERT "rbtree testing");

	prandom_seed_state(&rnd, 3141592653589793238ULL);
	init();

	time1 = get_cycles();

	for (i = 0; i < perf_loops; i++) {
		for (j = 0; j < nnodes; j++)
			insert(nodes + j, &root);
		for (j = 0; j < nnodes; j++)
			erase(nodes + j, &root);
	}

	time2 = get_cycles();
	time = time2 - time1;

	time = div_u64(time, perf_loops);
	printk(" -> %llu cycles\n", (unsigned long long)time);

	for (i = 0; i < check_loops; i++) {
		init();
		for (j = 0; j < nnodes; j++) {
			check(j);
			insert(nodes + j, &root);
		}
		for (j = 0; j < nnodes; j++) {
			check(nnodes - j);
			erase(nodes + j, &root);
		}
		check(0);
	}

	printk(KERN_ALERT "augmented rbtree testing");

	init();

	time1 = get_cycles();

	for (i = 0; i < perf_loops; i++) {
		for (j = 0; j < nnodes; j++)
			insert_augmented(nodes + j, &root);
		for (j = 0; j < nnodes; j++)
			erase_augmented(nodes + j, &root);
	}

	time2 = get_cycles();
	time = time2 - time1;

	time = div_u64(time, perf_loops);
	printk(" -> %llu cycles\n", (unsigned long long)time);

	for (i = 0; i < check_loops; i++) {
		init();
		for (j = 0; j < nnodes; j++) {
			check_augmented(j);
			insert_augmented(nodes + j, &root);
		}
		for (j = 0; j < nnodes; j++) {
			check_augmented(nnodes - j);
			erase_augmented(nodes + j, &root);
		}
		check_augmented(0);
	}

	kfree(nodes);

	return -EAGAIN; /* Fail will directly unload the module */
}

static void __exit rbtree_test_exit(void)
{
	printk(KERN_ALERT "test exit\n");
}

module_init(rbtree_test_init)
module_exit(rbtree_test_exit)

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Michel Lespinasse");
MODULE_DESCRIPTION("Red Black Tree test");
Commit	Line	Data
910a742d	1	#include <linux/module.h>
223f8911	2	#include <linux/moduleparam.h>
9c079add	3	#include <linux/rbtree_augmented.h>
910a742d	4	#include <linux/random.h>
223f8911	5	#include <linux/slab.h>
910a742d ML	6	#include <asm/timex.h>
910a742d ML	7
223f8911 DB	8	#define __param(type, name, init, msg) \
	9	static type name = init; \
	10	module_param(name, type, 0444); \
	11	MODULE_PARM_DESC(name, msg);
	12
	13	__param(int, nnodes, 100, "Number of nodes in the rb-tree");
	14	__param(int, perf_loops, 100000, "Number of iterations modifying the rb-tree");
	15	__param(int, check_loops, 100, "Number of iterations modifying and verifying the rb-tree");
910a742d ML	16
910a742d ML	17	struct test_node {
910a742d	18	u32 key;
dbf128cb	19	struct rb_node rb;
dadf9353 ML	20
	21	/* following fields used for testing augmented rbtree functionality */
	22	u32 val;
	23	u32 augmented;
910a742d ML	24	};
	25
	26	static struct rb_root root = RB_ROOT;
223f8911	27	static struct test_node *nodes = NULL;
910a742d ML	28
	29	static struct rnd_state rnd;
	30
	31	static void insert(struct test_node node, struct rb_root root)
	32	{
	33	struct rb_node *new = &root->rb_node, parent = NULL;
dadf9353	34	u32 key = node->key;
910a742d ML	35
	36	while (*new) {
	37	parent = *new;
dadf9353	38	if (key < rb_entry(parent, struct test_node, rb)->key)
910a742d ML	39	new = &parent->rb_left;
	40	else
	41	new = &parent->rb_right;
	42	}
	43
	44	rb_link_node(&node->rb, parent, new);
	45	rb_insert_color(&node->rb, root);
	46	}
	47
	48	static inline void erase(struct test_node node, struct rb_root root)
	49	{
	50	rb_erase(&node->rb, root);
	51	}
	52
dadf9353 ML	53	static inline u32 augment_recompute(struct test_node *node)
	54	{
	55	u32 max = node->val, child_augmented;
	56	if (node->rb.rb_left) {
	57	child_augmented = rb_entry(node->rb.rb_left, struct test_node,
	58	rb)->augmented;
	59	if (max < child_augmented)
	60	max = child_augmented;
	61	}
	62	if (node->rb.rb_right) {
	63	child_augmented = rb_entry(node->rb.rb_right, struct test_node,
	64	rb)->augmented;
	65	if (max < child_augmented)
	66	max = child_augmented;
	67	}
	68	return max;
	69	}
	70
3908836a ML	71	RB_DECLARE_CALLBACKS(static, augment_callbacks, struct test_node, rb,
3908836a ML	72	u32, augmented, augment_recompute)
14b94af0	73
dadf9353 ML	74	static void insert_augmented(struct test_node node, struct rb_root root)
dadf9353 ML	75	{
14b94af0	76	struct rb_node *new = &root->rb_node, rb_parent = NULL;
dadf9353	77	u32 key = node->key;
14b94af0 ML	78	u32 val = node->val;
14b94af0 ML	79	struct test_node *parent;
dadf9353 ML	80
dadf9353 ML	81	while (*new) {
14b94af0 ML	82	rb_parent = *new;
	83	parent = rb_entry(rb_parent, struct test_node, rb);
	84	if (parent->augmented < val)
	85	parent->augmented = val;
	86	if (key < parent->key)
	87	new = &parent->rb.rb_left;
dadf9353	88	else
14b94af0	89	new = &parent->rb.rb_right;
dadf9353 ML	90	}
dadf9353 ML	91
14b94af0 ML	92	node->augmented = val;
	93	rb_link_node(&node->rb, rb_parent, new);
	94	rb_insert_augmented(&node->rb, root, &augment_callbacks);
dadf9353 ML	95	}
	96
	97	static void erase_augmented(struct test_node node, struct rb_root root)
	98	{
14b94af0	99	rb_erase_augmented(&node->rb, root, &augment_callbacks);
dadf9353 ML	100	}
dadf9353 ML	101
910a742d ML	102	static void init(void)
	103	{
	104	int i;
223f8911	105	for (i = 0; i < nnodes; i++) {
496f2f93 AM	106	nodes[i].key = prandom_u32_state(&rnd);
496f2f93 AM	107	nodes[i].val = prandom_u32_state(&rnd);
dadf9353	108	}
910a742d ML	109	}
	110
	111	static bool is_red(struct rb_node *rb)
	112	{
	113	return !(rb->__rb_parent_color & 1);
	114	}
	115
	116	static int black_path_count(struct rb_node *rb)
	117	{
	118	int count;
	119	for (count = 0; rb; rb = rb_parent(rb))
	120	count += !is_red(rb);
	121	return count;
	122	}
	123
964fe94d CS	124	static void check_postorder_foreach(int nr_nodes)
	125	{
	126	struct test_node cur, n;
	127	int count = 0;
	128	rbtree_postorder_for_each_entry_safe(cur, n, &root, rb)
	129	count++;
	130
	131	WARN_ON_ONCE(count != nr_nodes);
	132	}
	133
a791a62f CS	134	static void check_postorder(int nr_nodes)
	135	{
	136	struct rb_node *rb;
	137	int count = 0;
	138	for (rb = rb_first_postorder(&root); rb; rb = rb_next_postorder(rb))
	139	count++;
	140
	141	WARN_ON_ONCE(count != nr_nodes);
	142	}
	143
910a742d ML	144	static void check(int nr_nodes)
	145	{
	146	struct rb_node *rb;
4130f0ef	147	int count = 0, blacks = 0;
910a742d ML	148	u32 prev_key = 0;
	149
	150	for (rb = rb_first(&root); rb; rb = rb_next(rb)) {
	151	struct test_node *node = rb_entry(rb, struct test_node, rb);
	152	WARN_ON_ONCE(node->key < prev_key);
	153	WARN_ON_ONCE(is_red(rb) &&
	154	(!rb_parent(rb) \|\| is_red(rb_parent(rb))));
	155	if (!count)
	156	blacks = black_path_count(rb);
	157	else
	158	WARN_ON_ONCE((!rb->rb_left \|\| !rb->rb_right) &&
	159	blacks != black_path_count(rb));
	160	prev_key = node->key;
	161	count++;
	162	}
4130f0ef	163
910a742d	164	WARN_ON_ONCE(count != nr_nodes);
4130f0ef	165	WARN_ON_ONCE(count < (1 << black_path_count(rb_last(&root))) - 1);
a791a62f CS	166
a791a62f CS	167	check_postorder(nr_nodes);
964fe94d	168	check_postorder_foreach(nr_nodes);
910a742d ML	169	}
910a742d ML	170
dadf9353 ML	171	static void check_augmented(int nr_nodes)
	172	{
	173	struct rb_node *rb;
	174
	175	check(nr_nodes);
	176	for (rb = rb_first(&root); rb; rb = rb_next(rb)) {
	177	struct test_node *node = rb_entry(rb, struct test_node, rb);
	178	WARN_ON_ONCE(node->augmented != augment_recompute(node));
	179	}
	180	}
	181
c75aaa8e	182	static int __init rbtree_test_init(void)
910a742d ML	183	{
	184	int i, j;
	185	cycles_t time1, time2, time;
	186
223f8911 DB	187	nodes = kmalloc(nnodes * sizeof(*nodes), GFP_KERNEL);
	188	if (!nodes)
	189	return -ENOMEM;
	190
910a742d ML	191	printk(KERN_ALERT "rbtree testing");
910a742d ML	192
496f2f93	193	prandom_seed_state(&rnd, 3141592653589793238ULL);
910a742d ML	194	init();
	195
	196	time1 = get_cycles();
	197
223f8911 DB	198	for (i = 0; i < perf_loops; i++) {
223f8911 DB	199	for (j = 0; j < nnodes; j++)
910a742d	200	insert(nodes + j, &root);
223f8911	201	for (j = 0; j < nnodes; j++)
910a742d ML	202	erase(nodes + j, &root);
	203	}
	204
	205	time2 = get_cycles();
	206	time = time2 - time1;
	207
223f8911	208	time = div_u64(time, perf_loops);
910a742d ML	209	printk(" -> %llu cycles\n", (unsigned long long)time);
910a742d ML	210
223f8911	211	for (i = 0; i < check_loops; i++) {
910a742d	212	init();
223f8911	213	for (j = 0; j < nnodes; j++) {
910a742d ML	214	check(j);
	215	insert(nodes + j, &root);
	216	}
223f8911 DB	217	for (j = 0; j < nnodes; j++) {
223f8911 DB	218	check(nnodes - j);
910a742d ML	219	erase(nodes + j, &root);
	220	}
	221	check(0);
	222	}
	223
dadf9353 ML	224	printk(KERN_ALERT "augmented rbtree testing");
	225
	226	init();
	227
	228	time1 = get_cycles();
	229
223f8911 DB	230	for (i = 0; i < perf_loops; i++) {
223f8911 DB	231	for (j = 0; j < nnodes; j++)
dadf9353	232	insert_augmented(nodes + j, &root);
223f8911	233	for (j = 0; j < nnodes; j++)
dadf9353 ML	234	erase_augmented(nodes + j, &root);
	235	}
	236
	237	time2 = get_cycles();
	238	time = time2 - time1;
	239
223f8911	240	time = div_u64(time, perf_loops);
dadf9353 ML	241	printk(" -> %llu cycles\n", (unsigned long long)time);
dadf9353 ML	242
223f8911	243	for (i = 0; i < check_loops; i++) {
dadf9353	244	init();
223f8911	245	for (j = 0; j < nnodes; j++) {
dadf9353 ML	246	check_augmented(j);
	247	insert_augmented(nodes + j, &root);
	248	}
223f8911 DB	249	for (j = 0; j < nnodes; j++) {
223f8911 DB	250	check_augmented(nnodes - j);
dadf9353 ML	251	erase_augmented(nodes + j, &root);
	252	}
	253	check_augmented(0);
	254	}
	255
223f8911 DB	256	kfree(nodes);
223f8911 DB	257
910a742d ML	258	return -EAGAIN; /* Fail will directly unload the module */
	259	}
	260
c75aaa8e	261	static void __exit rbtree_test_exit(void)
910a742d ML	262	{
	263	printk(KERN_ALERT "test exit\n");
	264	}
	265
	266	module_init(rbtree_test_init)
	267	module_exit(rbtree_test_exit)
	268
	269	MODULE_LICENSE("GPL");
	270	MODULE_AUTHOR("Michel Lespinasse");
	271	MODULE_DESCRIPTION("Red Black Tree test");