timers: fix priority queue removeAt

nodejs · Nov 12, 2018 · 2203dcc · 2203dcc
1 parent 8884a98
commit 2203dcc
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 21 deletions.
diff --git a/lib/internal/priority_queue.js b/lib/internal/priority_queue.js
@@ -28,25 +28,13 @@ module.exports = class PriorityQueue {
 
   insert(value) {
     const heap = this[kHeap];
-    let pos = ++this[kSize];
+    const pos = ++this[kSize];
+    heap[pos] = value;
 
     if (heap.length === pos)
       heap.length *= 2;
 
-    const compare = this[kCompare];
-    const setPosition = this[kSetPosition];
-    while (pos > 1) {
-      const parent = heap[pos / 2 | 0];
-      if (compare(parent, value) <= 0)
-        break;
-      heap[pos] = parent;
-      if (setPosition !== undefined)
-        setPosition(parent, pos);
-      pos = pos / 2 | 0;
-    }
-    heap[pos] = value;
-    if (setPosition !== undefined)
-      setPosition(value, pos);
+    this.percolateUp(pos);
   }
 
   peek() {
@@ -77,18 +65,37 @@ module.exports = class PriorityQueue {
       setPosition(item, pos);
   }
 
+  percolateUp(pos) {
+    const heap = this[kHeap];
+    const compare = this[kCompare];
+    const setPosition = this[kSetPosition];
+    const item = heap[pos];
+
+    while (pos > 1) {
+      const parent = heap[pos / 2 | 0];
+      if (compare(parent, item) <= 0)
+        break;
+      heap[pos] = parent;
+      if (setPosition !== undefined)
+        setPosition(parent, pos);
+      pos = pos / 2 | 0;
+    }
+    heap[pos] = item;
+    if (setPosition !== undefined)
+      setPosition(item, pos);
+  }
+
   removeAt(pos) {
     const heap = this[kHeap];
     const size = --this[kSize];
     heap[pos] = heap[size + 1];
     heap[size + 1] = undefined;
 
-    if (size > 0) {
-      // If not removing the last item, update the shifted item's position.
-      if (pos <= size && this[kSetPosition] !== undefined)
-        this[kSetPosition](heap[pos], pos);
-
-      this.percolateDown(1);
+    if (size > 0 && pos <= size) {
+      if (pos > 1 && this[kCompare](heap[pos / 2 | 0], heap[pos]) > 0)
+        this.percolateUp(pos);
+      else
+        this.percolateDown(pos);
     }
   }
 

diff --git a/test/parallel/test-priority-queue.js b/test/parallel/test-priority-queue.js
@@ -131,3 +131,33 @@ const PriorityQueue = require('internal/priority_queue');
 
   assert.strictEqual(queue.shift(), undefined);
 }
+
+
+{
+  // Checks that removeAt respects binary heap properties
+  const queue = new PriorityQueue((a, b) => {
+    return a.value - b.value;
+  }, (node, pos) => (node.position = pos));
+
+  const i3 = { value: 3, position: null };
+  const i7 = { value: 7, position: null };
+  const i8 = { value: 8, position: null };
+
+  queue.insert({ value: 1, position: null });
+  queue.insert({ value: 6, position: null });
+  queue.insert({ value: 2, position: null });
+  queue.insert(i7);
+  queue.insert(i8);
+  queue.insert(i3);
+
+  assert.strictEqual(i7.position, 4);
+  queue.removeAt(4);
+
+  // 3 should percolate up to swap with 6 (up)
+  assert.strictEqual(i3.position, 2);
+
+  queue.removeAt(2);
+
+  // 8 should swap places with 6 (down)
+  assert.strictEqual(i8.position, 4);
+}