[PATCH] btree: Add custom allocator interface.

From: Stuart Longland <[email protected]>

This provides a way for btree to be used with external allocator
libraries such as the tal or talloc modules.
---
 ccan/btree/btree.c | 194 +++++++++++++++++++++++++++++++----------------------
 ccan/btree/btree.h |  29 +++++---
 2 files changed, 136 insertions(+), 87 deletions(-)

diff --git a/ccan/btree/btree.c b/ccan/btree/btree.c
index 636edbc..5bb2435 100644
--- a/ccan/btree/btree.c
+++ b/ccan/btree/btree.c
@@ -29,12 +29,15 @@
 #define MAX (BTREE_ITEM_MAX)
 #define MIN (BTREE_ITEM_MAX >> 1)
 
-static struct btree_node *node_alloc(int internal);
+static struct btree_node *node_alloc(const struct btree_allocator* alloc, int internal);
 static void node_delete(struct btree_node *node, struct btree *btree);
+static void node_free(struct btree_node *node);
 
 static void branch_begin(btree_iterator iter);
 static void branch_end(btree_iterator iter);
 static void begin_end_lr(btree_iterator iter, struct btree_node *node, int lr);
+static void* default_malloc(const struct btree_allocator* alloc, size_t size);
+static void default_free(const struct btree_allocator* alloc, void* ptr);
 
 /*
  * If iter->node has parent, returns 1 and ascends the iterator such that
@@ -65,11 +68,19 @@ static int node_walk_forward(const struct btree_node *node,
 
 struct btree *btree_new(btree_search_t search)
 {
-	struct btree *btree = calloc(1, sizeof(struct btree));
-	struct btree_node *node = node_alloc(0);
+	return btree_new_with_alloc(search, &BTREE_DEFAULT_ALLOCATOR);
+}
+
+struct btree *btree_new_with_alloc(btree_search_t search,
+		const struct btree_allocator *alloc)
+{
+	struct btree *btree = (*(alloc->malloc))(alloc, sizeof(struct btree));
+	struct btree_node *node = node_alloc(alloc, 0);
 		node->parent = NULL;
 		node->count = 0;
 		node->depth = 0;
+	memset(btree, 0, sizeof(struct btree));
+	btree->alloc = alloc;
 	btree->root = node;
 	btree->search = search;
 	btree->multi = false;
@@ -79,16 +90,16 @@ struct btree *btree_new(btree_search_t search)
 void btree_delete(struct btree *btree)
 {
 	node_delete(btree->root, btree);
-	free(btree);
+	(*(btree->alloc->free))(btree->alloc, btree);
 }
 
 bool btree_insert(struct btree *btree, const void *item)
 {
 	btree_iterator iter;
-	
+
 	if (btree_find_last(btree, item, iter) && !btree->multi)
 		return false;
-	
+
 	btree_insert_at(iter, item);
 	return true;
 }
@@ -98,41 +109,41 @@ bool btree_remove(struct btree *btree, const void *key)
 	btree_iterator iter;
 	bool success = false;
 	bool multi = btree->multi;
-	
+
 	do {
 		if (btree_find_first(btree, key, iter)) {
 			btree_remove_at(iter);
 			success = true;
 		}
 	} while (multi);
-	
+
 	return success;
 }
 
 void *btree_lookup(struct btree *btree, const void *key)
 {
 	btree_iterator iter;
-	
+
 	if (btree_find_first(btree, key, iter))
 		return iter->item;
-	
+
 	return NULL;
 }
 
 int btree_begin_end_lr(const struct btree *btree, btree_iterator iter, int lr)
 {
 	struct btree_node *node;
-	
+
 	iter->btree = (struct btree *)btree;
 	begin_end_lr(iter, btree->root, lr);
-	
+
 	/* Set iter->item if any items exist. */
 	node = iter->node;
 	if (node->count) {
 		iter->item = (void*)node->item[iter->k - lr];
 		return 1;
 	}
-	
+
 	return 0;
 }
 
@@ -147,7 +158,7 @@ int btree_deref(btree_iterator iter)
 			}
 		} while (iter->k >= iter->node->count);
 	}
