rspamd/src/libutil/heap.c

/*-
 * Copyright 2016 Vsevolod Stakhov
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "config.h"
#include "libutil/heap.h"

struct rspamd_min_heap {
	GPtrArray *ar;
};

#define __SWAP(a, b)            \
	do {                        \
		__typeof__(a) _a = (a); \
		__typeof__(b) _b = (b); \
		a = _b;                 \
		b = _a;                 \
	} while (0)
#define heap_swap(h, e1, e2)                                                  \
	do {                                                                      \
		__SWAP((h)->ar->pdata[(e1)->idx - 1], (h)->ar->pdata[(e2)->idx - 1]); \
		__SWAP((e1)->idx, (e2)->idx);                                         \
	} while (0)

#define min_elt(e1, e2) ((e1)->pri <= (e2)->pri ? (e1) : (e2))

/*
 * Swims element added (or changed) to preserve heap's invariant
 */
static void
rspamd_min_heap_swim(struct rspamd_min_heap *heap,
					 struct rspamd_min_heap_elt *elt)
{
	struct rspamd_min_heap_elt *parent;

	while (elt->idx > 1) {
		parent = g_ptr_array_index(heap->ar, elt->idx / 2 - 1);

		if (parent->pri > elt->pri) {
			heap_swap(heap, elt, parent);
		}
		else {
			break;
		}
	}
}

/*
 * Sinks the element popped (or changed) to preserve heap's invariant
 */
static void
rspamd_min_heap_sink(struct rspamd_min_heap *heap,
					 struct rspamd_min_heap_elt *elt)
{
	struct rspamd_min_heap_elt *c1, *c2, *m;

	while (elt->idx * 2 < heap->ar->len) {
		c1 = g_ptr_array_index(heap->ar, elt->idx * 2 - 1);
		c2 = g_ptr_array_index(heap->ar, elt->idx * 2);
		m = min_elt(c1, c2);

		if (elt->pri > m->pri) {
			heap_swap(heap, elt, m);
		}
		else {
			break;
		}
	}

	if (elt->idx * 2 - 1 < heap->ar->len) {
		m = g_ptr_array_index(heap->ar, elt->idx * 2 - 1);
		if (elt->pri > m->pri) {
			heap_swap(heap, elt, m);
		}
	}
}

struct rspamd_min_heap *
rspamd_min_heap_create(gsize reserved_size)
{
	struct rspamd_min_heap *heap;

	heap = g_malloc(sizeof(*heap));
	heap->ar = g_ptr_array_sized_new(reserved_size);

	return heap;
}

void rspamd_min_heap_push(struct rspamd_min_heap *heap,
						  struct rspamd_min_heap_elt *elt)
{
	g_assert(heap != NULL);
	g_assert(elt != NULL);

	/* Add to the end */
	elt->idx = heap->ar->len + 1;
	g_ptr_array_add(heap->ar, elt);
	/* Now swim it up */
	rspamd_min_heap_swim(heap, elt);
}

struct rspamd_min_heap_elt *
rspamd_min_heap_pop(struct rspamd_min_heap *heap)
{
	struct rspamd_min_heap_elt *elt, *last;

	g_assert(heap != NULL);

	if (heap->ar->len == 0) {
		return NULL;
	}

	elt = g_ptr_array_index(heap->ar, 0);
	last = g_ptr_array_index(heap->ar, heap->ar->len - 1);

	if (elt != last) {
		/* Now replace elt with the last element and sink it if needed */
		heap_swap(heap, elt, last);
		g_ptr_array_remove_index_fast(heap->ar, heap->ar->len - 1);
		rspamd_min_heap_sink(heap, last);
	}
	else {
		g_ptr_array_remove_index_fast(heap->ar, heap->ar->len - 1);
	}


	return elt;
}

void rspamd_min_heap_update_elt(struct rspamd_min_heap *heap,
								struct rspamd_min_heap_elt *elt, unsigned int npri)
{
	unsigned int oldpri;

	g_assert(heap != NULL);
	g_assert(elt->idx > 0 && elt->idx <= heap->ar->len);

	oldpri = elt->pri;
	elt->pri = npri;

	if (npri > oldpri) {
		/* We might need to sink */
		rspamd_min_heap_sink(heap, elt);
	}
	else if (npri < oldpri) {
		/* We might need to swim */
		rspamd_min_heap_swim(heap, elt);
	}
}

void rspamd_min_heap_remove_elt(struct rspamd_min_heap *heap,
								struct rspamd_min_heap_elt *elt)
{
	struct rspamd_min_heap_elt *first;

	g_assert(heap != NULL);
	g_assert(elt->idx > 0 && elt->idx <= heap->ar->len);

	first = g_ptr_array_index(heap->ar, 0);

	if (elt != first) {
		elt->pri = first->pri - 1;
		rspamd_min_heap_swim(heap, elt);
	}

	/* Now the desired element is on the top of queue */
	(void) rspamd_min_heap_pop(heap);
}

void rspamd_min_heap_destroy(struct rspamd_min_heap *heap)
{
	if (heap) {
		g_ptr_array_free(heap->ar, TRUE);
		g_free(heap);
	}
}

struct rspamd_min_heap_elt *
rspamd_min_heap_index(struct rspamd_min_heap *heap, unsigned int idx)
{
	g_assert(heap != NULL);
	g_assert(idx < heap->ar->len);

	return g_ptr_array_index(heap->ar, idx);
}