diff options
Diffstat (limited to 'src/libutil/heap.c')
-rw-r--r-- | src/libutil/heap.c | 197 |
1 files changed, 197 insertions, 0 deletions
diff --git a/src/libutil/heap.c b/src/libutil/heap.c new file mode 100644 index 0000000..8ce70cf --- /dev/null +++ b/src/libutil/heap.c @@ -0,0 +1,197 @@ +/*- + * 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, guint npri) +{ + guint 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, guint idx) +{ + g_assert(heap != NULL); + g_assert(idx < heap->ar->len); + + return g_ptr_array_index(heap->ar, idx); +} |