1 /* Copyright (c) 2008-2018 Dovecot authors, see the included COPYING file */
2 
3 #include "lib.h"
4 #include "array.h"
5 #include "priorityq.h"
6 
7 /* Macros for moving inside an array implementation of binary tree where
8    [0] is the root. */
9 #define PARENT_IDX(idx) \
10 	(((idx) - 1) / 2)
11 #define LEFT_CHILD_IDX(idx) \
12 	((idx) * 2 + 1)
13 #define RIGHT_CHILD_IDX(idx) \
14 	((idx) * 2 + 2)
15 
16 struct priorityq {
17 	priorityq_cmp_callback_t *cmp_callback;
18 
19 	ARRAY(struct priorityq_item *) items;
20 };
21 
22 struct priorityq *
priorityq_init(priorityq_cmp_callback_t * cmp_callback,unsigned int init_size)23 priorityq_init(priorityq_cmp_callback_t *cmp_callback, unsigned int init_size)
24 {
25 	struct priorityq *pq;
26 
27 	pq = i_new(struct priorityq, 1);
28 	pq->cmp_callback = cmp_callback;
29 	i_array_init(&pq->items, init_size);
30 	return pq;
31 }
32 
priorityq_deinit(struct priorityq ** _pq)33 void priorityq_deinit(struct priorityq **_pq)
34 {
35 	struct priorityq *pq = *_pq;
36 
37 	*_pq = NULL;
38 	array_free(&pq->items);
39 	i_free(pq);
40 }
41 
priorityq_count(const struct priorityq * pq)42 unsigned int priorityq_count(const struct priorityq *pq)
43 {
44 	return array_count(&pq->items);
45 }
46 
heap_items_swap(struct priorityq_item ** items,unsigned int idx1,unsigned int idx2)47 static void heap_items_swap(struct priorityq_item **items,
48 			    unsigned int idx1, unsigned int idx2)
49 {
50 	struct priorityq_item *tmp;
51 
52 	/* swap the item indexes */
53 	i_assert(items[idx1]->idx == idx1);
54 	i_assert(items[idx2]->idx == idx2);
55 
56 	items[idx1]->idx = idx2;
57 	items[idx2]->idx = idx1;
58 
59 	/* swap the item pointers */
60 	tmp = items[idx1];
61 	items[idx1] = items[idx2];
62 	items[idx2] = tmp;
63 }
64 
65 static unsigned int
heap_item_bubble_up(struct priorityq * pq,unsigned int idx)66 heap_item_bubble_up(struct priorityq *pq, unsigned int idx)
67 {
68 	struct priorityq_item **items;
69 	unsigned int parent_idx, count;
70 
71 	items = array_get_modifiable(&pq->items, &count);
72 	while (idx > 0) {
73 		parent_idx = PARENT_IDX(idx);
74 
75 		i_assert(idx < count);
76 		if (pq->cmp_callback(items[idx], items[parent_idx]) >= 0)
77 			break;
78 
79 		/* wrong order - swap */
80 		heap_items_swap(items, idx, parent_idx);
81 		idx = parent_idx;
82 	}
83 	return idx;
84 }
85 
heap_item_bubble_down(struct priorityq * pq,unsigned int idx)86 static void heap_item_bubble_down(struct priorityq *pq, unsigned int idx)
87 {
88 	struct priorityq_item **items;
89 	unsigned int left_idx, right_idx, min_child_idx, count;
90 
91 	items = array_get_modifiable(&pq->items, &count);
92 	while ((left_idx = LEFT_CHILD_IDX(idx)) < count) {
93 		right_idx = RIGHT_CHILD_IDX(idx);
94 		if (right_idx >= count ||
95 		    pq->cmp_callback(items[left_idx], items[right_idx]) < 0)
96 			min_child_idx = left_idx;
97 		else
98 			min_child_idx = right_idx;
99 
100 		if (pq->cmp_callback(items[min_child_idx], items[idx]) >= 0)
101 			break;
102 
103 		/* wrong order - swap */
104 		heap_items_swap(items, idx, min_child_idx);
105 		idx = min_child_idx;
106 	}
107 }
108 
priorityq_add(struct priorityq * pq,struct priorityq_item * item)109 void priorityq_add(struct priorityq *pq, struct priorityq_item *item)
110 {
111 	item->idx = array_count(&pq->items);
112 	array_push_back(&pq->items, &item);
113 	(void)heap_item_bubble_up(pq, item->idx);
114 }
115 
priorityq_remove_idx(struct priorityq * pq,unsigned int idx)116 static void priorityq_remove_idx(struct priorityq *pq, unsigned int idx)
117 {
118 	struct priorityq_item **items;
119 	unsigned int count;
120 
121 	items = array_get_modifiable(&pq->items, &count);
122 	i_assert(idx < count);
123 
124 	/* move last item over the removed one and fix the heap */
125 	count--;
126 	heap_items_swap(items, idx, count);
127 	array_delete(&pq->items, count, 1);
128 
129 	if (count > 0 && idx != count) {
130 		if (idx > 0)
131 			idx = heap_item_bubble_up(pq, idx);
132 		heap_item_bubble_down(pq, idx);
133 	}
134 }
135 
priorityq_remove(struct priorityq * pq,struct priorityq_item * item)136 void priorityq_remove(struct priorityq *pq, struct priorityq_item *item)
137 {
138 	priorityq_remove_idx(pq, item->idx);
139 	item->idx = UINT_MAX;
140 }
141 
priorityq_peek(struct priorityq * pq)142 struct priorityq_item *priorityq_peek(struct priorityq *pq)
143 {
144 	struct priorityq_item *const *items;
145 
146 	if (array_count(&pq->items) == 0)
147 		return NULL;
148 
149 	items = array_front(&pq->items);
150 	return items[0];
151 }
152 
priorityq_pop(struct priorityq * pq)153 struct priorityq_item *priorityq_pop(struct priorityq *pq)
154 {
155 	struct priorityq_item *item;
156 
157 	item = priorityq_peek(pq);
158 	if (item != NULL) {
159 		priorityq_remove_idx(pq, 0);
160 		item->idx = UINT_MAX;
161 	}
162 	return item;
163 }
164 
priorityq_items(struct priorityq * pq)165 struct priorityq_item *const *priorityq_items(struct priorityq *pq)
166 {
167 	if (array_count(&pq->items) == 0)
168 		return NULL;
169 
170 	return array_front(&pq->items);
171 }
172