Implement missing prio heap functions (MIT-licensed)

[lttng-modules.git] / lib / prio_heap / prio_heap.c

diff --git a/lib/prio_heap/prio_heap.c b/lib/prio_heap/prio_heap.c
index 8945c2a019804f7d0afe9e126b8c4a1b4dcce79d..0c9bb607510829a0ee62bd3a919d41e4b4b7440e 100644 (file)
--- a/lib/prio_heap/prio_heap.c
+++ b/lib/prio_heap/prio_heap.c
@@ -1,23 +1,34 @@
 /*
- * LICENSING: this file is copied from the Linux kernel. We should therefore
- * assume a GPLv2 license for the code that comes from the Linux mainline.
- */
-
-/*
- * Static-sized priority heap containing pointers. Based on CLR, chapter 7.
+ * prio_heap.c
+ *
+ * Static-sized priority heap containing pointers. Based on CLRS,
+ * chapter 6.
+ *
+ * Copyright 2011 - Mathieu Desnoyers <mathieu.desnoyers@efficios.com>
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
  */
 
 #include <linux/slab.h>
 #include <linux/prio_heap.h>
 
-int heap_init(struct ptr_heap *heap, size_t size, gfp_t gfp_mask,
-             int (*gt)(void *, void *))
+int heap_init(struct ptr_heap *heap, size_t size,
+              gfp_t gfpmask, int gt(void *a, void *b))
 {
-       heap->ptrs = kmalloc(size, gfp_mask);
+       WARN_ON_ONCE(size == 0);
+       heap->ptrs = kmalloc(size * sizeof(void *), gfpmask);
        if (!heap->ptrs)
                return -ENOMEM;
        heap->size = 0;
-       heap->max = size / sizeof(void *);
+       heap->max = size;
        heap->gt = gt;
        return 0;
 }
@@ -70,27 +81,29 @@ void *heap_replace_max(struct ptr_heap *heap, void *p)
 void *heap_insert(struct ptr_heap *heap, void *p)
 {
        void **ptrs = heap->ptrs;
-       int pos;
+       void *tmp = NULL;
 
        if (heap->size < heap->max) {
-               /* Heap insertion */
-               pos = heap->size++;
-               while (pos > 0 && heap->gt(p, ptrs[(pos-1)/2])) {
-                       ptrs[pos] = ptrs[(pos-1)/2];
-                       pos = (pos-1)/2;
-               }
-               ptrs[pos] = p;
+               /* Add the element to the end */
+               heap->ptrs[heap->size++] = p;
+               /* rebalance */
+               heapify(heap, 0);
                return NULL;
        }
 
-       /* The heap is full, so something will have to be dropped */
-
-       /* If the new pointer is greater than the current max, drop it */
-       if (heap->gt(p, ptrs[0]))
-               return p;
-
-       /* Replace the current max and heapify */
-       return heap_replace_max(heap, p);
+       /*
+        * Full. We need to replace the largest (if we are
+        * smaller or equal to this element).
+        */
+       if (heap->gt(ptrs[0], p)) {
+               tmp = ptrs[0];
+               ptrs[0] = p;
+               /* rebalance */
+               heapify(heap, 0);
+       } else {
+               tmp = p;
+       }
+       return tmp;
 }
 
 void *heap_remove(struct ptr_heap *heap)