Branch data     Line data    Source code

       1                 :            : /*
       2                 :            :  * This file is part of the MicroPython 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/objlist.h"
      28                 :            : #include "py/runtime.h"
      29                 :            : 
      30                 :            : #if MICROPY_PY_HEAPQ
      31                 :            : 
      32                 :            : // the algorithm here is modelled on CPython's heapq.py
      33                 :            : 
      34                 :         48 : static mp_obj_list_t *heapq_get_heap(mp_obj_t heap_in) {
      35   [ -  +  -  +  :         48 :     if (!mp_obj_is_type(heap_in, &mp_type_list)) {
          -  +  -  +  +  
                -  +  + ]
      36                 :          2 :         mp_raise_TypeError(MP_ERROR_TEXT("heap must be a list"));
      37                 :            :     }
      38                 :         46 :     return MP_OBJ_TO_PTR(heap_in);
      39                 :            : }
      40                 :            : 
      41                 :         46 : static void heapq_heap_siftdown(mp_obj_list_t *heap, mp_uint_t start_pos, mp_uint_t pos) {
      42                 :         46 :     mp_obj_t item = heap->items[pos];
      43         [ +  + ]:         62 :     while (pos > start_pos) {
      44                 :         48 :         mp_uint_t parent_pos = (pos - 1) >> 1;
      45                 :         48 :         mp_obj_t parent = heap->items[parent_pos];
      46         [ +  + ]:         48 :         if (mp_binary_op(MP_BINARY_OP_LESS, item, parent) == mp_const_true) {
      47                 :         16 :             heap->items[pos] = parent;
      48                 :         16 :             pos = parent_pos;
      49                 :            :         } else {
      50                 :            :             break;
      51                 :            :         }
      52                 :            :     }
      53                 :         46 :     heap->items[pos] = item;
      54                 :         46 : }
      55                 :            : 
      56                 :         34 : static void heapq_heap_siftup(mp_obj_list_t *heap, mp_uint_t pos) {
      57                 :         34 :     mp_uint_t start_pos = pos;
      58                 :         34 :     mp_uint_t end_pos = heap->len;
      59                 :         34 :     mp_obj_t item = heap->items[pos];
      60         [ +  + ]:         86 :     for (mp_uint_t child_pos = 2 * pos + 1; child_pos < end_pos; child_pos = 2 * pos + 1) {
      61                 :            :         // choose right child if it's <= left child
      62   [ +  +  +  + ]:         52 :         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) {
      63                 :         26 :             child_pos += 1;
      64                 :            :         }
      65                 :            :         // bubble up the smaller child
      66                 :         52 :         heap->items[pos] = heap->items[child_pos];
      67                 :         52 :         pos = child_pos;
      68                 :            :     }
      69                 :         34 :     heap->items[pos] = item;
      70                 :         34 :     heapq_heap_siftdown(heap, start_pos, pos);
      71                 :         34 : }
      72                 :            : 
      73                 :         14 : static mp_obj_t mod_heapq_heappush(mp_obj_t heap_in, mp_obj_t item) {
      74                 :         14 :     mp_obj_list_t *heap = heapq_get_heap(heap_in);
      75                 :         12 :     mp_obj_list_append(heap_in, item);
      76                 :         12 :     heapq_heap_siftdown(heap, 0, heap->len - 1);
      77                 :         12 :     return mp_const_none;
      78                 :            : }
      79                 :            : static MP_DEFINE_CONST_FUN_OBJ_2(mod_heapq_heappush_obj, mod_heapq_heappush);
      80                 :            : 
      81                 :         32 : static mp_obj_t mod_heapq_heappop(mp_obj_t heap_in) {
      82                 :         32 :     mp_obj_list_t *heap = heapq_get_heap(heap_in);
      83         [ +  + ]:         32 :     if (heap->len == 0) {
      84                 :          2 :         mp_raise_msg(&mp_type_IndexError, MP_ERROR_TEXT("empty heap"));
      85                 :            :     }
      86                 :         30 :     mp_obj_t item = heap->items[0];
      87                 :         30 :     heap->len -= 1;
      88                 :         30 :     heap->items[0] = heap->items[heap->len];
      89                 :         30 :     heap->items[heap->len] = MP_OBJ_NULL; // so we don't retain a pointer
      90         [ +  + ]:         30 :     if (heap->len) {
      91                 :         26 :         heapq_heap_siftup(heap, 0);
      92                 :            :     }
      93                 :         30 :     return item;
      94                 :            : }
      95                 :            : static MP_DEFINE_CONST_FUN_OBJ_1(mod_heapq_heappop_obj, mod_heapq_heappop);
      96                 :            : 
      97                 :          2 : static mp_obj_t mod_heapq_heapify(mp_obj_t heap_in) {
      98                 :          2 :     mp_obj_list_t *heap = heapq_get_heap(heap_in);
      99         [ +  + ]:         10 :     for (mp_uint_t i = heap->len / 2; i > 0;) {
     100                 :          8 :         heapq_heap_siftup(heap, --i);
     101                 :            :     }
     102                 :          2 :     return mp_const_none;
     103                 :            : }
     104                 :            : static MP_DEFINE_CONST_FUN_OBJ_1(mod_heapq_heapify_obj, mod_heapq_heapify);
     105                 :            : 
     106                 :            : #if !MICROPY_ENABLE_DYNRUNTIME
     107                 :            : static const mp_rom_map_elem_t mp_module_heapq_globals_table[] = {
     108                 :            :     { MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_heapq) },
     109                 :            :     { MP_ROM_QSTR(MP_QSTR_heappush), MP_ROM_PTR(&mod_heapq_heappush_obj) },
     110                 :            :     { MP_ROM_QSTR(MP_QSTR_heappop), MP_ROM_PTR(&mod_heapq_heappop_obj) },
     111                 :            :     { MP_ROM_QSTR(MP_QSTR_heapify), MP_ROM_PTR(&mod_heapq_heapify_obj) },
     112                 :            : };
     113                 :            : 
     114                 :            : static MP_DEFINE_CONST_DICT(mp_module_heapq_globals, mp_module_heapq_globals_table);
     115                 :            : 
     116                 :            : const mp_obj_module_t mp_module_heapq = {
     117                 :            :     .base = { &mp_type_module },
     118                 :            :     .globals = (mp_obj_dict_t *)&mp_module_heapq_globals,
     119                 :            : };
     120                 :            : 
     121                 :            : MP_REGISTER_EXTENSIBLE_MODULE(MP_QSTR_heapq, mp_module_heapq);
     122                 :            : #endif
     123                 :            : 
     124                 :            : #endif // MICROPY_PY_HEAPQ