1 files changed, 76 insertions, 74 deletions

diff --git a/extent-cache.c b/extent-cache.c
index 3dd6434e..a09fe870 100644
--- a/extent-cache.c
+++ b/extent-cache.c
@@ -20,65 +20,42 @@
 #include "kerncompat.h"
 #include "extent-cache.h"
 
+struct cache_extent_search_range {
+	u64 start;
+	u64 size;
+};
+
 void cache_tree_init(struct cache_tree *tree)
 {
-	tree->root.rb_node = NULL;
+	tree->root = RB_ROOT;
 }
 
-static struct rb_node *tree_insert(struct rb_root *root, u64 offset,
-				   u64 size, struct rb_node *node)
+static int cache_tree_comp_range(struct rb_node *node, void *data)
 {
-	struct rb_node ** p = &root->rb_node;
-	struct rb_node * parent = NULL;
 	struct cache_extent *entry;
+	struct cache_extent_search_range *range;
 
-	while(*p) {
-		parent = *p;
-		entry = rb_entry(parent, struct cache_extent, rb_node);
-
-		if (offset + size <= entry->start)
-			p = &(*p)->rb_left;
-		else if (offset >= entry->start + entry->size)
-			p = &(*p)->rb_right;
-		else
-			return parent;
-	}
+	range = (struct cache_extent_search_range *)data;
+	entry = rb_entry(node, struct cache_extent, rb_node);
 
-	entry = rb_entry(parent, struct cache_extent, rb_node);
-	rb_link_node(node, parent, p);
-	rb_insert_color(node, root);
-	return NULL;
+	if (entry->start + entry->size <= range->start)
+		return 1;
+	else if (range->start + range->size <= entry->start)
+		return -1;
+	else
+		return 0;
 }
 
-static struct rb_node *__tree_search(struct rb_root *root, u64 offset,
-				     u64 size, struct rb_node **prev_ret)
+static int cache_tree_comp_nodes(struct rb_node *node1, struct rb_node *node2)
 {
-	struct rb_node * n = root->rb_node;
-	struct rb_node *prev = NULL;
 	struct cache_extent *entry;
-	struct cache_extent *prev_entry = NULL;
-
-	while(n) {
-		entry = rb_entry(n, struct cache_extent, rb_node);
-		prev = n;
-		prev_entry = entry;
-
-		if (offset + size <= entry->start)
-			n = n->rb_left;
-		else if (offset >= entry->start + entry->size)
-			n = n->rb_right;
-		else
-			return n;
-	}
-	if (!prev_ret)
-		return NULL;
+	struct cache_extent_search_range range;
 
-	while(prev && offset >= prev_entry->start + prev_entry->size) {
-		prev = rb_next(prev);
-		prev_entry = rb_entry(prev, struct cache_extent, rb_node);
-	}
-	*prev_ret = prev;
-	return NULL;
+	entry = rb_entry(node2, struct cache_extent, rb_node);
+	range.start = entry->start;
+	range.size = entry->size;
+
+	return cache_tree_comp_range(node1, (void *)&range);
 }
 
 struct cache_extent *alloc_cache_extent(u64 start, u64 size)
@@ -87,63 +64,79 @@ struct cache_extent *alloc_cache_extent(u64 start, u64 size)
 
 	if (!pe)
 		return pe;
+
 	pe->start = start;
 	pe->size = size;
 	return pe;
 }
 
-int insert_existing_cache_extent(struct cache_tree *tree,
-				 struct cache_extent *pe)
+int insert_cache_extent(struct cache_tree *tree, struct cache_extent *pe)
 {
-	struct rb_node *found;
-
-	found = tree_insert(&tree->root, pe->start, pe->size, &pe->rb_node);
-	if (found)
-		return -EEXIST;
-
-	return 0;
+	return rb_insert(&tree->root, &pe->rb_node, cache_tree_comp_nodes);
 }
 
-int insert_cache_extent(struct cache_tree *tree, u64 start, u64 size)
+int add_cache_extent(struct cache_tree *tree, u64 start, u64 size)
 {
 	struct cache_extent *pe = alloc_cache_extent(start, size);
 	int ret;
-	ret = insert_existing_cache_extent(tree, pe);
+
+	if (!pe) {
+		fprintf(stderr, "memory allocation failed\n");
+		exit(1);
+	}
+
+	ret = insert_cache_extent(tree, pe);
 	if (ret)
 		free(pe);
+
 	return ret;
 }
 
 struct cache_extent *find_cache_extent(struct cache_tree *tree,
-					   u64 start, u64 size)
+				       u64 start, u64 size)
 {
-	struct rb_node *prev;
-	struct rb_node *ret;
+	struct rb_node *node;
 	struct cache_extent *entry;
-	ret = __tree_search(&tree->root, start, size, &prev);
-	if (!ret)
+	struct cache_extent_search_range range;
+
+	range.start = start;
+	range.size = size;
+	node = rb_search(&tree->root, &range, cache_tree_comp_range, NULL);
+	if (!node)
 		return NULL;
 
-	entry = rb_entry(ret, struct cache_extent, rb_node);
+	entry = rb_entry(node, struct cache_extent, rb_node);
 	return entry;
 }
 
-struct cache_extent *find_first_cache_extent(struct cache_tree *tree,
-						 u64 start)
+struct cache_extent *find_first_cache_extent(struct cache_tree *tree, u64 start)
 {
-	struct rb_node *prev;
-	struct rb_node *ret;
+	struct rb_node *next;
+	struct rb_node *node;
 	struct cache_extent *entry;
+	struct cache_extent_search_range range;
 
-	ret = __tree_search(&tree->root, start, 1, &prev);
-	if (!ret)
-		ret = prev;
-	if (!ret)
+	range.start = start;
+	range.size = 1;
+	node = rb_search(&tree->root, &range, cache_tree_comp_range, &next);
+	if (!node)
+		node = next;
+	if (!node)
 		return NULL;
-	entry = rb_entry(ret, struct cache_extent, rb_node);
+
+	entry = rb_entry(node, struct cache_extent, rb_node);
 	return entry;
 }
 
+struct cache_extent *first_cache_extent(struct cache_tree *tree)
+{
+	struct rb_node *node = rb_first(&tree->root);
+
+	if (!node)
+		return NULL;
+	return rb_entry(node, struct cache_extent, rb_node);
+}
+
 struct cache_extent *prev_cache_extent(struct cache_extent *pe)
 {
 	struct rb_node *node = rb_prev(&pe->rb_node);
@@ -162,9 +155,18 @@ struct cache_extent *next_cache_extent(struct cache_extent *pe)
 	return rb_entry(node, struct cache_extent, rb_node);
 }
 
-void remove_cache_extent(struct cache_tree *tree,
-				 struct cache_extent *pe)
+void remove_cache_extent(struct cache_tree *tree, struct cache_extent *pe)
 {
 	rb_erase(&pe->rb_node, &tree->root);
 }
 
+void cache_tree_free_extents(struct cache_tree *tree,
+			     free_cache_extent free_func)
+{
+	struct cache_extent *ce;
+
+	while ((ce = first_cache_extent(tree))) {
+		remove_cache_extent(tree, ce);
+		free_func(ce);
+	}
+}