Skip to content

Commit 0cbc072

Browse files
pfalconpfalcon
committed
extmod/moduheapq: Adhoc changes to support ordering by utime.ticks_ms().
As required for further elaboration of uasyncio, like supporting baremetal systems with wraparound timesources. This is not intended to be public interface, and likely will be further refactored in the future.
1 parent 3c0da6a commit 0cbc072

3 files changed

Lines changed: 145 additions & 14 deletions

File tree

‎extmod/moduheapq.c‎

Lines changed: 50 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -28,9 +28,12 @@
2828
#include "py/objlist.h"
2929
#include "py/runtime0.h"
3030
#include "py/runtime.h"
31+
#include "py/smallint.h"
3132

3233
#if MICROPY_PY_UHEAPQ
3334

35+
#define MODULO MICROPY_PY_UTIME_TICKS_PERIOD
36+
3437
// the algorithm here is modelled on CPython's heapq.py
3538

3639
STATIC mp_obj_list_t *get_heap(mp_obj_t heap_in) {
@@ -40,12 +43,33 @@ STATIC mp_obj_list_t *get_heap(mp_obj_t heap_in) {
4043
return MP_OBJ_TO_PTR(heap_in);
4144
}
4245

43-
STATIC void heap_siftdown(mp_obj_list_t *heap, mp_uint_t start_pos, mp_uint_t pos) {
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) {
4462
mp_obj_t item = heap->items[pos];
4563
while (pos > start_pos) {
4664
mp_uint_t parent_pos = (pos - 1) >> 1;
4765
mp_obj_t parent = heap->items[parent_pos];
48-
if (mp_binary_op(MP_BINARY_OP_LESS, item, parent) == mp_const_true) {
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) {
4973
heap->items[pos] = parent;
5074
pos = parent_pos;
5175
} else {
@@ -55,32 +79,43 @@ STATIC void heap_siftdown(mp_obj_list_t *heap, mp_uint_t start_pos, mp_uint_t po
5579
heap->items[pos] = item;
5680
}
5781

58-
STATIC void heap_siftup(mp_obj_list_t *heap, mp_uint_t pos) {
82+
STATIC void heap_siftup(mp_obj_list_t *heap, mp_uint_t pos, bool timecmp) {
5983
mp_uint_t start_pos = pos;
6084
mp_uint_t end_pos = heap->len;
6185
mp_obj_t item = heap->items[pos];
6286
for (mp_uint_t child_pos = 2 * pos + 1; child_pos < end_pos; child_pos = 2 * pos + 1) {
6387
// choose right child if it's <= left child
64-
if (child_pos + 1 < end_pos && mp_binary_op(MP_BINARY_OP_LESS, heap->items[child_pos], heap->items[child_pos + 1]) == mp_const_false) {
65-
child_pos += 1;
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+
}
6698
}
6799
// bubble up the smaller child
68100
heap->items[pos] = heap->items[child_pos];
69101
pos = child_pos;
70102
}
71103
heap->items[pos] = item;
72-
heap_siftdown(heap, start_pos, pos);
104+
heap_siftdown(heap, start_pos, pos, timecmp);
73105
}
74106

75-
STATIC mp_obj_t mod_uheapq_heappush(mp_obj_t heap_in, mp_obj_t item) {
107+
STATIC mp_obj_t mod_uheapq_heappush(size_t n_args, const mp_obj_t *args) {
108+
mp_obj_t heap_in = args[0];
76109
mp_obj_list_t *heap = get_heap(heap_in);
77-
mp_obj_list_append(heap_in, item);
78-
heap_siftdown(heap, 0, heap->len - 1);
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);
79113
return mp_const_none;
80114
}
81-
STATIC MP_DEFINE_CONST_FUN_OBJ_2(mod_uheapq_heappush_obj, mod_uheapq_heappush);
115+
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mod_uheapq_heappush_obj, 2, 3, mod_uheapq_heappush);
82116

83-
STATIC mp_obj_t mod_uheapq_heappop(mp_obj_t heap_in) {
117+
STATIC mp_obj_t mod_uheapq_heappop(size_t n_args, const mp_obj_t *args) {
118+
mp_obj_t heap_in = args[0];
84119
mp_obj_list_t *heap = get_heap(heap_in);
85120
if (heap->len == 0) {
86121
nlr_raise(mp_obj_new_exception_msg(&mp_type_IndexError, "empty heap"));
@@ -90,16 +125,17 @@ STATIC mp_obj_t mod_uheapq_heappop(mp_obj_t heap_in) {
90125
heap->items[0] = heap->items[heap->len];
91126
heap->items[heap->len] = MP_OBJ_NULL; // so we don't retain a pointer
92127
if (heap->len) {
93-
heap_siftup(heap, 0);
128+
bool is_timeq = (n_args > 1 && args[1] == mp_const_true);
129+
heap_siftup(heap, 0, is_timeq);
94130
}
95131
return item;
96132
}
97-
STATIC MP_DEFINE_CONST_FUN_OBJ_1(mod_uheapq_heappop_obj, mod_uheapq_heappop);
133+
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mod_uheapq_heappop_obj, 1, 2, mod_uheapq_heappop);
98134

99135
STATIC mp_obj_t mod_uheapq_heapify(mp_obj_t heap_in) {
100136
mp_obj_list_t *heap = get_heap(heap_in);
101137
for (mp_uint_t i = heap->len / 2; i > 0;) {
102-
heap_siftup(heap, --i);
138+
heap_siftup(heap, --i, false);
103139
}
104140
return mp_const_none;
105141
}

‎tests/extmod/uheapq_timeq.py‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
# Test adhoc extension to uheapq to support wraparound
2+
# time (utime.ticks_ms() style) task queue.
3+
from utime import ticks_add, ticks_diff
4+
import uheapq as heapq
5+
6+
DEBUG = 0
7+
8+
MAX = ticks_add(0, -1)
9+
MODULO_HALF = MAX // 2 + 1
10+
11+
if DEBUG:
12+
def dprint(*v):
13+
print(*v)
14+
else:
15+
def dprint(*v):
16+
pass
17+
18+
# Try not to crash on invalid data
19+
h = []
20+
heapq.heappush(h, 1)
21+
try:
22+
heapq.heappush(h, 2, True)
23+
assert False
24+
except TypeError:
25+
pass
26+
27+
heapq.heappush(h, 2)
28+
try:
29+
heapq.heappop(h, True)
30+
assert False
31+
except TypeError:
32+
pass
33+
34+
35+
def pop_all(h):
36+
l = []
37+
while h:
38+
l.append(heapq.heappop(h, True))
39+
dprint(l)
40+
return l
41+
42+
def add(h, v):
43+
heapq.heappush(h, (v, None), True)
44+
45+
h = []
46+
add(h, 0)
47+
add(h, MAX)
48+
add(h, MAX - 1)
49+
add(h, 101)
50+
add(h, 100)
51+
add(h, MAX - 2)
52+
dprint(h)
53+
l = pop_all(h)
54+
for i in range(len(l) - 1):
55+
diff = ticks_diff(l[i + 1][0], l[i][0])
56+
assert diff > 0
57+
58+
def edge_case(edge, offset):
59+
h = []
60+
add(h, ticks_add(0, offset))
61+
add(h, ticks_add(edge, offset))
62+
dprint(h)
63+
l = pop_all(h)
64+
diff = ticks_diff(l[1][0], l[0][0])
65+
dprint(diff, diff > 0)
66+
return diff
67+
68+
dprint("===")
69+
diff = edge_case(MODULO_HALF - 1, 0)
70+
assert diff == MODULO_HALF - 1
71+
assert edge_case(MODULO_HALF - 1, 100) == diff
72+
assert edge_case(MODULO_HALF - 1, -100) == diff
73+
74+
# We expect diff to be always positive, per the definition of heappop() which should return
75+
# the smallest value.
76+
# This is the edge case where this invariant breaks, due to assymetry of two's-complement
77+
# range - there's one more negative integer than positive, so heappushing values like below
78+
# will then make ticks_diff() return the minimum negative value. We could make heappop
79+
# return them in a different order, but ticks_diff() result would be the same. Conclusion:
80+
# never add to a heap values where (a - b) == MODULO_HALF (and which are >= MODULO_HALF
81+
# ticks apart in real time of course).
82+
dprint("===")
83+
diff = edge_case(MODULO_HALF, 0)
84+
assert diff == -MODULO_HALF
85+
assert edge_case(MODULO_HALF, 100) == diff
86+
assert edge_case(MODULO_HALF, -100) == diff
87+
88+
dprint("===")
89+
diff = edge_case(MODULO_HALF + 1, 0)
90+
assert diff == MODULO_HALF - 1
91+
assert edge_case(MODULO_HALF + 1, 100) == diff
92+
assert edge_case(MODULO_HALF + 1, -100) == diff
93+
94+
print("OK")

‎tests/extmod/uheapq_timeq.py.exp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
OK

0 commit comments

Comments
 (0)