Make prio heap dynamically sized

Signed-off-by: Mathieu Desnoyers <mathieu.desnoyers@efficios.com>

diff --git a/lib/prio_heap/prio_heap.c b/lib/prio_heap/prio_heap.c
index ae47ebb8e4ba8f0b378e9a31b71b56ba4e4c207f..d4fadda3be8a53229fd5926a96b009a48eb5f1db 100644 (file)
--- a/lib/prio_heap/prio_heap.c
+++ b/lib/prio_heap/prio_heap.c
@@ -1,8 +1,7 @@
 /*
  * prio_heap.c
  *
- * Static-sized priority heap containing pointers. Based on CLRS,
- * chapter 6.
+ * Priority heap containing pointers. Based on CLRS, chapter 6.
  *
  * Copyright 2011 - Mathieu Desnoyers <mathieu.desnoyers@efficios.com>
  *
@@ -18,29 +17,102 @@
  */
 
 #include <linux/slab.h>
-#include <linux/prio_heap.h>
+#include "prio_heap.h"
 
-int heap_init(struct ptr_heap *heap, size_t size,
-              gfp_t gfpmask, int gt(void *a, void *b))
+static
+size_t parent(size_t i)
 {
-       WARN_ON_ONCE(size == 0);
-       heap->ptrs = kmalloc(size * sizeof(void *), gfpmask);
-       if (!heap->ptrs)
+       return i >> 1;
+}
+
+static
+size_t left(size_t i)
+{
+       return i << 1;
+}
+
+static
+size_t right(size_t i)
+{
+       return (i << 1) + 1;
+}
+
+static
+int heap_grow(struct ptr_heap *heap, size_t new_len)
+{
+       void **new_ptrs;
+
+       if (heap->alloc_len >= new_len)
+               return 0;
+
+       heap->alloc_len = max_t(size_t, new_len, heap->alloc_len << 1);
+       new_ptrs = kmalloc(heap->alloc_len * sizeof(void *), heap->gfpmask);
+       if (!new_ptrs)
                return -ENOMEM;
-       heap->size = 0;
-       heap->max = size;
-       heap->gt = gt;
+       if (heap->ptrs)
+               memcpy(new_ptrs, heap->ptrs, heap->len * sizeof(void *));
+       kfree(heap->ptrs);
+       heap->ptrs = new_ptrs;
+       return 0;
+}
+
+static
+int heap_set_len(struct ptr_heap *heap, size_t new_len)
+{
+       int ret;
+
+       ret = heap_grow(heap, new_len);
+       if (ret)
+               return ret;
+       heap->len = new_len;
        return 0;
 }
 
+int heap_init(struct ptr_heap *heap, size_t alloc_len,
+             gfp_t gfpmask, int gt(void *a, void *b))
+{
+       heap->ptrs = NULL;
+       heap->len = 0;
+       heap->alloc_len = 0;
+       heap->gt = gt;
+       /*
+        * Minimum size allocated is 1 entry to ensure memory allocation
+        * never fails within heap_replace_max.
+        */
+       return heap_grow(heap, min_t(size_t, 1, alloc_len));
+}
+
 void heap_free(struct ptr_heap *heap)
 {
        kfree(heap->ptrs);
 }
 
