Annotate prio heap likely branch

[babeltrace.git] / lib / prio_heap.c

diff --git a/lib/prio_heap.c b/lib/prio_heap.c
index f7f351e2a5d799dfabca90700ed66d6faaab2d64..a37e64c89542d10f017f8da835e01272dab2e2ef 100644 (file)
--- a/lib/prio_heap.c
+++ b/lib/prio_heap.c
@@ -18,31 +18,46 @@
  */
 
 #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
 
+#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;
 }
 
 /*
@@ -53,14 +68,14 @@ 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;
@@ -73,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;
@@ -104,35 +119,34 @@ 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;
 
-       if (!heap->len) {
+       if (unlikely(!heap->len)) {
                (void) heap_set_len(heap, 1);
                heap->ptrs[0] = p;
+               check_heap(heap);
                return NULL;
        }
 
@@ -150,31 +164,17 @@ int heap_insert(struct ptr_heap *heap, void *p)
        int ret;
 
        ret = heap_set_len(heap, heap->len + 1);
-       if (ret)
+       if (unlikely(ret))
                return ret;
        ptrs = heap->ptrs;
-       /* Add the element to the end */
-       ptrs[heap->len - 1] = p;
        pos = heap->len - 1;
-       /* Bubble it up to the appropriate position. */
-       for (;;) {
-               if (pos > 0 && heap->gt(ptrs[pos], ptrs[parent(pos)])) {
-                       void *tmp;
-
-                       /* Need to exchange */
-                       tmp = ptrs[pos];
-                       ptrs[pos] = ptrs[parent(pos)];
-                       ptrs[parent(pos)] = tmp;
-                       pos = parent(pos);
-                       /*
-                        * No need to rebalance: if we are larger than
-                        * our parent, we are necessarily larger than
-                        * its other child.
-                        */
-               } else {
-                       break;
-               }
+       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;
 }
 
@@ -189,7 +189,8 @@ void *heap_remove(struct ptr_heap *heap)
        }
        /* Shrink, replace the current max by previous last entry and heapify */
        heap_set_len(heap, heap->len - 1);
-       return heap_replace_max(heap, 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)
@@ -197,17 +198,19 @@ void *heap_cherrypick(struct ptr_heap *heap, void *p)
        size_t pos, len = heap->len;
 
        for (pos = 0; pos < len; pos++)
-               if (heap->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);
+               check_heap(heap);
                return heap->ptrs[0];
        }
        /* Replace p with previous last entry and heapify. */
        heap_set_len(heap, heap->len - 1);
-       heap->ptrs[pos] = heap->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;
 }