1 /**
2  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3  * SPDX-License-Identifier: Apache-2.0.
4  */
5 
6 #include <aws/common/priority_queue.h>
7 
8 #include <string.h>
9 
10 #define PARENT_OF(index) (((index)&1) ? (index) >> 1 : (index) > 1 ? ((index)-2) >> 1 : 0)
11 #define LEFT_OF(index) (((index) << 1) + 1)
12 #define RIGHT_OF(index) (((index) << 1) + 2)
13 
s_swap(struct aws_priority_queue * queue,size_t a,size_t b)14 static void s_swap(struct aws_priority_queue *queue, size_t a, size_t b) {
15     AWS_PRECONDITION(aws_priority_queue_is_valid(queue));
16     AWS_PRECONDITION(a < queue->container.length);
17     AWS_PRECONDITION(b < queue->container.length);
18     AWS_PRECONDITION(aws_priority_queue_backpointer_index_valid(queue, a));
19     AWS_PRECONDITION(aws_priority_queue_backpointer_index_valid(queue, b));
20 
21     aws_array_list_swap(&queue->container, a, b);
22 
23     /* Invariant: If the backpointer array is initialized, we have enough room for all elements */
24     if (!AWS_IS_ZEROED(queue->backpointers)) {
25         AWS_ASSERT(queue->backpointers.length > a);
26         AWS_ASSERT(queue->backpointers.length > b);
27 
28         struct aws_priority_queue_node **bp_a = &((struct aws_priority_queue_node **)queue->backpointers.data)[a];
29         struct aws_priority_queue_node **bp_b = &((struct aws_priority_queue_node **)queue->backpointers.data)[b];
30 
31         struct aws_priority_queue_node *tmp = *bp_a;
32         *bp_a = *bp_b;
33         *bp_b = tmp;
34 
35         if (*bp_a) {
36             (*bp_a)->current_index = a;
37         }
38 
39         if (*bp_b) {
40             (*bp_b)->current_index = b;
41         }
42     }
43     AWS_POSTCONDITION(aws_priority_queue_is_valid(queue));
44     AWS_POSTCONDITION(aws_priority_queue_backpointer_index_valid(queue, a));
45     AWS_POSTCONDITION(aws_priority_queue_backpointer_index_valid(queue, b));
46 }
47 
48 /* Precondition: with the exception of the given root element, the container must be
49  * in heap order */
s_sift_down(struct aws_priority_queue * queue,size_t root)50 static bool s_sift_down(struct aws_priority_queue *queue, size_t root) {
51     AWS_PRECONDITION(aws_priority_queue_is_valid(queue));
52     AWS_PRECONDITION(root < queue->container.length);
53 
54     bool did_move = false;
55 
56     size_t len = aws_array_list_length(&queue->container);
57 
58     while (LEFT_OF(root) < len) {
59         size_t left = LEFT_OF(root);
60         size_t right = RIGHT_OF(root);
61         size_t first = root;
62         void *first_item = NULL, *other_item = NULL;
63 
64         aws_array_list_get_at_ptr(&queue->container, &first_item, root);
65         aws_array_list_get_at_ptr(&queue->container, &other_item, left);
66 
67         if (queue->pred(first_item, other_item) > 0) {
68             first = left;
69             first_item = other_item;
70         }
71 
72         if (right < len) {
73             aws_array_list_get_at_ptr(&queue->container, &other_item, right);
74 
75             /* choose the larger/smaller of the two in case of a max/min heap
76              * respectively */
77             if (queue->pred(first_item, other_item) > 0) {
78                 first = right;
79                 first_item = other_item;
80             }
81         }
82 
83         if (first != root) {
84             s_swap(queue, first, root);
85             did_move = true;
86             root = first;
87         } else {
88             break;
89         }
90     }
91 
92     AWS_POSTCONDITION(aws_priority_queue_is_valid(queue));
93     return did_move;
94 }
95 
96 /* Precondition: Elements prior to the specified index must be in heap order. */
s_sift_up(struct aws_priority_queue * queue,size_t index)97 static bool s_sift_up(struct aws_priority_queue *queue, size_t index) {
98     AWS_PRECONDITION(aws_priority_queue_is_valid(queue));
99     AWS_PRECONDITION(index < queue->container.length);
100 
101     bool did_move = false;
102 
103     void *parent_item = NULL, *child_item = NULL;
104     size_t parent = PARENT_OF(index);
105     while (index) {
106         /*
107          * These get_ats are guaranteed to be successful; if they are not, we have
108          * serious state corruption, so just abort.
109          */
110 
111         if (aws_array_list_get_at_ptr(&queue->container, &parent_item, parent) ||
112             aws_array_list_get_at_ptr(&queue->container, &child_item, index)) {
113             abort();
114         }
115 
116         if (queue->pred(parent_item, child_item) > 0) {
117             s_swap(queue, index, parent);
118             did_move = true;
119             index = parent;
120             parent = PARENT_OF(index);
121         } else {
122             break;
123         }
124     }
125 
126     AWS_POSTCONDITION(aws_priority_queue_is_valid(queue));
127     return did_move;
128 }
129 
130 /*
131  * Precondition: With the exception of the given index, the heap condition holds for all elements.
132  * In particular, the parent of the current index is a predecessor of all children of the current index.
133  */
s_sift_either(struct aws_priority_queue * queue,size_t index)134 static void s_sift_either(struct aws_priority_queue *queue, size_t index) {
135     AWS_PRECONDITION(aws_priority_queue_is_valid(queue));
136     AWS_PRECONDITION(index < queue->container.length);
137 
138     if (!index || !s_sift_up(queue, index)) {
139         s_sift_down(queue, index);
140     }
141 
142     AWS_POSTCONDITION(aws_priority_queue_is_valid(queue));
143 }
144 
aws_priority_queue_init_dynamic(struct aws_priority_queue * queue,struct aws_allocator * alloc,size_t default_size,size_t item_size,aws_priority_queue_compare_fn * pred)145 int aws_priority_queue_init_dynamic(
146     struct aws_priority_queue *queue,
147     struct aws_allocator *alloc,
148     size_t default_size,
149     size_t item_size,
150     aws_priority_queue_compare_fn *pred) {
151 
152     AWS_FATAL_PRECONDITION(queue != NULL);
153     AWS_FATAL_PRECONDITION(alloc != NULL);
154     AWS_FATAL_PRECONDITION(item_size > 0);
155 
156     queue->pred = pred;
157     AWS_ZERO_STRUCT(queue->backpointers);
158 
159     int ret = aws_array_list_init_dynamic(&queue->container, alloc, default_size, item_size);
160     if (ret == AWS_OP_SUCCESS) {
161         AWS_POSTCONDITION(aws_priority_queue_is_valid(queue));
162     } else {
163         AWS_POSTCONDITION(AWS_IS_ZEROED(queue->container));
164         AWS_POSTCONDITION(AWS_IS_ZEROED(queue->backpointers));
165     }
166     return ret;
167 }
168 
aws_priority_queue_init_static(struct aws_priority_queue * queue,void * heap,size_t item_count,size_t item_size,aws_priority_queue_compare_fn * pred)169 void aws_priority_queue_init_static(
170     struct aws_priority_queue *queue,
171     void *heap,
172     size_t item_count,
173     size_t item_size,
174     aws_priority_queue_compare_fn *pred) {
175 
176     AWS_FATAL_PRECONDITION(queue != NULL);
177     AWS_FATAL_PRECONDITION(heap != NULL);
178     AWS_FATAL_PRECONDITION(item_count > 0);
179     AWS_FATAL_PRECONDITION(item_size > 0);
180 
181     queue->pred = pred;
182     AWS_ZERO_STRUCT(queue->backpointers);
183 
184     aws_array_list_init_static(&queue->container, heap, item_count, item_size);
185 
186     AWS_POSTCONDITION(aws_priority_queue_is_valid(queue));
187 }
188 
aws_priority_queue_backpointer_index_valid(const struct aws_priority_queue * const queue,size_t index)189 bool aws_priority_queue_backpointer_index_valid(const struct aws_priority_queue *const queue, size_t index) {
190     if (AWS_IS_ZEROED(queue->backpointers)) {
191         return true;
192     }
193     if (index < queue->backpointers.length) {
194         struct aws_priority_queue_node *node = ((struct aws_priority_queue_node **)queue->backpointers.data)[index];
195         return (node == NULL) || AWS_MEM_IS_WRITABLE(node, sizeof(struct aws_priority_queue_node));
196     }
197     return false;
198 }
199 
aws_priority_queue_backpointers_valid_deep(const struct aws_priority_queue * const queue)200 bool aws_priority_queue_backpointers_valid_deep(const struct aws_priority_queue *const queue) {
201     if (!queue) {
202         return false;
203     }
204     for (size_t i = 0; i < queue->backpointers.length; i++) {
205         if (!aws_priority_queue_backpointer_index_valid(queue, i)) {
206             return false;
207         }
208     }
209     return true;
210 }
211 
aws_priority_queue_backpointers_valid(const struct aws_priority_queue * const queue)212 bool aws_priority_queue_backpointers_valid(const struct aws_priority_queue *const queue) {
213     if (!queue) {
214         return false;
215     }
216 
217     /* Internal container validity */
218     bool backpointer_list_is_valid =
219         ((aws_array_list_is_valid(&queue->backpointers) && (queue->backpointers.current_size != 0) &&
220           (queue->backpointers.data != NULL)));
221 
222     /* Backpointer struct should either be zero or should be
223      * initialized to be at most as long as the container, and having
224      * as elements potentially null pointers to
225      * aws_priority_queue_nodes */
226     bool backpointer_list_item_size = queue->backpointers.item_size == sizeof(struct aws_priority_queue_node *);
227     bool lists_equal_lengths = queue->backpointers.length == queue->container.length;
228     bool backpointers_non_zero_current_size = queue->backpointers.current_size > 0;
229 
230     /* This check must be guarded, as it is not efficient, neither
231      * when running tests nor CBMC */
232 #if (AWS_DEEP_CHECKS == 1)
233     bool backpointers_valid_deep = aws_priority_queue_backpointers_valid_deep(queue);
234 #else
235     bool backpointers_valid_deep = true;
236 #endif
237     bool backpointers_zero =
238         (queue->backpointers.current_size == 0 && queue->backpointers.length == 0 && queue->backpointers.data == NULL);
239     bool backpointer_struct_is_valid =
240         backpointers_zero || (backpointer_list_item_size && lists_equal_lengths && backpointers_non_zero_current_size &&
241                               backpointers_valid_deep);
242 
243     return ((backpointer_list_is_valid && backpointer_struct_is_valid) || AWS_IS_ZEROED(queue->backpointers));
244 }
245 
aws_priority_queue_is_valid(const struct aws_priority_queue * const queue)246 bool aws_priority_queue_is_valid(const struct aws_priority_queue *const queue) {
247     /* Pointer validity checks */
248     if (!queue) {
249         return false;
250     }
251     bool pred_is_valid = (queue->pred != NULL);
252     bool container_is_valid = aws_array_list_is_valid(&queue->container);
253 
254     bool backpointers_valid = aws_priority_queue_backpointers_valid(queue);
255     return pred_is_valid && container_is_valid && backpointers_valid;
256 }
257 
aws_priority_queue_clean_up(struct aws_priority_queue * queue)258 void aws_priority_queue_clean_up(struct aws_priority_queue *queue) {
259     aws_array_list_clean_up(&queue->container);
260     if (!AWS_IS_ZEROED(queue->backpointers)) {
261         aws_array_list_clean_up(&queue->backpointers);
262     }
263 }
264 
aws_priority_queue_push(struct aws_priority_queue * queue,void * item)265 int aws_priority_queue_push(struct aws_priority_queue *queue, void *item) {
266     AWS_PRECONDITION(aws_priority_queue_is_valid(queue));
267     AWS_PRECONDITION(item && AWS_MEM_IS_READABLE(item, queue->container.item_size));
268     int rval = aws_priority_queue_push_ref(queue, item, NULL);
269     AWS_POSTCONDITION(aws_priority_queue_is_valid(queue));
270     return rval;
271 }
272 
aws_priority_queue_push_ref(struct aws_priority_queue * queue,void * item,struct aws_priority_queue_node * backpointer)273 int aws_priority_queue_push_ref(
274     struct aws_priority_queue *queue,
275     void *item,
276     struct aws_priority_queue_node *backpointer) {
277     AWS_PRECONDITION(aws_priority_queue_is_valid(queue));
278     AWS_PRECONDITION(item && AWS_MEM_IS_READABLE(item, queue->container.item_size));
279 
280     int err = aws_array_list_push_back(&queue->container, item);
281     if (err) {
282         AWS_POSTCONDITION(aws_priority_queue_is_valid(queue));
283         return err;
284     }
285     size_t index = aws_array_list_length(&queue->container) - 1;
286 
287     if (backpointer && !queue->backpointers.alloc) {
288         if (!queue->container.alloc) {
289             aws_raise_error(AWS_ERROR_UNSUPPORTED_OPERATION);
290             goto backpointer_update_failed;
291         }
292 
293         if (aws_array_list_init_dynamic(
294                 &queue->backpointers, queue->container.alloc, index + 1, sizeof(struct aws_priority_queue_node *))) {
295             goto backpointer_update_failed;
296         }
297 
298         /* When we initialize the backpointers array we need to zero out all existing entries */
299         memset(queue->backpointers.data, 0, queue->backpointers.current_size);
300     }
301 
302     /*
303      * Once we have any backpointers, we want to make sure we always have room in the backpointers array
304      * for all elements; otherwise, sift_down gets complicated if it runs out of memory when sifting an
305      * element with a backpointer down in the array.
306      */
307     if (!AWS_IS_ZEROED(queue->backpointers)) {
308         if (aws_array_list_set_at(&queue->backpointers, &backpointer, index)) {
309             goto backpointer_update_failed;
310         }
311     }
312 
313     if (backpointer) {
314         backpointer->current_index = index;
315     }
316 
317     s_sift_up(queue, aws_array_list_length(&queue->container) - 1);
318 
319     AWS_POSTCONDITION(aws_priority_queue_is_valid(queue));
320     return AWS_OP_SUCCESS;
321 
322 backpointer_update_failed:
323     /* Failed to initialize or grow the backpointer array, back out the node addition */
324     aws_array_list_pop_back(&queue->container);
325     AWS_POSTCONDITION(aws_priority_queue_is_valid(queue));
326     return AWS_OP_ERR;
327 }
328 
s_remove_node(struct aws_priority_queue * queue,void * item,size_t item_index)329 static int s_remove_node(struct aws_priority_queue *queue, void *item, size_t item_index) {
330     AWS_PRECONDITION(aws_priority_queue_is_valid(queue));
331     AWS_PRECONDITION(item && AWS_MEM_IS_WRITABLE(item, queue->container.item_size));
332     if (aws_array_list_get_at(&queue->container, item, item_index)) {
333         /* shouldn't happen, but if it does we've already raised an error... */
334         AWS_POSTCONDITION(aws_priority_queue_is_valid(queue));
335         return AWS_OP_ERR;
336     }
337 
338     size_t swap_with = aws_array_list_length(&queue->container) - 1;
339     struct aws_priority_queue_node *backpointer = NULL;
340 
341     if (item_index != swap_with) {
342         s_swap(queue, item_index, swap_with);
343     }
344 
345     aws_array_list_pop_back(&queue->container);
346 
347     if (!AWS_IS_ZEROED(queue->backpointers)) {
348         aws_array_list_get_at(&queue->backpointers, &backpointer, swap_with);
349         if (backpointer) {
350             backpointer->current_index = SIZE_MAX;
351         }
352         aws_array_list_pop_back(&queue->backpointers);
353     }
354 
355     if (item_index != swap_with) {
356         s_sift_either(queue, item_index);
357     }
358 
359     AWS_POSTCONDITION(aws_priority_queue_is_valid(queue));
360     return AWS_OP_SUCCESS;
361 }
362 
aws_priority_queue_remove(struct aws_priority_queue * queue,void * item,const struct aws_priority_queue_node * node)363 int aws_priority_queue_remove(
364     struct aws_priority_queue *queue,
365     void *item,
366     const struct aws_priority_queue_node *node) {
367     AWS_PRECONDITION(aws_priority_queue_is_valid(queue));
368     AWS_PRECONDITION(item && AWS_MEM_IS_WRITABLE(item, queue->container.item_size));
369     AWS_PRECONDITION(node && AWS_MEM_IS_READABLE(node, sizeof(struct aws_priority_queue_node)));
370     AWS_ERROR_PRECONDITION(
371         node->current_index < aws_array_list_length(&queue->container), AWS_ERROR_PRIORITY_QUEUE_BAD_NODE);
372     AWS_ERROR_PRECONDITION(queue->backpointers.data, AWS_ERROR_PRIORITY_QUEUE_BAD_NODE);
373 
374     int rval = s_remove_node(queue, item, node->current_index);
375     AWS_POSTCONDITION(aws_priority_queue_is_valid(queue));
376     return rval;
377 }
378 
aws_priority_queue_pop(struct aws_priority_queue * queue,void * item)379 int aws_priority_queue_pop(struct aws_priority_queue *queue, void *item) {
380     AWS_PRECONDITION(aws_priority_queue_is_valid(queue));
381     AWS_PRECONDITION(item && AWS_MEM_IS_WRITABLE(item, queue->container.item_size));
382     AWS_ERROR_PRECONDITION(aws_array_list_length(&queue->container) != 0, AWS_ERROR_PRIORITY_QUEUE_EMPTY);
383 
384     int rval = s_remove_node(queue, item, 0);
385     AWS_POSTCONDITION(aws_priority_queue_is_valid(queue));
386     return rval;
387 }
388 
aws_priority_queue_top(const struct aws_priority_queue * queue,void ** item)389 int aws_priority_queue_top(const struct aws_priority_queue *queue, void **item) {
390     AWS_ERROR_PRECONDITION(aws_array_list_length(&queue->container) != 0, AWS_ERROR_PRIORITY_QUEUE_EMPTY);
391     return aws_array_list_get_at_ptr(&queue->container, item, 0);
392 }
393 
aws_priority_queue_size(const struct aws_priority_queue * queue)394 size_t aws_priority_queue_size(const struct aws_priority_queue *queue) {
395     return aws_array_list_length(&queue->container);
396 }
397 
aws_priority_queue_capacity(const struct aws_priority_queue * queue)398 size_t aws_priority_queue_capacity(const struct aws_priority_queue *queue) {
399     return aws_array_list_capacity(&queue->container);
400 }
401