-static void heapify(struct ptr_heap *heap, int pos)
+static void heapify(struct ptr_heap *heap, size_t i)
 {
-       /* TODO */
+       void **ptrs = heap->ptrs;
+       size_t l, r, largest;
+
+       for (;;) {
+               l = left(i);
+               r = right(i);
+               if (l <= heap->len && ptrs[l] > ptrs[i])
+                       largest = l;
+               else
+                       largest = i;
+               if (r <= heap->len && 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 {
+                       break;
+               }
+       }
 }
 
 void *heap_replace_max(struct ptr_heap *heap, void *p)
@@ -48,9 +120,9 @@ void *heap_replace_max(struct ptr_heap *heap, void *p)
        void *res;
        void **ptrs = heap->ptrs;
 
-       if (!heap->size) {
+       if (!heap->len) {
+               (void) heap_set_len(heap, 1);
                ptrs[0] = p;
-               heap->size = 1;
                return NULL;
        }
 
@@ -61,68 +133,54 @@ void *heap_replace_max(struct ptr_heap *heap, void *p)
        return res;
 }
 
-void *heap_insert(struct ptr_heap *heap, void *p)
+int heap_insert(struct ptr_heap *heap, void *p)
 {
        void **ptrs = heap->ptrs;
-       void *tmp = NULL;
-
-       if (heap->size < heap->max) {
-               /* Add the element to the end */
-               heap->ptrs[heap->size++] = p;
-               /* rebalance */
-               heapify(heap, 0);
-               return NULL;
-       }
+       int ret;
 
-       /*
-        * 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;
+       ret = heap_set_len(heap, heap->len + 1);
+       if (ret)
+               return ret;
+       /* Add the element to the end */
+       ptrs[heap->len - 1] = p;
+       /* rebalance */
+       heapify(heap, 0);
+       return 0;
 }
 
 void *heap_remove(struct ptr_heap *heap)
 {
        void **ptrs = heap->ptrs;
 
-       switch (heap->size) {
+       switch (heap->len) {
        case 0:
                return NULL;
        case 1:
-               heap->size = 0;
+               (void) heap_set_len(heap, 0);
                return ptrs[0];
        }
-
        /* Shrink, replace the current max by previous last entry and heapify */
-       return heap_replace_max(heap, ptrs[--heap->size]);
+       heap_set_len(heap, heap->len - 1);
+       return heap_replace_max(heap, ptrs[heap->len - 1]);
 }
 
 void *heap_cherrypick(struct ptr_heap *heap, void *p)
 {
        void **ptrs = heap->ptrs;
-       size_t pos, size = heap->size;
+       size_t pos, len = heap->len;
 
-       for (pos = 0; pos < size; pos++)
+       for (pos = 0; pos < len; pos++)
                if (ptrs[pos] == p)
                        goto found;
        return NULL;
 found:
-       if (heap->size == 1) {
-               heap->size = 0;
+       if (heap->len == 1) {
+               (void) heap_set_len(heap, 0);
                return ptrs[0];
        }
-       /*
-        * Replace p with previous last entry and heapify.
-        */
-       ptrs[pos] = ptrs[--heap->size];
+       /* Replace p with previous last entry and heapify. */
+       heap_set_len(heap, heap->len - 1);
+       ptrs[pos] = ptrs[heap->len - 1];
        heapify(heap, pos);
        return p;
 }

diff --git a/lib/prio_heap/prio_heap.h b/lib/prio_heap/prio_heap.h
index 3a3188534545dc18c86fb6a7785e925572b980b2..2674db21fb8496d5e5a7ceacd5ad26ce1af756f6 100644 (file)
--- a/lib/prio_heap/prio_heap.h
+++ b/lib/prio_heap/prio_heap.h
@@ -4,8 +4,7 @@
 /*
  * prio_heap.h
  *
- * Static-sized priority heap containing pointers. Based on CLRS,
- * chapter 6.
+ * Priority heap containing pointers. Based on CLRS, chapter 6.
  *
  * Copyright 2011 - Mathieu Desnoyers <mathieu.desnoyers@efficios.com>
  *
@@ -23,9 +22,10 @@
 #include <linux/gfp.h>
 
 struct ptr_heap {
-       size_t size, max;
+       size_t len, alloc_len;
        void **ptrs;
        int (*gt)(void *a, void *b);
+       gfp_t gfpmask;
 };
 
 /**
@@ -37,20 +37,21 @@ struct ptr_heap {
  */
 static inline void *heap_maximum(const struct ptr_heap *heap)
 {
-       return heap->size ? heap->ptrs[0] : NULL;
+       return heap->len ? heap->ptrs[0] : NULL;
 }
 
 /**
  * heap_init - initialize the heap
  * @heap: the heap to initialize
- * @size: maximum number of elements
+ * @alloc_len: number of elements initially allocated
  * @gfp: allocation flags
  * @gt: function to compare the elements
  *
  * Returns -ENOMEM if out of memory.
  */
-extern int heap_init(struct ptr_heap *heap, size_t size,
-                    gfp_t gfpmask, int gt(void *a, void *b));
+extern int heap_init(struct ptr_heap *heap,
+                    size_t alloc_len, gfp_t gfpmask,
+                    int gt(void *a, void *b));
 
 /**
  * heap_free - free the heap
@@ -61,12 +62,13 @@ extern void heap_free(struct ptr_heap *heap);
 /**
  * heap_insert - insert an element into the heap
  * @heap: the heap to be operated on
+ * @p: the element to add
  *
- * Insert an element into the heap. If the heap is full, return the
- * largest element between those previously present in the heap and the
- * element being added, else return NULL.
+ * Insert an element into the heap.
+ *
+ * Returns -ENOMEM if out of memory.
  */
-extern void *heap_insert(struct ptr_heap *heap, void *p);
+extern int heap_insert(struct ptr_heap *heap, void *p);
 
 /**
  * heap_remove - remove the largest element from the heap
@@ -98,7 +100,7 @@ extern void *heap_cherrypick(struct ptr_heap *heap, void *p);
  * the heap is empty.
  *
  * This is the equivalent of calling heap_remove() and then heap_insert(), but
- * it only rebalances the heap once.
+ * it only rebalances the heap once. It never allocates memory.
  */
 extern void *heap_replace_max(struct ptr_heap *heap, void *p);