Skip to content

Commit ef23399

Browse files
pfalconpfalcon
committed
extmod/modutimeq: Copy of current moduheapq with timeq support for refactoring.
1 parent 1731868 commit ef23399

1 file changed

Lines changed: 158 additions & 0 deletions

File tree

‎extmod/modutimeq.c‎

Lines changed: 158 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,158 @@
1+
/*
2+
* This file is part of the Micro Python project, http://micropython.org/
3+
*
4+
* The MIT License (MIT)
5+
*
6+
* Copyright (c) 2014 Damien P. George
7+
*
8+
* Permission is hereby granted, free of charge, to any person obtaining a copy
9+
* of this software and associated documentation files (the "Software"), to deal
10+
* in the Software without restriction, including without limitation the rights
11+
* to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
12+
* copies of the Software, and to permit persons to whom the Software is
13+
* furnished to do so, subject to the following conditions:
14+
*
15+
* The above copyright notice and this permission notice shall be included in
16+
* all copies or substantial portions of the Software.
17+
*
18+
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
19+
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
20+
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
21+
* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
22+
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
23+
* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
24+
* THE SOFTWARE.
25+
*/
26+
27+
#include "py/nlr.h"
28+
#include "py/objlist.h"
29+
#include "py/runtime0.h"
30+
#include "py/runtime.h"
31+
#include "py/smallint.h"
32+
33+
#if MICROPY_PY_UHEAPQ
34+
35+
#define MODULO MICROPY_PY_UTIME_TICKS_PERIOD
36+
37+
// the algorithm here is modelled on CPython's heapq.py
38+
39+
STATIC mp_obj_list_t *get_heap(mp_obj_t heap_in) {
40+
if (!MP_OBJ_IS_TYPE(heap_in, &mp_type_list)) {
41+
nlr_raise(mp_obj_new_exception_msg(&mp_type_TypeError, "heap must be a list"));
42+
}
43+
return MP_OBJ_TO_PTR(heap_in);
44+
}
45+
46+
STATIC bool time_less_than(mp_obj_t item, mp_obj_t parent) {
47+
if (!MP_OBJ_IS_TYPE(item, &mp_type_tuple) || !MP_OBJ_IS_TYPE(parent, &mp_type_tuple)) {
48+
mp_raise_TypeError("");
49+
}
50+
mp_obj_tuple_t *item_p = MP_OBJ_TO_PTR(item);
51+
mp_obj_tuple_t *parent_p = MP_OBJ_TO_PTR(parent);
52+
mp_uint_t item_tm = MP_OBJ_SMALL_INT_VALUE(item_p->items[0]);
53+
mp_uint_t parent_tm = MP_OBJ_SMALL_INT_VALUE(parent_p->items[0]);
54+
mp_uint_t res = parent_tm - item_tm;
55+
if ((mp_int_t)res < 0) {
56+
res += MODULO;
57+
}
58+
return res < (MODULO / 2);
59+
}
60+
61+
STATIC void heap_siftdown(mp_obj_list_t *heap, mp_uint_t start_pos, mp_uint_t pos, bool timecmp) {
62+
mp_obj_t item = heap->items[pos];
63+
while (pos > start_pos) {
64+
mp_uint_t parent_pos = (pos - 1) >> 1;
65+
mp_obj_t parent = heap->items[parent_pos];
66+
bool lessthan;
67+
if (MP_UNLIKELY(timecmp)) {
68+
lessthan = time_less_than(item, parent);
69+
} else {
70+
lessthan = (mp_binary_op(MP_BINARY_OP_LESS, item, parent) == mp_const_true);
71+
}
72+
if (lessthan) {
73+
heap->items[pos] = parent;
74+
pos = parent_pos;
75+
} else {
76+
break;
77+
}
78+
}
79+
heap->items[pos] = item;
80+
}
81+
82+
STATIC void heap_siftup(mp_obj_list_t *heap, mp_uint_t pos, bool timecmp) {
83+
mp_uint_t start_pos = pos;
84+
mp_uint_t end_pos = heap->len;
85+
mp_obj_t item = heap->items[pos];
86+
for (mp_uint_t child_pos = 2 * pos + 1; child_pos < end_pos; child_pos = 2 * pos + 1) {
87+
// choose right child if it's <= left child
88+
if (child_pos + 1 < end_pos) {
89+
bool lessthan;
90+
if (MP_UNLIKELY(timecmp)) {
91+
lessthan = time_less_than(heap->items[child_pos], heap->items[child_pos + 1]);
92+
} else {
93+
lessthan = (mp_binary_op(MP_BINARY_OP_LESS, heap->items[child_pos], heap->items[child_pos + 1]) == mp_const_true);
94+
}
95+
if (!lessthan) {
96+
child_pos += 1;
97+
}
98+
}
99+
// bubble up the smaller child
100+
heap->items[pos] = heap->items[child_pos];
101+
pos = child_pos;
102+
}
103+
heap->items[pos] = item;
104+
heap_siftdown(heap, start_pos, pos, timecmp);
105+
}
106+
107+
STATIC mp_obj_t mod_utimeq_heappush(size_t n_args, const mp_obj_t *args) {
108+
mp_obj_t heap_in = args[0];
109+
mp_obj_list_t *heap = get_heap(heap_in);
110+
mp_obj_list_append(heap_in, args[1]);
111+
bool is_timeq = (n_args > 2 && args[2] == mp_const_true);
112+
heap_siftdown(heap, 0, heap->len - 1, is_timeq);
113+
return mp_const_none;
114+
}
115+
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mod_utimeq_heappush_obj, 2, 3, mod_utimeq_heappush);
116+
117+
STATIC mp_obj_t mod_utimeq_heappop(size_t n_args, const mp_obj_t *args) {
118+
mp_obj_t heap_in = args[0];
119+
mp_obj_list_t *heap = get_heap(heap_in);
120+
if (heap->len == 0) {
121+
nlr_raise(mp_obj_new_exception_msg(&mp_type_IndexError, "empty heap"));
122+
}
123+
mp_obj_t item = heap->items[0];
124+
heap->len -= 1;
125+
heap->items[0] = heap->items[heap->len];
126+
heap->items[heap->len] = MP_OBJ_NULL; // so we don't retain a pointer
127+
if (heap->len) {
128+
bool is_timeq = (n_args > 1 && args[1] == mp_const_true);
129+
heap_siftup(heap, 0, is_timeq);
130+
}
131+
return item;
132+
}
133+
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mod_utimeq_heappop_obj, 1, 2, mod_utimeq_heappop);
134+
135+
STATIC mp_obj_t mod_utimeq_heapify(mp_obj_t heap_in) {
136+
mp_obj_list_t *heap = get_heap(heap_in);
137+
for (mp_uint_t i = heap->len / 2; i > 0;) {
138+
heap_siftup(heap, --i, false);
139+
}
140+
return mp_const_none;
141+
}
142+
STATIC MP_DEFINE_CONST_FUN_OBJ_1(mod_utimeq_heapify_obj, mod_utimeq_heapify);
143+
144+
STATIC const mp_rom_map_elem_t mp_module_utimeq_globals_table[] = {
145+
{ MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_utimeq) },
146+
{ MP_ROM_QSTR(MP_QSTR_heappush), MP_ROM_PTR(&mod_utimeq_heappush_obj) },
147+
{ MP_ROM_QSTR(MP_QSTR_heappop), MP_ROM_PTR(&mod_utimeq_heappop_obj) },
148+
{ MP_ROM_QSTR(MP_QSTR_heapify), MP_ROM_PTR(&mod_utimeq_heapify_obj) },
149+
};
150+
151+
STATIC MP_DEFINE_CONST_DICT(mp_module_utimeq_globals, mp_module_utimeq_globals_table);
152+
153+
const mp_obj_module_t mp_module_utimeq = {
154+
.base = { &mp_type_module },
155+
.globals = (mp_obj_dict_t*)&mp_module_utimeq_globals,
156+
};
157+
158+
#endif //MICROPY_PY_UHEAPQ

0 commit comments

Comments
 (0)