-	
+
 	iter->item = (void*)iter->node->item[iter->k];
 	return 1;
 }
@@ -165,7 +176,7 @@ int btree_prev(btree_iterator iter)
 			}
 		} while (iter->k == 0);
 	}
-	
+
 	iter->item = (void*)iter->node->item[--iter->k];
 	return 1;
 }
@@ -188,28 +199,28 @@ int btree_find_lr(const struct btree *btree, const void *key,
 	unsigned int k;
 	unsigned int depth;
 	int found = 0;
-	
+
 	iter->btree = (struct btree *)btree;
 	iter->item = NULL;
-	
+
 	depth = node->depth;
 	for (;;) {
 		int f = 0;
 		k = btree->search(key, node->item, node->count, lr, &f);
-		
+
 		if (f) {
 			iter->item = (void*)node->item[k - lr];
 			found = 1;
 		}
 		if (!depth--)
 			break;
-		
+
 		node = node->branch[k];
 	}
-	
+
 	iter->node = node;
 	iter->k = k;
-	
+
 	return found;
 }
 
@@ -231,10 +242,10 @@ void btree_insert_at(btree_iterator iter, const void *item)
 	struct btree_node *xr = NULL;
 	struct btree_node *p;
 	struct btree *btree = iter->btree;
-	
+
 	/* btree_insert_at always sets iter->item to item. */
 	iter->item = (void*)item;
-	
+
 	/*
 	 * If node is not a leaf, fall to the end of the left branch of item[k]
 	 * so that it will be a leaf. This does not modify the iterator's logical
@@ -242,7 +253,7 @@ void btree_insert_at(btree_iterator iter, const void *item)
 	 */
 	if (iter->node->depth)
 		branch_end(iter);
-	
+
 	/*
 	 * First try inserting item into this node.
 	 * If it's too big, split it, and repeat by
@@ -254,23 +265,23 @@ void btree_insert_at(btree_iterator iter, const void *item)
 	} else {
 		for (;;) {
 			node_split(&x, &xr, iter->node, iter->k);
-			
+
 			if (!ascend(iter))
 				break;
-			
+
 			if (iter->node->count < MAX) {
 				node_insert(x, xr, iter->node, iter->k);
 				goto finished;
 			}
 		}
-		
+
 		/*
 		 * If splitting came all the way up to the root, create a new root whose
 		 * left branch is the current root, median is x, and right branch is the
 		 * half split off from the root.
 		 */
 		assert(iter->node == btree->root);
-		p = node_alloc(1);
+		p = node_alloc(btree->alloc, 1);
 		p->parent = NULL;
 		p->count = 1;
 		p->depth = btree->root->depth + 1;
@@ -283,7 +294,7 @@ void btree_insert_at(btree_iterator iter, const void *item)
 			xr->k = 1;
 		btree->root = p;
 	}
-	
+
 finished:
 	btree->count++;
 	iter->node = NULL;
