5658c1fd1037aa5546ae90f22be4e35c89d6bf65
[blender.git] / source / blender / blenlib / intern / BLI_heap.c
1 /*
2  * ***** BEGIN GPL LICENSE BLOCK *****
3  *
4  * This program is free software; you can redistribute it and/or
5  * modify it under the terms of the GNU General Public License
6  * as published by the Free Software Foundation; either version 2
7  * of the License, or (at your option) any later version.
8  *
9  * This program is distributed in the hope that it will be useful,
10  * but WITHOUT ANY WARRANTY; without even the implied warranty of
11  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12  * GNU General Public License for more details.
13  *
14  * You should have received a copy of the GNU General Public License
15  * along with this program; if not, write to the Free Software Foundation,
16  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
17  *
18  * Contributor(s): Brecht Van Lommel
19  *
20  * ***** END GPL LICENSE BLOCK *****
21  */
22
23 /** \file blender/blenlib/intern/BLI_heap.c
24  *  \ingroup bli
25  *
26  * A min-heap / priority queue ADT.
27  */
28
29 #include <stdlib.h>
30 #include <string.h>
31
32 #include "MEM_guardedalloc.h"
33
34 #include "BLI_utildefines.h"
35 #include "BLI_heap.h"
36 #include "BLI_strict_flags.h"
37
38 /***/
39
40 struct HeapNode {
41         void *ptr;
42         float value;
43         uint  index;
44 };
45
46 struct HeapNode_Chunk {
47         struct HeapNode_Chunk *prev;
48         uint size;
49         uint bufsize;
50         struct HeapNode buf[0];
51 };
52
53 /**
54  * Number of nodes to include per #HeapNode_Chunk when no reserved size is passed,
55  * or we allocate past the reserved number.
56  *
57  * \note Optimize number for 64kb allocs.
58  * \note keep type in sync with tot_nodes in heap_node_alloc_chunk.
59  */
60 #define HEAP_CHUNK_DEFAULT_NUM \
61         ((uint)((MEM_SIZE_OPTIMAL((1 << 16) - sizeof(struct HeapNode_Chunk))) / sizeof(HeapNode)))
62
63 struct Heap {
64         uint size;
65         uint bufsize;
66         HeapNode **tree;
67
68         struct {
69                 /* Always keep at least one chunk (never NULL) */
70                 struct HeapNode_Chunk *chunk;
71                 /* when NULL, allocate a new chunk */
72                 HeapNode *free;
73         } nodes;
74 };
75
76 /** \name Internal Functions
77  * \{ */
78
79 #define HEAP_PARENT(i) (((i) - 1) >> 1)
80 #define HEAP_LEFT(i)   (((i) << 1) + 1)
81 #define HEAP_RIGHT(i)  (((i) << 1) + 2)
82 #define HEAP_COMPARE(a, b) ((a)->value < (b)->value)
83
84 #if 0  /* UNUSED */
85 #define HEAP_EQUALS(a, b) ((a)->value == (b)->value)
86 #endif
87
88 BLI_INLINE void heap_swap(Heap *heap, const uint i, const uint j)
89 {
90
91 #if 0
92         SWAP(uint,  heap->tree[i]->index, heap->tree[j]->index);
93         SWAP(HeapNode *,    heap->tree[i],        heap->tree[j]);
94 #else
95         HeapNode **tree = heap->tree;
96         union {
97                 uint  index;
98                 HeapNode     *node;
99         } tmp;
100         SWAP_TVAL(tmp.index, tree[i]->index, tree[j]->index);
101         SWAP_TVAL(tmp.node,  tree[i],        tree[j]);
102 #endif
103 }
104
105 static void heap_down(Heap *heap, uint i)
106 {
107         /* size won't change in the loop */
108         const uint size = heap->size;
109
110         while (1) {
111                 const uint l = HEAP_LEFT(i);
112                 const uint r = HEAP_RIGHT(i);
113                 uint smallest = i;
114
115                 if ((l < size) && HEAP_COMPARE(heap->tree[l], heap->tree[smallest])) {
116                         smallest = l;
117                 }
118                 if ((r < size) && HEAP_COMPARE(heap->tree[r], heap->tree[smallest])) {
119                         smallest = r;
120                 }
121
122                 if (smallest == i) {
123                         break;
124                 }
125
126                 heap_swap(heap, i, smallest);
127                 i = smallest;
128         }
129 }
130
131 static void heap_up(Heap *heap, uint i)
132 {
133         while (i > 0) {
134                 const uint p = HEAP_PARENT(i);
135
136                 if (HEAP_COMPARE(heap->tree[p], heap->tree[i])) {
137                         break;
138                 }
139                 heap_swap(heap, p, i);
140                 i = p;
141         }
142 }
143
144 /** \} */
145
146
147 /** \name Internal Memory Management
148  * \{ */
149
150 static struct HeapNode_Chunk *heap_node_alloc_chunk(
151         uint tot_nodes, struct HeapNode_Chunk *chunk_prev)
152 {
153         struct HeapNode_Chunk *chunk = MEM_mallocN(
154                 sizeof(struct HeapNode_Chunk) + (sizeof(HeapNode) * tot_nodes), __func__);
155         chunk->prev = chunk_prev;
156         chunk->bufsize = tot_nodes;
157         chunk->size = 0;
158         return chunk;
159 }
160
161 static struct HeapNode *heap_node_alloc(Heap *heap)
162 {
163         HeapNode *node;
164
165         if (heap->nodes.free) {
166                 node = heap->nodes.free;
167                 heap->nodes.free = heap->nodes.free->ptr;
168         }
169         else {
170                 struct HeapNode_Chunk *chunk = heap->nodes.chunk;
171                 if (UNLIKELY(chunk->size == chunk->bufsize)) {
172                         chunk = heap->nodes.chunk = heap_node_alloc_chunk(HEAP_CHUNK_DEFAULT_NUM, chunk);
173                 }
174                 node = &chunk->buf[chunk->size++];
175         }
176
177         return node;
178 }
179
180 static void heap_node_free(Heap *heap, HeapNode *node)
181 {
182         node->ptr = heap->nodes.free;
183         heap->nodes.free = node;
184 }
185
186 /** \} */
187
188
189 /** \name Public Heap API
190  * \{ */
191
192 /**
193  * Creates a new heap. Removed nodes are recycled, so memory usage will not shrink.
194  *
195  * \note Use when the size of the heap is known in advance.
196  */
197 Heap *BLI_heap_new_ex(uint tot_reserve)
198 {
199         Heap *heap = MEM_mallocN(sizeof(Heap), __func__);
200         /* ensure we have at least one so we can keep doubling it */
201         heap->size = 0;
202         heap->bufsize = MAX2(1u, tot_reserve);
203         heap->tree = MEM_mallocN(heap->bufsize * sizeof(HeapNode *), "BLIHeapTree");
204
205         heap->nodes.chunk = heap_node_alloc_chunk((tot_reserve > 1) ? tot_reserve : HEAP_CHUNK_DEFAULT_NUM, NULL);
206         heap->nodes.free = NULL;
207
208         return heap;
209 }
210
211 Heap *BLI_heap_new(void)
212 {
213         return BLI_heap_new_ex(1);
214 }
215
216 void BLI_heap_free(Heap *heap, HeapFreeFP ptrfreefp)
217 {
218         if (ptrfreefp) {
219                 uint i;
220
221                 for (i = 0; i < heap->size; i++) {
222                         ptrfreefp(heap->tree[i]->ptr);
223                 }
224         }
225
226         struct HeapNode_Chunk *chunk = heap->nodes.chunk;
227         do {
228                 struct HeapNode_Chunk *chunk_prev;
229                 chunk_prev = chunk->prev;
230                 MEM_freeN(chunk);
231                 chunk = chunk_prev;
232         } while (chunk);
233
234         MEM_freeN(heap->tree);
235         MEM_freeN(heap);
236 }
237
238 void BLI_heap_clear(Heap *heap, HeapFreeFP ptrfreefp)
239 {
240         if (ptrfreefp) {
241                 uint i;
242
243                 for (i = 0; i < heap->size; i++) {
244                         ptrfreefp(heap->tree[i]->ptr);
245                 }
246         }
247         heap->size = 0;
248
249         /* Remove all except the last chunk */
250         while (heap->nodes.chunk->prev) {
251                 struct HeapNode_Chunk *chunk_prev = heap->nodes.chunk->prev;
252                 MEM_freeN(heap->nodes.chunk);
253                 heap->nodes.chunk = chunk_prev;
254         }
255         heap->nodes.chunk->size = 0;
256         heap->nodes.free = NULL;
257 }
258
259 /**
260  * Insert heap node with a value (often a 'cost') and pointer into the heap,
261  * duplicate values are allowed.
262  */
263 HeapNode *BLI_heap_insert(Heap *heap, float value, void *ptr)
264 {
265         HeapNode *node;
266
267         if (UNLIKELY(heap->size >= heap->bufsize)) {
268                 heap->bufsize *= 2;
269                 heap->tree = MEM_reallocN(heap->tree, heap->bufsize * sizeof(*heap->tree));
270         }
271
272         node = heap_node_alloc(heap);
273
274         node->ptr = ptr;
275         node->value = value;
276         node->index = heap->size;
277
278         heap->tree[node->index] = node;
279
280         heap->size++;
281
282         heap_up(heap, node->index);
283
284         return node;
285 }
286
287 /**
288  * Convenience function since this is a common pattern.
289  */
290 void BLI_heap_insert_or_update(Heap *heap, HeapNode **node_p, float value, void *ptr)
291 {
292         if (*node_p == NULL) {
293                 *node_p = BLI_heap_insert(heap, value, ptr);
294         }
295         else {
296                 BLI_heap_node_value_update_ptr(heap, *node_p, value, ptr);
297         }
298 }
299
300
301 bool BLI_heap_is_empty(const Heap *heap)
302 {
303         return (heap->size == 0);
304 }
305
306 uint BLI_heap_len(const Heap *heap)
307 {
308         return heap->size;
309 }
310
311 /**
312  * Return the top node of the heap.
313  * This is the node with the lowest value.
314  */
315 HeapNode *BLI_heap_top(const Heap *heap)
316 {
317         return heap->tree[0];
318 }
319
320 /**
321  * Pop the top node off the heap and return it's pointer.
322  */
323 void *BLI_heap_pop_min(Heap *heap)
324 {
325         BLI_assert(heap->size != 0);
326
327         void *ptr = heap->tree[0]->ptr;
328
329         heap_node_free(heap, heap->tree[0]);
330
331         if (--heap->size) {
332                 heap_swap(heap, 0, heap->size);
333                 heap_down(heap, 0);
334         }
335
336         return ptr;
337 }
338
339 void BLI_heap_remove(Heap *heap, HeapNode *node)
340 {
341         BLI_assert(heap->size != 0);
342
343         uint i = node->index;
344
345         while (i > 0) {
346                 uint p = HEAP_PARENT(i);
347                 heap_swap(heap, p, i);
348                 i = p;
349         }
350
351         BLI_heap_pop_min(heap);
352 }
353
354 /**
355  * Can be used to avoid #BLI_heap_remove, #BLI_heap_insert calls,
356  * balancing the tree still has a performance cost,
357  * but is often much less than remove/insert, difference is most noticable with large heaps.
358  */
359 void BLI_heap_node_value_update(Heap *heap, HeapNode *node, float value)
360 {
361         if (value < node->value) {
362                 node->value = value;
363                 heap_up(heap, node->index);
364         }
365         else if (value > node->value) {
366                 node->value = value;
367                 heap_down(heap, node->index);
368         }
369 }
370
371 void BLI_heap_node_value_update_ptr(Heap *heap, HeapNode *node, float value, void *ptr)
372 {
373         node->ptr = ptr; /* only difference */
374         if (value < node->value) {
375                 node->value = value;
376                 heap_up(heap, node->index);
377         }
378         else if (value > node->value) {
379                 node->value = value;
380                 heap_down(heap, node->index);
381         }
382 }
383
384 float BLI_heap_node_value(const HeapNode *node)
385 {
386         return node->value;
387 }
388
389 void *BLI_heap_node_ptr(const HeapNode *node)
390 {
391         return node->ptr;
392 }
393
394 static bool heap_is_minheap(const Heap *heap, uint root)
395 {
396         if (root < heap->size) {
397                 const uint l = HEAP_LEFT(root);
398                 if (l < heap->size) {
399                         if (HEAP_COMPARE(heap->tree[l], heap->tree[root]) || !heap_is_minheap(heap, l)) {
400                                 return false;
401                         }
402                 }
403                 const uint r = HEAP_RIGHT(root);
404                 if (r < heap->size) {
405                         if (HEAP_COMPARE(heap->tree[r], heap->tree[root]) || !heap_is_minheap(heap, r)) {
406                                 return false;
407                         }
408                 }
409         }
410         return true;
411 }
412 /**
413  * Only for checking internal errors (gtest).
414  */
415 bool BLI_heap_is_valid(const Heap *heap)
416 {
417         return heap_is_minheap(heap, 0);
418 }
419
420 /** \} */