From a29e892a9ba3f8bf865787fde21bfd12bf8bdf26 Mon Sep 17 00:00:00 2001
From: Mathieu Desnoyers <mathieu.desnoyers@efficios.com>
Date: Mon, 23 May 2011 17:37:51 -0400
Subject: [PATCH] Fix prio heap insertion and reallocation handling

Signed-off-by: Mathieu Desnoyers <mathieu.desnoyers@efficios.com>
---
 lib/prio_heap/prio_heap.c | 48 ++++++++++++++++++++++++++-------------
 1 file changed, 32 insertions(+), 16 deletions(-)

diff --git a/lib/prio_heap/prio_heap.c b/lib/prio_heap/prio_heap.c
index b518c552..c987d5de 100644
--- a/lib/prio_heap/prio_heap.c
+++ b/lib/prio_heap/prio_heap.c
@@ -19,7 +19,7 @@
 #include <linux/slab.h>
 #include "prio_heap.h"
 
-static __attribute__((unused))
+static
 size_t parent(size_t i)
 {
 	return i >> 1;
@@ -37,6 +37,9 @@ size_t right(size_t i)
 	return (i << 1) + 1;
 }
 
+/*
+ * Copy of heap->ptrs pointer is invalid after heap_grow.
+ */
 static
 int heap_grow(struct ptr_heap *heap, size_t new_len)
 {
@@ -118,69 +121,82 @@ static void heapify(struct ptr_heap *heap, size_t i)
 void *heap_replace_max(struct ptr_heap *heap, void *p)
 {
 	void *res;
-	void **ptrs = heap->ptrs;
 
 	if (!heap->len) {
 		(void) heap_set_len(heap, 1);
-		ptrs[0] = p;
+		heap->ptrs[0] = p;
 		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)
 		return ret;
+	ptrs = heap->ptrs;
 	/* Add the element to the end */
 	ptrs[heap->len - 1] = p;
-	/* rebalance */
-	heapify(heap, 0);
+	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);
+			/* rebalance */
+			heapify(heap, pos);
+		} else {
+			break;
+		}
+	}
 	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]);
+	return heap_replace_max(heap, heap->ptrs[heap->len - 1]);
 }
 
 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 (heap->ptrs[pos] == p)
 			goto found;
 	return NULL;
 found:
 	if (heap->len == 1) {
 		(void) heap_set_len(heap, 0);
-		return ptrs[0];
+		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];
+	heap->ptrs[pos] = heap->ptrs[heap->len - 1];
 	heapify(heap, pos);
 	return p;
 }
-- 
2.34.1