1 /*
2  * blkmap64_rb.c --- Simple rb-tree implementation for bitmaps
3  *
4  * (C)2010 Red Hat, Inc., Lukas Czerner <lczerner@redhat.com>
5  *
6  * %Begin-Header%
7  * This file may be redistributed under the terms of the GNU Public
8  * License.
9  * %End-Header%
10  */
11 
12 #include "config.h"
13 #include <stdio.h>
14 #include <string.h>
15 #if HAVE_UNISTD_H
16 #include <unistd.h>
17 #endif
18 #include <fcntl.h>
19 #include <time.h>
20 #if HAVE_SYS_STAT_H
21 #include <sys/stat.h>
22 #endif
23 #if HAVE_SYS_TYPES_H
24 #include <sys/types.h>
25 #endif
26 #if HAVE_LINUX_TYPES_H
27 #include <linux/types.h>
28 #endif
29 
30 #include "ext2_fs.h"
31 #include "ext2fsP.h"
32 #include "bmap64.h"
33 #include "rbtree.h"
34 
35 #include <limits.h>
36 
37 struct bmap_rb_extent {
38 	struct rb_node node;
39 	__u64 start;
40 	__u64 count;
41 };
42 
43 struct ext2fs_rb_private {
44 	struct rb_root root;
45 	struct bmap_rb_extent *wcursor;
46 	struct bmap_rb_extent *rcursor;
47 	struct bmap_rb_extent *rcursor_next;
48 #ifdef ENABLE_BMAP_STATS_OPS
49 	__u64 mark_hit;
50 	__u64 test_hit;
51 #endif
52 };
53 
node_to_extent(struct rb_node * node)54 inline static struct bmap_rb_extent *node_to_extent(struct rb_node *node)
55 {
56 	/*
57 	 * This depends on the fact the struct rb_node is at the
58 	 * beginning of the bmap_rb_extent structure.  We use this
59 	 * instead of the ext2fs_rb_entry macro because it causes gcc
60 	 * -Wall to generate a huge amount of noise.
61 	 */
62 	return (struct bmap_rb_extent *) node;
63 }
64 
65 static int rb_insert_extent(__u64 start, __u64 count,
66 			    struct ext2fs_rb_private *);
67 static void rb_get_new_extent(struct bmap_rb_extent **, __u64, __u64);
68 
69 /* #define DEBUG_RB */
70 
71 #ifdef DEBUG_RB
print_tree(struct rb_root * root)72 static void print_tree(struct rb_root *root)
73 {
74 	struct rb_node *node = NULL;
75 	struct bmap_rb_extent *ext;
76 
77 	printf("\t\t\t=================================\n");
78 	node = ext2fs_rb_first(root);
79 	for (node = ext2fs_rb_first(root); node != NULL;
80 	     node = ext2fs_rb_next(node)) {
81 		ext = node_to_extent(node);
82 		printf("\t\t\t--> (%llu -> %llu)\n",
83 			ext->start, ext->start + ext->count);
84 	}
85 	printf("\t\t\t=================================\n");
86 }
87 
check_tree(struct rb_root * root,const char * msg)88 static void check_tree(struct rb_root *root, const char *msg)
89 {
90 	struct rb_node *node;
91 	struct bmap_rb_extent *ext, *old = NULL;
92 
93 	for (node = ext2fs_rb_first(root); node;
94 	     node = ext2fs_rb_next(node)) {
95 		ext = node_to_extent(node);
96 		if (ext->count == 0) {
97 			printf("Tree Error: count is zero\n");
98 			printf("extent: %llu -> %llu (%llu)\n", ext->start,
99 				ext->start + ext->count, ext->count);
100 			goto err_out;
101 		}
102 		if (ext->start + ext->count < ext->start) {
103 			printf("Tree Error: start or count is crazy\n");
104 			printf("extent: %llu -> %llu (%llu)\n", ext->start,
105 				ext->start + ext->count, ext->count);
106 			goto err_out;
107 		}
108 
109 		if (old) {
110 			if (old->start > ext->start) {
111 				printf("Tree Error: start is crazy\n");
112 				printf("extent: %llu -> %llu (%llu)\n",
113 					old->start, old->start + old->count,
114 					old->count);
115 				printf("extent next: %llu -> %llu (%llu)\n",
116 					ext->start, ext->start + ext->count,
117 					ext->count);
118 				goto err_out;
119 			}
120 			if ((old->start + old->count) >= ext->start) {
121 				printf("Tree Error: extent is crazy\n");
122 				printf("extent: %llu -> %llu (%llu)\n",
123 					old->start, old->start + old->count,
124 					old->count);
125 				printf("extent next: %llu -> %llu (%llu)\n",
126 					ext->start, ext->start + ext->count,
127 					ext->count);
128 				goto err_out;
129 			}
130 		}
131 		old = ext;
132 	}
133 	return;
134 
135 err_out:
136 	printf("%s\n", msg);
137 	print_tree(root);
138 	exit(1);
139 }
140 #else
141 #define check_tree(root, msg) do {} while (0)
142 #define print_tree(root) do {} while (0)
143 #endif
144 
rb_get_new_extent(struct bmap_rb_extent ** ext,__u64 start,__u64 count)145 static void rb_get_new_extent(struct bmap_rb_extent **ext, __u64 start,
146 			     __u64 count)
147 {
148 	struct bmap_rb_extent *new_ext;
149 	int retval;
150 
151 	retval = ext2fs_get_mem(sizeof (struct bmap_rb_extent),
152 				&new_ext);
153 	if (retval)
154 		abort();
155 
156 	new_ext->start = start;
157 	new_ext->count = count;
158 	*ext = new_ext;
159 }
160 
161 inline
rb_free_extent(struct ext2fs_rb_private * bp,struct bmap_rb_extent * ext)162 static void rb_free_extent(struct ext2fs_rb_private *bp,
163 			   struct bmap_rb_extent *ext)
164 {
165 	if (bp->wcursor == ext)
166 		bp->wcursor = NULL;
167 	if (bp->rcursor == ext)
168 		bp->rcursor = NULL;
169 	if (bp->rcursor_next == ext)
170 		bp->rcursor_next = NULL;
171 	ext2fs_free_mem(&ext);
172 }
173 
rb_alloc_private_data(ext2fs_generic_bitmap bitmap)174 static errcode_t rb_alloc_private_data (ext2fs_generic_bitmap bitmap)
175 {
176 	struct ext2fs_rb_private *bp;
177 	errcode_t	retval;
178 
179 	retval = ext2fs_get_mem(sizeof (struct ext2fs_rb_private), &bp);
180 	if (retval)
181 		return retval;
182 
183 	bp->root = RB_ROOT;
184 	bp->rcursor = NULL;
185 	bp->rcursor_next = NULL;
186 	bp->wcursor = NULL;
187 
188 #ifdef ENABLE_BMAP_STATS_OPS
189 	bp->test_hit = 0;
190 	bp->mark_hit = 0;
191 #endif
192 
193 	bitmap->private = (void *) bp;
194 	return 0;
195 }
196 
rb_new_bmap(ext2_filsys fs EXT2FS_ATTR ((unused)),ext2fs_generic_bitmap bitmap)197 static errcode_t rb_new_bmap(ext2_filsys fs EXT2FS_ATTR((unused)),
198 			     ext2fs_generic_bitmap bitmap)
199 {
200 	errcode_t	retval;
201 
202 	retval = rb_alloc_private_data (bitmap);
203 	if (retval)
204 		return retval;
205 
206 	return 0;
207 }
208 
rb_free_tree(struct rb_root * root)209 static void rb_free_tree(struct rb_root *root)
210 {
211 	struct bmap_rb_extent *ext;
212 	struct rb_node *node, *next;
213 
214 	for (node = ext2fs_rb_first(root); node; node = next) {
215 		next = ext2fs_rb_next(node);
216 		ext = node_to_extent(node);
217 		ext2fs_rb_erase(node, root);
218 		ext2fs_free_mem(&ext);
219 	}
220 }
221 
rb_free_bmap(ext2fs_generic_bitmap bitmap)222 static void rb_free_bmap(ext2fs_generic_bitmap bitmap)
223 {
224 	struct ext2fs_rb_private *bp;
225 
226 	bp = (struct ext2fs_rb_private *) bitmap->private;
227 
228 	rb_free_tree(&bp->root);
229 	ext2fs_free_mem(&bp);
230 	bp = 0;
231 }
232 
rb_copy_bmap(ext2fs_generic_bitmap src,ext2fs_generic_bitmap dest)233 static errcode_t rb_copy_bmap(ext2fs_generic_bitmap src,
234 			      ext2fs_generic_bitmap dest)
235 {
236 	struct ext2fs_rb_private *src_bp, *dest_bp;
237 	struct bmap_rb_extent *src_ext, *dest_ext;
238 	struct rb_node *dest_node, *src_node, *dest_last, **n;
239 	errcode_t retval = 0;
240 
241 	retval = rb_alloc_private_data (dest);
242 	if (retval)
243 		return retval;
244 
245 	src_bp = (struct ext2fs_rb_private *) src->private;
246 	dest_bp = (struct ext2fs_rb_private *) dest->private;
247 	src_bp->rcursor = NULL;
248 	dest_bp->rcursor = NULL;
249 
250 	src_node = ext2fs_rb_first(&src_bp->root);
251 	while (src_node) {
252 		src_ext = node_to_extent(src_node);
253 		retval = ext2fs_get_mem(sizeof (struct bmap_rb_extent),
254 					&dest_ext);
255 		if (retval)
256 			break;
257 
258 		memcpy(dest_ext, src_ext, sizeof(struct bmap_rb_extent));
259 
260 		dest_node = &dest_ext->node;
261 		n = &dest_bp->root.rb_node;
262 
263 		dest_last = NULL;
264 		if (*n) {
265 			dest_last = ext2fs_rb_last(&dest_bp->root);
266 			n = &(dest_last)->rb_right;
267 		}
268 
269 		ext2fs_rb_link_node(dest_node, dest_last, n);
270 		ext2fs_rb_insert_color(dest_node, &dest_bp->root);
271 
272 		src_node = ext2fs_rb_next(src_node);
273 	}
274 
275 	return retval;
276 }
277 
rb_truncate(__u64 new_max,struct rb_root * root)278 static void rb_truncate(__u64 new_max, struct rb_root *root)
279 {
280 	struct bmap_rb_extent *ext;
281 	struct rb_node *node;
282 
283 	node = ext2fs_rb_last(root);
284 	while (node) {
285 		ext = node_to_extent(node);
286 
287 		if ((ext->start + ext->count - 1) <= new_max)
288 			break;
289 		else if (ext->start > new_max) {
290 			ext2fs_rb_erase(node, root);
291 			ext2fs_free_mem(&ext);
292 			node = ext2fs_rb_last(root);
293 			continue;
294 		} else
295 			ext->count = new_max - ext->start + 1;
296 	}
297 }
298 
rb_resize_bmap(ext2fs_generic_bitmap bmap,__u64 new_end,__u64 new_real_end)299 static errcode_t rb_resize_bmap(ext2fs_generic_bitmap bmap,
300 				__u64 new_end, __u64 new_real_end)
301 {
302 	struct ext2fs_rb_private *bp;
303 
304 	bp = (struct ext2fs_rb_private *) bmap->private;
305 	bp->rcursor = NULL;
306 	bp->wcursor = NULL;
307 
308 	rb_truncate(((new_end < bmap->end) ? new_end : bmap->end) - bmap->start,
309 		    &bp->root);
310 
311 	bmap->end = new_end;
312 	bmap->real_end = new_real_end;
313 
314 	if (bmap->end < bmap->real_end)
315 		rb_insert_extent(bmap->end + 1 - bmap->start,
316 				 bmap->real_end - bmap->end, bp);
317 	return 0;
318 
319 }
320 
321 inline static int
rb_test_bit(struct ext2fs_rb_private * bp,__u64 bit)322 rb_test_bit(struct ext2fs_rb_private *bp, __u64 bit)
323 {
324 	struct bmap_rb_extent *rcursor, *next_ext = NULL;
325 	struct rb_node *parent = NULL, *next;
326 	struct rb_node **n = &bp->root.rb_node;
327 	struct bmap_rb_extent *ext;
328 
329 	rcursor = bp->rcursor;
330 	if (!rcursor)
331 		goto search_tree;
332 
333 	if (bit >= rcursor->start && bit < rcursor->start + rcursor->count) {
334 #ifdef ENABLE_BMAP_STATS_OPS
335 		bp->test_hit++;
336 #endif
337 		return 1;
338 	}
339 
340 	next_ext = bp->rcursor_next;
341 	if (!next_ext) {
342 		next = ext2fs_rb_next(&rcursor->node);
343 		if (next)
344 			next_ext = node_to_extent(next);
345 		bp->rcursor_next = next_ext;
346 	}
347 	if (next_ext) {
348 		if ((bit >= rcursor->start + rcursor->count) &&
349 		    (bit < next_ext->start)) {
350 #ifdef BMAP_STATS_OPS
351 			bp->test_hit++;
352 #endif
353 			return 0;
354 		}
355 	}
356 	bp->rcursor = NULL;
357 	bp->rcursor_next = NULL;
358 
359 	rcursor = bp->wcursor;
360 	if (!rcursor)
361 		goto search_tree;
362 
363 	if (bit >= rcursor->start && bit < rcursor->start + rcursor->count)
364 		return 1;
365 
366 search_tree:
367 
368 	while (*n) {
369 		parent = *n;
370 		ext = node_to_extent(parent);
371 		if (bit < ext->start)
372 			n = &(*n)->rb_left;
373 		else if (bit >= (ext->start + ext->count))
374 			n = &(*n)->rb_right;
375 		else {
376 			bp->rcursor = ext;
377 			bp->rcursor_next = NULL;
378 			return 1;
379 		}
380 	}
381 	return 0;
382 }
383 
rb_insert_extent(__u64 start,__u64 count,struct ext2fs_rb_private * bp)384 static int rb_insert_extent(__u64 start, __u64 count,
385 			    struct ext2fs_rb_private *bp)
386 {
387 	struct rb_root *root = &bp->root;
388 	struct rb_node *parent = NULL, **n = &root->rb_node;
389 	struct rb_node *new_node, *node, *next;
390 	struct bmap_rb_extent *new_ext;
391 	struct bmap_rb_extent *ext;
392 	int retval = 0;
393 
394 	bp->rcursor_next = NULL;
395 	ext = bp->wcursor;
396 	if (ext) {
397 		if (start >= ext->start &&
398 		    start <= (ext->start + ext->count)) {
399 #ifdef ENABLE_BMAP_STATS_OPS
400 			bp->mark_hit++;
401 #endif
402 			goto got_extent;
403 		}
404 	}
405 
406 	while (*n) {
407 		parent = *n;
408 		ext = node_to_extent(parent);
409 
410 		if (start < ext->start) {
411 			n = &(*n)->rb_left;
412 		} else if (start > (ext->start + ext->count)) {
413 			n = &(*n)->rb_right;
414 		} else {
415 got_extent:
416 			if ((start + count) <= (ext->start + ext->count))
417 				return 1;
418 
419 			if ((ext->start + ext->count) == start)
420 				retval = 0;
421 			else
422 				retval = 1;
423 
424 			count += (start - ext->start);
425 			start = ext->start;
426 			new_ext = ext;
427 			new_node = &ext->node;
428 
429 			goto skip_insert;
430 		}
431 	}
432 
433 	rb_get_new_extent(&new_ext, start, count);
434 
435 	new_node = &new_ext->node;
436 	ext2fs_rb_link_node(new_node, parent, n);
437 	ext2fs_rb_insert_color(new_node, root);
438 	bp->wcursor = new_ext;
439 
440 	node = ext2fs_rb_prev(new_node);
441 	if (node) {
442 		ext = node_to_extent(node);
443 		if ((ext->start + ext->count) == start) {
444 			start = ext->start;
445 			count += ext->count;
446 			ext2fs_rb_erase(node, root);
447 			rb_free_extent(bp, ext);
448 		}
449 	}
450 
451 skip_insert:
452 	/* See if we can merge extent to the right */
453 	for (node = ext2fs_rb_next(new_node); node != NULL; node = next) {
454 		next = ext2fs_rb_next(node);
455 		ext = node_to_extent(node);
456 
457 		if ((ext->start + ext->count) <= start)
458 			continue;
459 
460 		/* No more merging */
461 		if ((start + count) < ext->start)
462 			break;
463 
464 		/* ext is embedded in new_ext interval */
465 		if ((start + count) >= (ext->start + ext->count)) {
466 			ext2fs_rb_erase(node, root);
467 			rb_free_extent(bp, ext);
468 			continue;
469 		} else {
470 		/* merge ext with new_ext */
471 			count += ((ext->start + ext->count) -
472 				  (start + count));
473 			ext2fs_rb_erase(node, root);
474 			rb_free_extent(bp, ext);
475 			break;
476 		}
477 	}
478 
479 	new_ext->start = start;
480 	new_ext->count = count;
481 
482 	return retval;
483 }
484 
rb_remove_extent(__u64 start,__u64 count,struct ext2fs_rb_private * bp)485 static int rb_remove_extent(__u64 start, __u64 count,
486 			    struct ext2fs_rb_private *bp)
487 {
488 	struct rb_root *root = &bp->root;
489 	struct rb_node *parent = NULL, **n = &root->rb_node;
490 	struct rb_node *node;
491 	struct bmap_rb_extent *ext;
492 	__u64 new_start, new_count;
493 	int retval = 0;
494 
495 	if (ext2fs_rb_empty_root(root))
496 		return 0;
497 
498 	while (*n) {
499 		parent = *n;
500 		ext = node_to_extent(parent);
501 		if (start < ext->start) {
502 			n = &(*n)->rb_left;
503 			continue;
504 		} else if (start >= (ext->start + ext->count)) {
505 			n = &(*n)->rb_right;
506 			continue;
507 		}
508 
509 		if ((start > ext->start) &&
510 		    (start + count) < (ext->start + ext->count)) {
511 			/* We have to split extent into two */
512 			new_start = start + count;
513 			new_count = (ext->start + ext->count) - new_start;
514 
515 			ext->count = start - ext->start;
516 
517 			rb_insert_extent(new_start, new_count, bp);
518 			return 1;
519 		}
520 
521 		if ((start + count) >= (ext->start + ext->count)) {
522 			ext->count = start - ext->start;
523 			retval = 1;
524 		}
525 
526 		if (0 == ext->count) {
527 			parent = ext2fs_rb_next(&ext->node);
528 			ext2fs_rb_erase(&ext->node, root);
529 			rb_free_extent(bp, ext);
530 			break;
531 		}
532 
533 		if (start == ext->start) {
534 			ext->start += count;
535 			ext->count -= count;
536 			return 1;
537 		}
538 	}
539 
540 	/* See if we should delete or truncate extent on the right */
541 	for (; parent != NULL; parent = node) {
542 		node = ext2fs_rb_next(parent);
543 		ext = node_to_extent(parent);
544 		if ((ext->start + ext->count) <= start)
545 			continue;
546 
547 		/* No more extents to be removed/truncated */
548 		if ((start + count) < ext->start)
549 			break;
550 
551 		/* The entire extent is within the region to be removed */
552 		if ((start + count) >= (ext->start + ext->count)) {
553 			ext2fs_rb_erase(parent, root);
554 			rb_free_extent(bp, ext);
555 			retval = 1;
556 			continue;
557 		} else {
558 			/* modify the last extent in reigon to be removed */
559 			ext->count -= ((start + count) - ext->start);
560 			ext->start = start + count;
561 			retval = 1;
562 			break;
563 		}
564 	}
565 
566 	return retval;
567 }
568 
rb_mark_bmap(ext2fs_generic_bitmap bitmap,__u64 arg)569 static int rb_mark_bmap(ext2fs_generic_bitmap bitmap, __u64 arg)
570 {
571 	struct ext2fs_rb_private *bp;
572 	int retval;
573 
574 	bp = (struct ext2fs_rb_private *) bitmap->private;
575 	arg -= bitmap->start;
576 
577 	retval = rb_insert_extent(arg, 1, bp);
578 	check_tree(&bp->root, __func__);
579 	return retval;
580 }
581 
rb_unmark_bmap(ext2fs_generic_bitmap bitmap,__u64 arg)582 static int rb_unmark_bmap(ext2fs_generic_bitmap bitmap, __u64 arg)
583 {
584 	struct ext2fs_rb_private *bp;
585 	int retval;
586 
587 	bp = (struct ext2fs_rb_private *) bitmap->private;
588 	arg -= bitmap->start;
589 
590 	retval = rb_remove_extent(arg, 1, bp);
591 	check_tree(&bp->root, __func__);
592 
593 	return retval;
594 }
595 
596 inline
rb_test_bmap(ext2fs_generic_bitmap bitmap,__u64 arg)597 static int rb_test_bmap(ext2fs_generic_bitmap bitmap, __u64 arg)
598 {
599 	struct ext2fs_rb_private *bp;
600 
601 	bp = (struct ext2fs_rb_private *) bitmap->private;
602 	arg -= bitmap->start;
603 
604 	return rb_test_bit(bp, arg);
605 }
606 
rb_mark_bmap_extent(ext2fs_generic_bitmap bitmap,__u64 arg,unsigned int num)607 static void rb_mark_bmap_extent(ext2fs_generic_bitmap bitmap, __u64 arg,
608 				unsigned int num)
609 {
610 	struct ext2fs_rb_private *bp;
611 
612 	bp = (struct ext2fs_rb_private *) bitmap->private;
613 	arg -= bitmap->start;
614 
615 	rb_insert_extent(arg, num, bp);
616 	check_tree(&bp->root, __func__);
617 }
618 
rb_unmark_bmap_extent(ext2fs_generic_bitmap bitmap,__u64 arg,unsigned int num)619 static void rb_unmark_bmap_extent(ext2fs_generic_bitmap bitmap, __u64 arg,
620 				  unsigned int num)
621 {
622 	struct ext2fs_rb_private *bp;
623 
624 	bp = (struct ext2fs_rb_private *) bitmap->private;
625 	arg -= bitmap->start;
626 
627 	rb_remove_extent(arg, num, bp);
628 	check_tree(&bp->root, __func__);
629 }
630 
rb_test_clear_bmap_extent(ext2fs_generic_bitmap bitmap,__u64 start,unsigned int len)631 static int rb_test_clear_bmap_extent(ext2fs_generic_bitmap bitmap,
632 				     __u64 start, unsigned int len)
633 {
634 	struct rb_node *parent = NULL, **n;
635 	struct rb_node *node, *next;
636 	struct ext2fs_rb_private *bp;
637 	struct bmap_rb_extent *ext;
638 	int retval = 1;
639 
640 	bp = (struct ext2fs_rb_private *) bitmap->private;
641 	n = &bp->root.rb_node;
642 	start -= bitmap->start;
643 
644 	if (len == 0 || ext2fs_rb_empty_root(&bp->root))
645 		return 1;
646 
647 	/*
648 	 * If we find nothing, we should examine whole extent, but
649 	 * when we find match, the extent is not clean, thus be return
650 	 * false.
651 	 */
652 	while (*n) {
653 		parent = *n;
654 		ext = node_to_extent(parent);
655 		if (start < ext->start) {
656 			n = &(*n)->rb_left;
657 		} else if (start >= (ext->start + ext->count)) {
658 			n = &(*n)->rb_right;
659 		} else {
660 			/*
661 			 * We found extent int the tree -> extent is not
662 			 * clean
663 			 */
664 			return 0;
665 		}
666 	}
667 
668 	node = parent;
669 	while (node) {
670 		next = ext2fs_rb_next(node);
671 		ext = node_to_extent(node);
672 		node = next;
673 
674 		if ((ext->start + ext->count) <= start)
675 			continue;
676 
677 		/* No more merging */
678 		if ((start + len) <= ext->start)
679 			break;
680 
681 		retval = 0;
682 		break;
683 	}
684 	return retval;
685 }
686 
rb_set_bmap_range(ext2fs_generic_bitmap bitmap,__u64 start,size_t num,void * in)687 static errcode_t rb_set_bmap_range(ext2fs_generic_bitmap bitmap,
688 				     __u64 start, size_t num, void *in)
689 {
690 	struct ext2fs_rb_private *bp;
691 	unsigned char *cp = in;
692 	size_t i;
693 	int first_set = -1;
694 
695 	bp = (struct ext2fs_rb_private *) bitmap->private;
696 
697 	for (i = 0; i < num; i++) {
698 		if ((i & 7) == 0) {
699 			unsigned char c = cp[i/8];
700 			if (c == 0xFF) {
701 				if (first_set == -1)
702 					first_set = i;
703 				i += 7;
704 				continue;
705 			}
706 			if ((c == 0x00) && (first_set == -1)) {
707 				i += 7;
708 				continue;
709 			}
710 		}
711 		if (ext2fs_test_bit(i, in)) {
712 			if (first_set == -1)
713 				first_set = i;
714 			continue;
715 		}
716 		if (first_set == -1)
717 			continue;
718 
719 		rb_insert_extent(start + first_set - bitmap->start,
720 				 i - first_set, bp);
721 		check_tree(&bp->root, __func__);
722 		first_set = -1;
723 	}
724 	if (first_set != -1) {
725 		rb_insert_extent(start + first_set - bitmap->start,
726 				 num - first_set, bp);
727 		check_tree(&bp->root, __func__);
728 	}
729 
730 	return 0;
731 }
732 
rb_get_bmap_range(ext2fs_generic_bitmap bitmap,__u64 start,size_t num,void * out)733 static errcode_t rb_get_bmap_range(ext2fs_generic_bitmap bitmap,
734 				     __u64 start, size_t num, void *out)
735 {
736 
737 	struct rb_node *parent = NULL, *next, **n;
738 	struct ext2fs_rb_private *bp;
739 	struct bmap_rb_extent *ext;
740 	__u64 count, pos;
741 
742 	bp = (struct ext2fs_rb_private *) bitmap->private;
743 	n = &bp->root.rb_node;
744 	start -= bitmap->start;
745 
746 	if (ext2fs_rb_empty_root(&bp->root))
747 		return 0;
748 
749 	while (*n) {
750 		parent = *n;
751 		ext = node_to_extent(parent);
752 		if (start < ext->start) {
753 			n = &(*n)->rb_left;
754 		} else if (start >= (ext->start + ext->count)) {
755 			n = &(*n)->rb_right;
756 		} else
757 			break;
758 	}
759 
760 	memset(out, 0, (num + 7) >> 3);
761 
762 	for (; parent != NULL; parent = next) {
763 		next = ext2fs_rb_next(parent);
764 		ext = node_to_extent(parent);
765 
766 		pos = ext->start;
767 		count = ext->count;
768 		if (pos >= start + num)
769 			break;
770 		if (pos < start) {
771 			if (pos + count <  start)
772 				continue;
773 			count -= start - pos;
774 			pos = start;
775 		}
776 		if (pos + count > start + num)
777 			count = start + num - pos;
778 
779 		while (count > 0) {
780 			if ((count >= 8) &&
781 			    ((pos - start) % 8) == 0) {
782 				int nbytes = count >> 3;
783 				int offset = (pos - start) >> 3;
784 
785 				memset(((char *) out) + offset, 0xFF, nbytes);
786 				pos += nbytes << 3;
787 				count -= nbytes << 3;
788 				continue;
789 			}
790 			ext2fs_fast_set_bit64((pos - start), out);
791 			pos++;
792 			count--;
793 		}
794 	}
795 	return 0;
796 }
797 
rb_clear_bmap(ext2fs_generic_bitmap bitmap)798 static void rb_clear_bmap(ext2fs_generic_bitmap bitmap)
799 {
800 	struct ext2fs_rb_private *bp;
801 
802 	bp = (struct ext2fs_rb_private *) bitmap->private;
803 
804 	rb_free_tree(&bp->root);
805 	bp->rcursor = NULL;
806 	bp->rcursor_next = NULL;
807 	bp->wcursor = NULL;
808 	check_tree(&bp->root, __func__);
809 }
810 
rb_find_first_zero(ext2fs_generic_bitmap bitmap,__u64 start,__u64 end,__u64 * out)811 static errcode_t rb_find_first_zero(ext2fs_generic_bitmap bitmap,
812 				   __u64 start, __u64 end, __u64 *out)
813 {
814 	struct rb_node *parent = NULL, **n;
815 	struct ext2fs_rb_private *bp;
816 	struct bmap_rb_extent *ext;
817 
818 	bp = (struct ext2fs_rb_private *) bitmap->private;
819 	n = &bp->root.rb_node;
820 	start -= bitmap->start;
821 	end -= bitmap->start;
822 
823 	if (start > end)
824 		return EINVAL;
825 
826 	if (ext2fs_rb_empty_root(&bp->root))
827 		return ENOENT;
828 
829 	while (*n) {
830 		parent = *n;
831 		ext = node_to_extent(parent);
832 		if (start < ext->start) {
833 			n = &(*n)->rb_left;
834 		} else if (start >= (ext->start + ext->count)) {
835 			n = &(*n)->rb_right;
836 		} else if (ext->start + ext->count <= end) {
837 			*out = ext->start + ext->count + bitmap->start;
838 			return 0;
839 		} else
840 			return ENOENT;
841 	}
842 
843 	*out = start + bitmap->start;
844 	return 0;
845 }
846 
rb_find_first_set(ext2fs_generic_bitmap bitmap,__u64 start,__u64 end,__u64 * out)847 static errcode_t rb_find_first_set(ext2fs_generic_bitmap bitmap,
848 				   __u64 start, __u64 end, __u64 *out)
849 {
850 	struct rb_node *parent = NULL, **n;
851 	struct rb_node *node;
852 	struct ext2fs_rb_private *bp;
853 	struct bmap_rb_extent *ext;
854 
855 	bp = (struct ext2fs_rb_private *) bitmap->private;
856 	n = &bp->root.rb_node;
857 	start -= bitmap->start;
858 	end -= bitmap->start;
859 
860 	if (start > end)
861 		return EINVAL;
862 
863 	if (ext2fs_rb_empty_root(&bp->root))
864 		return ENOENT;
865 
866 	while (*n) {
867 		parent = *n;
868 		ext = node_to_extent(parent);
869 		if (start < ext->start) {
870 			n = &(*n)->rb_left;
871 		} else if (start >= (ext->start + ext->count)) {
872 			n = &(*n)->rb_right;
873 		} else {
874 			/* The start bit is set */
875 			*out = start + bitmap->start;
876 			return 0;
877 		}
878 	}
879 
880 	node = parent;
881 	ext = node_to_extent(node);
882 	if (ext->start < start) {
883 		node = ext2fs_rb_next(node);
884 		if (node == NULL)
885 			return ENOENT;
886 		ext = node_to_extent(node);
887 	}
888 	if (ext->start <= end) {
889 		*out = ext->start + bitmap->start;
890 		return 0;
891 	}
892 	return ENOENT;
893 }
894 
895 #ifdef ENABLE_BMAP_STATS
rb_print_stats(ext2fs_generic_bitmap bitmap)896 static void rb_print_stats(ext2fs_generic_bitmap bitmap)
897 {
898 	struct ext2fs_rb_private *bp;
899 	struct rb_node *node = NULL;
900 	struct bmap_rb_extent *ext;
901 	__u64 count = 0;
902 	__u64 max_size = 0;
903 	__u64 min_size = ULONG_MAX;
904 	__u64 size = 0, avg_size = 0;
905 	double eff;
906 #ifdef ENABLE_BMAP_STATS_OPS
907 	__u64 mark_all, test_all;
908 	double m_hit = 0.0, t_hit = 0.0;
909 #endif
910 
911 	bp = (struct ext2fs_rb_private *) bitmap->private;
912 
913 	for (node = ext2fs_rb_first(&bp->root); node != NULL;
914 	     node = ext2fs_rb_next(node)) {
915 		ext = node_to_extent(node);
916 		count++;
917 		if (ext->count > max_size)
918 			max_size = ext->count;
919 		if (ext->count < min_size)
920 			min_size = ext->count;
921 		size += ext->count;
922 	}
923 
924 	if (count)
925 		avg_size = size / count;
926 	if (min_size == ULONG_MAX)
927 		min_size = 0;
928 	eff = (double)((count * sizeof(struct bmap_rb_extent)) << 3) /
929 	      (bitmap->real_end - bitmap->start);
930 #ifdef ENABLE_BMAP_STATS_OPS
931 	mark_all = bitmap->stats.mark_count + bitmap->stats.mark_ext_count;
932 	test_all = bitmap->stats.test_count + bitmap->stats.test_ext_count;
933 	if (mark_all)
934 		m_hit = ((double)bp->mark_hit / mark_all) * 100;
935 	if (test_all)
936 		t_hit = ((double)bp->test_hit / test_all) * 100;
937 
938 	fprintf(stderr, "%16llu cache hits on test (%.2f%%)\n"
939 		"%16llu cache hits on mark (%.2f%%)\n",
940 		bp->test_hit, t_hit, bp->mark_hit, m_hit);
941 #endif
942 	fprintf(stderr, "%16llu extents (%llu bytes)\n",
943 		count, ((count * sizeof(struct bmap_rb_extent)) +
944 			sizeof(struct ext2fs_rb_private)));
945  	fprintf(stderr, "%16llu bits minimum size\n",
946 		min_size);
947 	fprintf(stderr, "%16llu bits maximum size\n"
948 		"%16llu bits average size\n",
949 		max_size, avg_size);
950 	fprintf(stderr, "%16llu bits set in bitmap (out of %llu)\n", size,
951 		bitmap->real_end - bitmap->start);
952 	fprintf(stderr,
953 		"%16.4lf memory / bitmap bit memory ratio (bitarray = 1)\n",
954 		eff);
955 }
956 #else
rb_print_stats(ext2fs_generic_bitmap bitmap EXT2FS_ATTR ((unused)))957 static void rb_print_stats(ext2fs_generic_bitmap bitmap EXT2FS_ATTR((unused)))
958 {
959 }
960 #endif
961 
962 struct ext2_bitmap_ops ext2fs_blkmap64_rbtree = {
963 	.type = EXT2FS_BMAP64_RBTREE,
964 	.new_bmap = rb_new_bmap,
965 	.free_bmap = rb_free_bmap,
966 	.copy_bmap = rb_copy_bmap,
967 	.resize_bmap = rb_resize_bmap,
968 	.mark_bmap = rb_mark_bmap,
969 	.unmark_bmap = rb_unmark_bmap,
970 	.test_bmap = rb_test_bmap,
971 	.test_clear_bmap_extent = rb_test_clear_bmap_extent,
972 	.mark_bmap_extent = rb_mark_bmap_extent,
973 	.unmark_bmap_extent = rb_unmark_bmap_extent,
974 	.set_bmap_range = rb_set_bmap_range,
975 	.get_bmap_range = rb_get_bmap_range,
976 	.clear_bmap = rb_clear_bmap,
977 	.print_stats = rb_print_stats,
978 	.find_first_zero = rb_find_first_zero,
979 	.find_first_set = rb_find_first_set,
980 };
981