[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;
	struct rb_node *node;

	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(" -> test 1 (latency of nnodes insert+delete): %llu cycles\n", (unsigned long long)time);

	for (i = 0; i < nnodes; i++)
		insert(nodes + i, &root);

	time1 = get_cycles();

	for (i = 0; i < perf_loops; i++) {
		for (node = rb_first(&root); node; node = rb_next(node))
			;
	}

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

	time = div_u64(time, perf_loops);
	printk(" -> test 2 (latency of inorder traversal): %llu cycles\n", (unsigned long long)time);

	for (i = 0; i < nnodes; i++)
		erase(nodes + i, &root);

	/* run checks */
	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(" -> test 1 (latency of nnodes insert+delete): %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;
977bd8d5	186	struct rb_node *node;
910a742d	187
223f8911 DB	188	nodes = kmalloc(nnodes * sizeof(*nodes), GFP_KERNEL);
	189	if (!nodes)
	190	return -ENOMEM;
	191
910a742d ML	192	printk(KERN_ALERT "rbtree testing");
910a742d ML	193
496f2f93	194	prandom_seed_state(&rnd, 3141592653589793238ULL);
910a742d ML	195	init();
	196
	197	time1 = get_cycles();
	198
223f8911 DB	199	for (i = 0; i < perf_loops; i++) {
223f8911 DB	200	for (j = 0; j < nnodes; j++)
910a742d	201	insert(nodes + j, &root);
223f8911	202	for (j = 0; j < nnodes; j++)
910a742d ML	203	erase(nodes + j, &root);
	204	}
	205
	206	time2 = get_cycles();
	207	time = time2 - time1;
	208
223f8911	209	time = div_u64(time, perf_loops);
977bd8d5	210	printk(" -> test 1 (latency of nnodes insert+delete): %llu cycles\n", (unsigned long long)time);
910a742d	211
977bd8d5 DB	212	for (i = 0; i < nnodes; i++)
	213	insert(nodes + i, &root);
	214
	215	time1 = get_cycles();
	216
	217	for (i = 0; i < perf_loops; i++) {
	218	for (node = rb_first(&root); node; node = rb_next(node))
	219	;
	220	}
	221
	222	time2 = get_cycles();
	223	time = time2 - time1;
	224
	225	time = div_u64(time, perf_loops);
	226	printk(" -> test 2 (latency of inorder traversal): %llu cycles\n", (unsigned long long)time);
	227
	228	for (i = 0; i < nnodes; i++)
	229	erase(nodes + i, &root);
	230
	231	/* run checks */
223f8911	232	for (i = 0; i < check_loops; i++) {
910a742d	233	init();
223f8911	234	for (j = 0; j < nnodes; j++) {
910a742d ML	235	check(j);
	236	insert(nodes + j, &root);
	237	}
223f8911 DB	238	for (j = 0; j < nnodes; j++) {
223f8911 DB	239	check(nnodes - j);
910a742d ML	240	erase(nodes + j, &root);
	241	}
	242	check(0);
	243	}
	244
dadf9353 ML	245	printk(KERN_ALERT "augmented rbtree testing");
	246
	247	init();
	248
	249	time1 = get_cycles();
	250
223f8911 DB	251	for (i = 0; i < perf_loops; i++) {
223f8911 DB	252	for (j = 0; j < nnodes; j++)
dadf9353	253	insert_augmented(nodes + j, &root);
223f8911	254	for (j = 0; j < nnodes; j++)
dadf9353 ML	255	erase_augmented(nodes + j, &root);
	256	}
	257
	258	time2 = get_cycles();
	259	time = time2 - time1;
	260
223f8911	261	time = div_u64(time, perf_loops);
977bd8d5	262	printk(" -> test 1 (latency of nnodes insert+delete): %llu cycles\n", (unsigned long long)time);
dadf9353	263
223f8911	264	for (i = 0; i < check_loops; i++) {
dadf9353	265	init();
223f8911	266	for (j = 0; j < nnodes; j++) {
dadf9353 ML	267	check_augmented(j);
	268	insert_augmented(nodes + j, &root);
	269	}
223f8911 DB	270	for (j = 0; j < nnodes; j++) {
223f8911 DB	271	check_augmented(nnodes - j);
dadf9353 ML	272	erase_augmented(nodes + j, &root);
	273	}
	274	check_augmented(0);
	275	}
	276
223f8911 DB	277	kfree(nodes);
223f8911 DB	278
910a742d ML	279	return -EAGAIN; /* Fail will directly unload the module */
	280	}
	281
c75aaa8e	282	static void __exit rbtree_test_exit(void)
910a742d ML	283	{
	284	printk(KERN_ALERT "test exit\n");
	285	}
	286
	287	module_init(rbtree_test_init)
	288	module_exit(rbtree_test_exit)
	289
	290	MODULE_LICENSE("GPL");
	291	MODULE_AUTHOR("Michel Lespinasse");
	292	MODULE_DESCRIPTION("Red Black Tree test");