From 7369e702eeca267896917095bb7cb34ac3b492d1 Mon Sep 17 00:00:00 2001 From: Mathieu Desnoyers Date: Mon, 23 May 2011 15:56:32 -0400 Subject: [PATCH] Make prio heap dynamically sized Signed-off-by: Mathieu Desnoyers --- lib/prio_heap/prio_heap.c | 158 ++++++++++++++++++++++++++------------ lib/prio_heap/prio_heap.h | 26 ++++--- 2 files changed, 122 insertions(+), 62 deletions(-) diff --git a/lib/prio_heap/prio_heap.c b/lib/prio_heap/prio_heap.c index ae47ebb8..d4fadda3 100644 --- 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 * @@ -18,29 +17,102 @@ */ #include -#include +#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 3a318853..2674db21 100644 --- 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 * @@ -23,9 +22,10 @@ #include 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); -- 2.34.1