@@ -293,10 +304,10 @@ int btree_remove_at(btree_iterator iter)
 {
 	struct btree *btree = iter->btree;
 	struct btree_node *root;
-	
+
 	if (!btree_deref(iter))
 		return 0;
-	
+
 	if (!iter->node->depth) {
 		node_remove_leaf_item(iter->node, iter->k);
 		if (iter->node->count >= MIN || !iter->node->parent)
@@ -307,21 +318,21 @@ int btree_remove_at(btree_iterator iter)
 		 * with its successor (which will always be in a leaf), then remove
 		 * the original copy of the successor.
 		 */
-		
+
 		/* Save pointer to condemned item. */
 		const void **x = &iter->node->item[iter->k];
-		
+
 		/* Descend to successor. */
 		iter->k++;
 		branch_begin(iter);
-		
+
 		/* Replace condemned item with successor. */
 		*x = iter->node->item[0];
-		
+
 		/* Remove successor. */
 		node_remove_leaf_item(iter->node, 0);
 	}
-	
+
 	/*
 	 * Restore nodes that fall under their minimum count.  This may
 	 * propagate all the way up to the root.
@@ -333,7 +344,7 @@ int btree_remove_at(btree_iterator iter)
 			break;
 		node_restore(iter->node, iter->k);
 	}
-	
+
 	/*
 	 * If combining came all the way up to the root, and it has no more
 	 * dividers, delete it and make its only branch the root.
@@ -344,9 +355,9 @@ int btree_remove_at(btree_iterator iter)
 	if (root->count == 0) {
 		btree->root = root->branch[0];
 		btree->root->parent = NULL;
-		free(root);
+		node_free(root);
 	}
-	
+
 finished:
 	btree->count--;
 	iter->node = NULL;
@@ -363,7 +374,7 @@ static int elevate(btree_iterator a, btree_iterator b)
 {
 	while (a->node->depth < b->node->depth)
 		ascend(a);
-	
+
 	if (a->k == b->k)
 		return -1;
 	return 0;
@@ -373,16 +384,16 @@ int btree_cmp_iters(const btree_iterator iter_a, const btree_iterator iter_b)
 {
 	btree_iterator a = {*iter_a}, b = {*iter_b};
 	int ad, bd;
-	
+
 	ad = btree_deref(a);
 	bd = btree_deref(b);
-	
+
 	/* Check cases where one or both iterators are at the end. */
 	if (!ad)
 		return bd ? 1 : 0;
 	if (!bd)
 		return ad ? -1 : 0;
-	
+
 	/* Bring iterators to the same depth. */
 	if (a->node->depth < b->node->depth) {
 		if (elevate(a, b))
@@ -391,19 +402,19 @@ int btree_cmp_iters(const btree_iterator iter_a, const btree_iterator iter_b)
 		if (elevate(b, a))
 			return 1;
 	}
-	
+
 	/* Bring iterators to the same node. */
 	while (a->node != b->node) {
 		ascend(a);
 		ascend(b);
 	}
-	
+
 	/* Now we can compare by k directly. */
 	if (a->k < b->k)
 		return -1;
 	if (a->k > b->k)
 		return 1;
-	
+
 	return 0;
 }
 
@@ -421,20 +432,23 @@ btree_search_implement
 
 /************************* Private functions *************************/
 
-static struct btree_node *node_alloc(int internal)
+static struct btree_node *node_alloc(const struct btree_allocator *alloc,
+		int internal)
 {
 	struct btree_node *node;
 	size_t isize = internal
 		? sizeof(struct btree_node*) * (BTREE_ITEM_MAX+1)
 		: 0;
-	node = malloc(sizeof(struct btree_node) + isize);
+	node = (*(alloc->malloc))(alloc, \
+			sizeof(struct btree_node) + isize);
+	node->alloc = alloc;
 	return node;
 }
 
 static void node_delete(struct btree_node *node, struct btree *btree)
 {
 	unsigned int i, count = node->count;
-	
+
 	if (!node->depth) {
 		if (btree->destroy) {
 			for (i=0; i<count; i++)
@@ -448,8 +462,14 @@ static void node_delete(struct btree_node *node, struct btree *btree)
 		}
 		node_delete(node->branch[count], btree);
 	}
-	
-	free(node);
+
+	node_free(node);
+}
+
+/* Free the allocated node using its allocator */
+static void node_free(struct btree_node *node)
+{
+	(*(node->alloc->free))(node->alloc, node);
 }
 
 /* Set iter to beginning of branch pointed to by iter. */
