X-Git-Url: http://git.efficios.com/?a=blobdiff_plain;f=lib%2Fprio_heap.c;h=a37e64c89542d10f017f8da835e01272dab2e2ef;hb=a3dde804badf7409ce9af45a70b9f33a2c42173a;hp=0a99a3d635c91bdb78f03a194ce76e86f13689ec;hpb=1eb0c69cf759f535b1a926e667ae3f3aed486c9f;p=babeltrace.git

diff --git a/lib/prio_heap.c b/lib/prio_heap.c
index 0a99a3d6..a37e64c8 100644
--- a/lib/prio_heap.c
+++ b/lib/prio_heap.c
@@ -18,46 +18,64 @@
  */
 
 #include <babeltrace/prio_heap.h>
+#include <babeltrace/babeltrace-internal.h>
 #include <errno.h>
 #include <stdlib.h>
 #include <string.h>
+#include <assert.h>
 
 #ifndef max_t
 #define max_t(type, a, b)	\
 	((type) (a) > (type) (b) ? (type) (a) : (type) (b))
 #endif
 
-static __attribute__((unused))
+#ifdef DEBUG_HEAP
+void check_heap(const struct ptr_heap *heap)
+{
+	size_t i;
+
+	if (!heap->len)
+		return;
+
+	for (i = 1; i < heap->len; i++)
+		assert(!heap->gt(heap->ptrs[i], heap->ptrs[0]));
+}
+#endif
+
+static
 size_t parent(size_t i)
 {
-	return i >> 1;
+	return (i - 1) >> 1;
 }
 
 static
 size_t left(size_t i)
 {
-	return i << 1;
+	return (i << 1) + 1;
 }
 
 static
 size_t right(size_t i)
 {
-	return (i << 1) + 1;
+	return (i << 1) + 2;
 }
 
+/*
+ * Copy of heap->ptrs pointer is invalid after heap_grow.
+ */
 static
 int heap_grow(struct ptr_heap *heap, size_t new_len)
 {
 	void **new_ptrs;
 
-	if (heap->alloc_len >= new_len)
+	if (likely(heap->alloc_len >= new_len))
 		return 0;
 
 	heap->alloc_len = max_t(size_t, new_len, heap->alloc_len << 1);
 	new_ptrs = calloc(heap->alloc_len, sizeof(void *));
-	if (!new_ptrs)
+	if (unlikely(!new_ptrs))
 		return -ENOMEM;
-	if (heap->ptrs)
+	if (likely(heap->ptrs))
 		memcpy(new_ptrs, heap->ptrs, heap->len * sizeof(void *));
 	free(heap->ptrs);
 	heap->ptrs = new_ptrs;
@@ -70,7 +88,7 @@ int heap_set_len(struct ptr_heap *heap, size_t new_len)
 	int ret;
 
 	ret = heap_grow(heap, new_len);
-	if (ret)
+	if (unlikely(ret))
 		return ret;
 	heap->len = new_len;
 	return 0;
@@ -101,94 +119,98 @@ static void heapify(struct ptr_heap *heap, size_t i)
 	size_t l, r, largest;
 
 	for (;;) {
+		void *tmp;
+
 		l = left(i);
 		r = right(i);
-		if (l <= heap->len && ptrs[l] > ptrs[i])
+		if (l < heap->len && heap->gt(ptrs[l], ptrs[i]))
 			largest = l;
 		else
 			largest = i;
-		if (r <= heap->len && ptrs[r] > ptrs[largest])
+		if (r < heap->len && heap->gt(ptrs[r], ptrs[largest]))
 			largest = r;
-		if (largest != i) {
-			void *tmp;
-
-			tmp = ptrs[i];
-			ptrs[i] = ptrs[largest];
-			ptrs[largest] = tmp;
-			i = largest;
-			continue;
-		} else {
+		if (unlikely(largest == i))
 			break;
-		}
+		tmp = ptrs[i];
+		ptrs[i] = ptrs[largest];
+		ptrs[largest] = tmp;
+		i = largest;
 	}
+	check_heap(heap);
 }
 
 void *heap_replace_max(struct ptr_heap *heap, void *p)
 {
 	void *res;
-	void **ptrs = heap->ptrs;
 
-	if (!heap->len) {
+	if (unlikely(!heap->len)) {
 		(void) heap_set_len(heap, 1);
-		ptrs[0] = p;
+		heap->ptrs[0] = p;
+		check_heap(heap);
 		return NULL;
 	}
 
 	/* Replace the current max and heapify */
-	res = ptrs[0];
-	ptrs[0] = p;
+	res = heap->ptrs[0];
+	heap->ptrs[0] = p;
 	heapify(heap, 0);
 	return res;
 }
 
 int heap_insert(struct ptr_heap *heap, void *p)
 {
-	void **ptrs = heap->ptrs;
+	void **ptrs;
+	size_t pos;
 	int ret;
 
 	ret = heap_set_len(heap, heap->len + 1);
-	if (ret)
+	if (unlikely(ret))
 		return ret;
-	/* Add the element to the end */
-	ptrs[heap->len - 1] = p;
-	/* rebalance */
-	heapify(heap, 0);
+	ptrs = heap->ptrs;
+	pos = heap->len - 1;
+	while (pos > 0 && heap->gt(p, ptrs[parent(pos)])) {
+		/* Move parent down until we find the right spot */
+		ptrs[pos] = ptrs[parent(pos)];
+		pos = parent(pos);
+	}
+	ptrs[pos] = p;
+	check_heap(heap);
 	return 0;
 }
 
 void *heap_remove(struct ptr_heap *heap)
 {
-	void **ptrs = heap->ptrs;
-
 	switch (heap->len) {
 	case 0:
 		return NULL;
 	case 1:
 		(void) heap_set_len(heap, 0);
-		return ptrs[0];
+		return heap->ptrs[0];
 	}
 	/* Shrink, replace the current max by previous last entry and heapify */
 	heap_set_len(heap, heap->len - 1);
-	return heap_replace_max(heap, ptrs[heap->len - 1]);
+	/* len changed. previous last entry is at heap->len */
+	return heap_replace_max(heap, heap->ptrs[heap->len]);
 }
 
 void *heap_cherrypick(struct ptr_heap *heap, void *p)
 {
-	void **ptrs = heap->ptrs;
 	size_t pos, len = heap->len;
 
 	for (pos = 0; pos < len; pos++)
-		if (ptrs[pos] == p)
+		if (unlikely(heap->ptrs[pos] == p))
 			goto found;
 	return NULL;
 found:
-	if (heap->len == 1) {
+	if (unlikely(heap->len == 1)) {
 		(void) heap_set_len(heap, 0);
-		return ptrs[0];
+		check_heap(heap);
+		return heap->ptrs[0];
 	}
 	/* Replace p with previous last entry and heapify. */
 	heap_set_len(heap, heap->len - 1);
-	ptrs[pos] = ptrs[heap->len - 1];
+	/* len changed. previous last entry is at heap->len */
+	heap->ptrs[pos] = heap->ptrs[heap->len];
 	heapify(heap, pos);
 	return p;
 }