py/pairheap: Properly unlink node on pop and delete.

This fixes a bug in the pairing-heap implementation when nodes are deleted
with mp_pairheap_delete and then reinserted later on.

2 files changed, 17 insertions, 6 deletions

diff --git a/py/pairheap.c b/py/pairheap.c
index a1eececa0..d3a011c4a 100644
--- a/py/pairheap.c
+++ b/py/pairheap.c
@@ -72,9 +72,11 @@ mp_pairheap_t *mp_pairheap_pairing(mp_pairheap_lt_t lt, mp_pairheap_t *child) {
     while (!NEXT_IS_RIGHTMOST_PARENT(child)) {
         mp_pairheap_t *n1 = child;
         child = child->next;
+        n1->next = NULL;
         if (!NEXT_IS_RIGHTMOST_PARENT(child)) {
             mp_pairheap_t *n2 = child;
             child = child->next;
+            n2->next = NULL;
             n1 = mp_pairheap_meld(lt, n1, n2);
         }
         heap = mp_pairheap_meld(lt, heap, n1);
@@ -87,7 +89,9 @@ mp_pairheap_t *mp_pairheap_pairing(mp_pairheap_lt_t lt, mp_pairheap_t *child) {
 mp_pairheap_t *mp_pairheap_delete(mp_pairheap_lt_t lt, mp_pairheap_t *heap, mp_pairheap_t *node) {
     // Simple case of the top being the node to delete
     if (node == heap) {
-        return mp_pairheap_pairing(lt, heap->child);
+        mp_pairheap_t *child = heap->child;
+        node->child = NULL;
+        return mp_pairheap_pairing(lt, child);
     }
 
     // Case where node is not in the heap
@@ -113,18 +117,22 @@ mp_pairheap_t *mp_pairheap_delete(mp_pairheap_lt_t lt, mp_pairheap_t *heap, mp_p
         node->next = NULL;
         return heap;
     } else if (node == parent->child) {
+        mp_pairheap_t *child = node->child;
         next = node->next;
+        node->child = NULL;
         node->next = NULL;
-        node = mp_pairheap_pairing(lt, node->child);
+        node = mp_pairheap_pairing(lt, child);
         parent->child = node;
     } else {
         mp_pairheap_t *n = parent->child;
         while (node != n->next) {
             n = n->next;
         }
+        mp_pairheap_t *child = node->child;
         next = node->next;
+        node->child = NULL;
         node->next = NULL;
-        node = mp_pairheap_pairing(lt, node->child);
+        node = mp_pairheap_pairing(lt, child);
         if (node == NULL) {
             node = n;
         } else {
diff --git a/py/pairheap.h b/py/pairheap.h
index 8a9138b17..16ae78809 100644
--- a/py/pairheap.h
+++ b/py/pairheap.h
@@ -35,6 +35,7 @@
 // Algorithmica 1:111-129, 1986.
 // https://www.cs.cmu.edu/~sleator/papers/pairing-heaps.pdf
 
+#include <assert.h>
 #include "py/obj.h"
 
 // This struct forms the nodes of the heap and is intended to be extended, by
@@ -77,14 +78,16 @@ static inline mp_pairheap_t *mp_pairheap_peek(mp_pairheap_lt_t lt, mp_pairheap_t
 
 // Push new node onto existing heap.  Returns the new heap.
 static inline mp_pairheap_t *mp_pairheap_push(mp_pairheap_lt_t lt, mp_pairheap_t *heap, mp_pairheap_t *node) {
-    node->child = NULL;
-    node->next = NULL;
+    assert(node->child == NULL && node->next == NULL);
     return mp_pairheap_meld(lt, node, heap); // node is first to be stable
 }
 
 // Pop the top off the heap, which must not be empty.  Returns the new heap.
 static inline mp_pairheap_t *mp_pairheap_pop(mp_pairheap_lt_t lt, mp_pairheap_t *heap) {
-    return mp_pairheap_pairing(lt, heap->child);
+    assert(heap->next == NULL);
+    mp_pairheap_t *child = heap->child;
+    heap->child = NULL;
+    return mp_pairheap_pairing(lt, child);
 }
 
 #endif // MICROPY_INCLUDED_PY_PAIRHEAP_H