No need to rebalance heap for insertion

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

/*
 * prio_heap.c
 *
 * 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 "prio_heap.h"

static
size_t parent(size_t i)
{
	return i >> 1;
}

static
size_t left(size_t i)
{
	return i << 1;
}

static
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)
{
	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;
	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, max_t(size_t, 1, alloc_len));
}

void heap_free(struct ptr_heap *heap)
{
	kfree(heap->ptrs);
}

static void heapify(struct ptr_heap *heap, size_t i)
{
	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)
{
	void *res;

	if (!heap->len) {
		(void) heap_set_len(heap, 1);
		heap->ptrs[0] = p;
		return NULL;
	}

	/* Replace the current max and heapify */
	res = heap->ptrs[0];
	heap->ptrs[0] = p;
	heapify(heap, 0);
	return res;
}

int heap_insert(struct ptr_heap *heap, void *p)
{
	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;
	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;
		}
	}
	return 0;
}

void *heap_remove(struct ptr_heap *heap)
{
	switch (heap->len) {
	case 0:
		return NULL;
	case 1:
		(void) heap_set_len(heap, 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, heap->ptrs[heap->len - 1]);
}

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)
			goto found;
	return NULL;
found:
	if (heap->len == 1) {
		(void) heap_set_len(heap, 0);
		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];
	heapify(heap, pos);
	return p;
}