2828#include "py/objlist.h"
2929#include "py/runtime0.h"
3030#include "py/runtime.h"
31- #include "py/smallint.h"
3231
3332#if MICROPY_PY_UHEAPQ
3433
35- #define MODULO MICROPY_PY_UTIME_TICKS_PERIOD
36-
3734// the algorithm here is modelled on CPython's heapq.py
3835
3936STATIC mp_obj_list_t * get_heap (mp_obj_t heap_in ) {
@@ -43,33 +40,12 @@ STATIC mp_obj_list_t *get_heap(mp_obj_t heap_in) {
4340 return MP_OBJ_TO_PTR (heap_in );
4441}
4542
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 ) {
43+ STATIC void heap_siftdown (mp_obj_list_t * heap , mp_uint_t start_pos , mp_uint_t pos ) {
6244 mp_obj_t item = heap -> items [pos ];
6345 while (pos > start_pos ) {
6446 mp_uint_t parent_pos = (pos - 1 ) >> 1 ;
6547 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 ) {
48+ if (mp_binary_op (MP_BINARY_OP_LESS , item , parent ) == mp_const_true ) {
7349 heap -> items [pos ] = parent ;
7450 pos = parent_pos ;
7551 } else {
@@ -79,43 +55,32 @@ STATIC void heap_siftdown(mp_obj_list_t *heap, mp_uint_t start_pos, mp_uint_t po
7955 heap -> items [pos ] = item ;
8056}
8157
82- STATIC void heap_siftup (mp_obj_list_t * heap , mp_uint_t pos , bool timecmp ) {
58+ STATIC void heap_siftup (mp_obj_list_t * heap , mp_uint_t pos ) {
8359 mp_uint_t start_pos = pos ;
8460 mp_uint_t end_pos = heap -> len ;
8561 mp_obj_t item = heap -> items [pos ];
8662 for (mp_uint_t child_pos = 2 * pos + 1 ; child_pos < end_pos ; child_pos = 2 * pos + 1 ) {
8763 // 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- }
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 ;
9866 }
9967 // bubble up the smaller child
10068 heap -> items [pos ] = heap -> items [child_pos ];
10169 pos = child_pos ;
10270 }
10371 heap -> items [pos ] = item ;
104- heap_siftdown (heap , start_pos , pos , timecmp );
72+ heap_siftdown (heap , start_pos , pos );
10573}
10674
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 ];
75+ STATIC mp_obj_t mod_uheapq_heappush (mp_obj_t heap_in , mp_obj_t item ) {
10976 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 );
77+ mp_obj_list_append (heap_in , item );
78+ heap_siftdown (heap , 0 , heap -> len - 1 );
11379 return mp_const_none ;
11480}
115- STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN (mod_uheapq_heappush_obj , 2 , 3 , mod_uheapq_heappush );
81+ STATIC MP_DEFINE_CONST_FUN_OBJ_2 (mod_uheapq_heappush_obj , mod_uheapq_heappush );
11682
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 ];
83+ STATIC mp_obj_t mod_uheapq_heappop (mp_obj_t heap_in ) {
11984 mp_obj_list_t * heap = get_heap (heap_in );
12085 if (heap -> len == 0 ) {
12186 nlr_raise (mp_obj_new_exception_msg (& mp_type_IndexError , "empty heap" ));
@@ -125,17 +90,16 @@ STATIC mp_obj_t mod_uheapq_heappop(size_t n_args, const mp_obj_t *args) {
12590 heap -> items [0 ] = heap -> items [heap -> len ];
12691 heap -> items [heap -> len ] = MP_OBJ_NULL ; // so we don't retain a pointer
12792 if (heap -> len ) {
128- bool is_timeq = (n_args > 1 && args [1 ] == mp_const_true );
129- heap_siftup (heap , 0 , is_timeq );
93+ heap_siftup (heap , 0 );
13094 }
13195 return item ;
13296}
133- STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN (mod_uheapq_heappop_obj , 1 , 2 , mod_uheapq_heappop );
97+ STATIC MP_DEFINE_CONST_FUN_OBJ_1 (mod_uheapq_heappop_obj , mod_uheapq_heappop );
13498
13599STATIC mp_obj_t mod_uheapq_heapify (mp_obj_t heap_in ) {
136100 mp_obj_list_t * heap = get_heap (heap_in );
137101 for (mp_uint_t i = heap -> len / 2 ; i > 0 ;) {
138- heap_siftup (heap , -- i , false );
102+ heap_siftup (heap , -- i );
139103 }
140104 return mp_const_none ;
141105}
0 commit comments