@@ -493,11 +513,11 @@ static void node_insert(const void *x, struct btree_node *xr,
 				struct btree_node *p, unsigned int k)
 {
 	unsigned int i;
-	
+
 	for (i = p->count; i-- > k;)
 		p->item[i+1] = p->item[i];
 	p->item[k] = x;
-	
+
 	if (p->depth) {
 		k++;
 		for (i = p->count+1; i-- > k;) {
@@ -508,7 +528,7 @@ static void node_insert(const void *x, struct btree_node *xr,
 		xr->parent = p;
 		xr->k = k;
 	}
-	
+
 	p->count++;
 }
 
@@ -524,7 +544,7 @@ static void node_split(const void **x, struct btree_node **xr,
 {
 	unsigned int i, split;
 	struct btree_node *l = p, *r;
-	
+
 	/*
 	 * If k <= MIN, item will be inserted into left subtree, so give l
 	 * fewer items initially.
@@ -535,17 +555,17 @@ static void node_split(const void **x, struct btree_node **xr,
 		split = MIN;
 	else
 		split = MIN + 1;
-	
+
 	/*
 	 * If l->depth is 0, allocate a leaf node.
 	 * Otherwise, allocate an internal node.
 	 */
-	r = node_alloc(l->depth);
-	
+	r = node_alloc(l->alloc, l->depth);
+
 	/* l and r will be siblings, so they will have the same parent and depth. */
 	r->parent = l->parent;
 	r->depth = l->depth;
-	
+
 	/*
 	 * Initialize items/branches of right side.
 	 * Do not initialize r's leftmost branch yet because we don't know
@@ -561,11 +581,11 @@ static void node_split(const void **x, struct btree_node **xr,
 			r->branch[i-split]->k = i-split;
 		}
 	}
-	
+
 	/* Update counts. */
 	l->count = split;
 	r->count = MAX - split;
-	
+
 	/*
 	 * The nodes are now split, but the key isn't inserted yet.
 	 *
@@ -576,7 +596,7 @@ static void node_split(const void **x, struct btree_node **xr,
 		node_insert(*x, *xr, l, k);
 	else
 		node_insert(*x, *xr, r, k - split);
-	
+
 	/*
 	 * Give l's rightmost branch to r because l's rightmost item
 	 * is going up to become the median.
@@ -586,7 +606,7 @@ static void node_split(const void **x, struct btree_node **xr,
 		r->branch[0]->parent = r;
 		r->branch[0]->k = 0;
 	}
-	
+
 	/*
 	 * Take up l's rightmost item to make it the median.
 	 * That item's right branch is now r.
@@ -643,24 +663,24 @@ static void move_left(struct btree_node *node, unsigned int k)
 {
 	struct btree_node *l = node->branch[k], *r = node->branch[k+1], *mv;
 	unsigned int i;
-	
+
 	l->item[l->count] = node->item[k];
 	node->item[k] = r->item[0];
 	for (i = 1; i < r->count; i++)
 		r->item[i-1] = r->item[i];
-	
+
 	if (r->depth) {
 		mv = r->branch[0];
 		l->branch[l->count+1] = mv;
 		mv->parent = l;
 		mv->k = l->count+1;
-		
+
 		for (i = 1; i <= r->count; i++) {
 			r->branch[i-1] = r->branch[i];
 			r->branch[i-1]->k = i-1;
 		}
 	}
-	
+
 	l->count++;
 	r->count--;
 }
@@ -669,12 +689,12 @@ static void move_right(struct btree_node *node, unsigned int k)
 {
 	struct btree_node *l = node->branch[k], *r = node->branch[k+1];
 	unsigned int i;
-	
+
 	for (i = r->count; i--;)
 		r->item[i+1] = r->item[i];
 	r->item[0] = node->item[k];
 	node->item[k] = l->item[l->count-1];
-	
+
 	if (r->depth) {
 		for (i = r->count+1; i--;) {
 			r->branch[i+1] = r->branch[i];
@@ -684,7 +704,7 @@ static void move_right(struct btree_node *node, unsigned int k)
 		r->branch[0]->parent = r;
 		r->branch[0]->k = 0;
 	}
-	
+
 	l->count--;
 	r->count++;
 }
@@ -695,12 +715,12 @@ static void combine(struct btree_node *node, unsigned int k)
 	struct btree_node *l = node->branch[k], *r = node->branch[k+1], *mv;
 	const void **o = &l->item[l->count];
 	unsigned int i;
-	
+
 	//append node->item[k] followed by right node's items to left node
 	*o++ = node->item[k];
 	for (i=0; i<r->count; i++)
 		*o++ = r->item[i];
-	
+
 	//if applicable, append right node's branches to left node
 	if (r->depth) {
 		for (i=0; i<=r->count; i++) {
@@ -710,25 +730,25 @@ static void combine(struct btree_node *node, unsigned int k)
 			mv->k = l->count + i + 1;
 		}
 	}
-	
+
 	//remove k and its right branch from parent node
 	for (i = k+1; i < node->count; i++) {
 		node->item[i-1] = node->item[i];
 		node->branch[i] = node->branch[i+1];
 		node->branch[i]->k = i;
 	}
-	
+
 	//don't forget to update the left and parent node's counts and to free the right node
 	l->count += r->count + 1;
 	node->count--;
-	free(r);
+	node_free(r);
 }
 
 static int node_walk_backward(const struct btree_node *node,
 				btree_action_t action, void *ctx)
 {
 	unsigned int i, count = node->count;
-	
+
 	if (!node->depth) {
 		for (i=count; i--;)
 			if (!action((void*)node->item[i], ctx))
@@ -743,7 +763,7 @@ static int node_walk_backward(const struct btree_node *node,
 				return 0;
 		}
 	}
-	
+
 	return 1;
 }
 
@@ -751,7 +771,7 @@ static int node_walk_forward(const struct btree_node *node,
 				btree_action_t action, void *ctx)
 {
 	unsigned int i, count = node->count;
-	
+
 	if (!node->depth) {
 		for (i=0; i<count; i++)
 			if (!action((void*)node->item[i], ctx))
@@ -766,6 +786,22 @@ static int node_walk_forward(const struct btree_node *node,
 		if (!node_walk_forward(node->branch[count], action, ctx))
 			return 0;
 	}
-	
+
 	return 1;
 }
+
+/* Default allocator implementation */
+const struct btree_allocator BTREE_DEFAULT_ALLOCATOR = {
+	.malloc = default_malloc,
+	.free = default_free,
+};
+
+static void* default_malloc(const struct btree_allocator* alloc, size_t size)
+{
+	return malloc(size);
+}
+
+static void default_free(const struct btree_allocator* alloc, void* ptr)
+{
+	free(ptr);
+}
diff --git a/ccan/btree/btree.h b/ccan/btree/btree.h
index fdf198d..a413be8 100644
--- a/ccan/btree/btree.h
+++ b/ccan/btree/btree.h
@@ -43,20 +43,30 @@ btree_lookup
  */
 #define BTREE_ITEM_MAX 20
 
+/* Allocator interface, for use with tal/talloc, etc */
+struct btree_allocator {
+	void *(*malloc)(const struct btree_allocator *alloc, size_t size);
+	void (*free)(const struct btree_allocator *alloc, void *ptr);
+};
+
+/* Default allocator using malloc/free in libc */
+extern const struct btree_allocator BTREE_DEFAULT_ALLOCATOR;
+
 struct btree_node {
 	struct btree_node *parent;
-	
+	const struct btree_allocator *alloc;
+
 	/* Number of items (rather than branches). */
 	unsigned char count;
-	
+
 	/* 0 if node is a leaf, 1 if it has leaf children, etc. */
 	unsigned char depth;
-	
+
 	/* node->parent->branch[node->k] == this */
 	unsigned char k;
-	
+
 	const void *item[BTREE_ITEM_MAX];
-	
+
 	/*
 	 * Allocated to BTREE_ITEM_MAX+1 items if this is
 	 * an internal node, 0 items if it is a leaf.
@@ -68,7 +78,7 @@ typedef struct btree_iterator_s {
 	struct btree *btree;
 	struct btree_node *node;
 	unsigned int k;
-	
+
 	/*
 	 * The relationship between item and (node, k) depends on what function
 	 * set it.  It is mainly for convenience.
@@ -105,10 +115,11 @@ typedef int (*btree_action_t)(void *item, void *ctx);
 struct btree {
 	struct btree_node *root;
 	size_t count; /* Total number of items in B-tree */
-	
+
 	btree_search_t search;
 	bool multi;
-	
+
+	const struct btree_allocator *alloc;
 	/*
 	 * These are set to NULL by default.
 	 *
@@ -123,6 +134,8 @@ struct btree {
 };
 
 struct btree *btree_new(btree_search_t search);
+struct btree *btree_new_with_alloc(btree_search_t search,
+		const struct btree_allocator *alloc);
 void btree_delete(struct btree *btree);
 
 /* Inserts an item into the btree.  If an item already exists that is equal
-- 
2.4.10

_______________________________________________
ccan mailing list
[email protected]
https://lists.ozlabs.org/listinfo/ccan