lib: add internal PriorityQueue class

An efficient JS implementation of a binary heap on top of an array with
worst-case O(log n) runtime for all operations, including arbitrary
item removal (unlike O(n) for most binary heap array implementations).

PR-URL: nodejs#20555
Fixes: nodejs#16105
Reviewed-By: Ben Noordhuis <[email protected]>
Reviewed-By: Jeremiah Senkpiel <[email protected]>
Reviewed-By: Matteo Collina <[email protected]>
Reviewed-By: James M Snell <[email protected]>
Reviewed-By: Ruben Bridgewater <[email protected]>

benchmark/util/priority-queue.js

@@ -0,0 +1,18 @@
+'use strict';
+
+const common = require('../common');
+
+const bench = common.createBenchmark(main, {
+  n: [1e6]
+}, { flags: ['--expose-internals'] });
+
+function main({ n, type }) {
+  const PriorityQueue = require('internal/priority_queue');
+  const queue = new PriorityQueue();
+  bench.start();
+  for (var i = 0; i < n; i++)
+    queue.insert(Math.random() * 1e7 | 0);
+  for (i = 0; i < n; i++)
+    queue.shift();
+  bench.end(n);
+}

lib/internal/priority_queue.js

@@ -0,0 +1,111 @@
+'use strict';
+
+const kCompare = Symbol('compare');
+const kHeap = Symbol('heap');
+const kSetPosition = Symbol('setPosition');
+const kSize = Symbol('size');
+
+// The PriorityQueue is a basic implementation of a binary heap that accepts
+// a custom sorting function via its constructor. This function is passed
+// the two nodes to compare, similar to the native Array#sort. Crucially
+// this enables priority queues that are based on a comparison of more than
+// just a single criteria.
+
+module.exports = class PriorityQueue {
+  constructor(comparator, setPosition) {
+    if (comparator !== undefined)
+      this[kCompare] = comparator;
+    if (setPosition !== undefined)
+      this[kSetPosition] = setPosition;
+
+    this[kHeap] = new Array(64);
+    this[kSize] = 0;
+  }
+
+  [kCompare](a, b) {
+    return a - b;
+  }
+
+  insert(value) {
+    const heap = this[kHeap];
+    let pos = ++this[kSize];
+
+    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);
+  }
+
+  peek() {
+    return this[kHeap][1];
+  }
+
+  percolateDown(pos) {
+    const compare = this[kCompare];
+    const setPosition = this[kSetPosition];
+    const heap = this[kHeap];
+    const size = this[kSize];
+    const item = heap[pos];
+
+    while (pos * 2 <= size) {
+      let childIndex = pos * 2 + 1;
+      if (childIndex > size || compare(heap[pos * 2], heap[childIndex]) < 0)
+        childIndex = pos * 2;
+      const child = heap[childIndex];
+      if (compare(item, child) <= 0)
+        break;
+      if (setPosition !== undefined)
+        setPosition(child, pos);
+      heap[pos] = child;
+      pos = childIndex;
+    }
+    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)
+      this.percolateDown(1);
+  }
+
+  remove(value) {
+    const heap = this[kHeap];
+    const pos = heap.indexOf(value);
+    if (pos < 1)
+      return false;
+
+    this.removeAt(pos);
+
+    return true;
+  }
+
+  shift() {
+    const heap = this[kHeap];
+    const value = heap[1];
+    if (value === undefined)
+      return;
+
+    this.removeAt(1);
+
+    return value;
+  }
+};

node.gyp

@@ -123,6 +123,7 @@
       'lib/internal/safe_globals.js',
       'lib/internal/net.js',
       'lib/internal/os.js',
+      'lib/internal/priority_queue.js',
       'lib/internal/process/esm_loader.js',
       'lib/internal/process/methods.js',
       'lib/internal/process/next_tick.js',

test/parallel/test-priority-queue.js

@@ -0,0 +1,97 @@
+// Flags: --expose-internals
+'use strict';
+
+require('../common');
+
+const assert = require('assert');
+const PriorityQueue = require('internal/priority_queue');
+
+{
+  // Checks that the queue is fundamentally correct.
+  const queue = new PriorityQueue();
+  for (let i = 15; i > 0; i--)
+    queue.insert(i);
+
+  for (let i = 1; i < 16; i++) {
+    assert.strictEqual(queue.peek(), i);
+    assert.strictEqual(queue.shift(), i);
+  }
+
+  assert.strictEqual(queue.shift(), undefined);
+
+  // Reverse the order.
+  for (let i = 1; i < 16; i++)
+    queue.insert(i);
+
+  for (let i = 1; i < 16; i++) {
+    assert.strictEqual(queue.shift(), i);
+  }
+
+  assert.strictEqual(queue.shift(), undefined);
+}
+
+{
+  // Checks that the queue is capable of resizing and fitting more elements.
+  const queue = new PriorityQueue();
+  for (let i = 2048; i > 0; i--)
+    queue.insert(i);
+
+  for (let i = 1; i < 2049; i++) {
+    assert.strictEqual(queue.shift(), i);
+  }
+
+  assert.strictEqual(queue.shift(), undefined);
+}
+
+{
+  // Checks that remove works as expected.
+  const queue = new PriorityQueue();
+  for (let i = 16; i > 0; i--)
+    queue.insert(i);
+
+  const removed = [5, 10, 15];
+  for (const id of removed)
+    assert(queue.remove(id));
+
+  assert(!queue.remove(100));
+  assert(!queue.remove(-100));
+
+  for (let i = 1; i < 17; i++) {
+    if (removed.indexOf(i) < 0)
+      assert.strictEqual(queue.shift(), i);
+  }
+
+  assert.strictEqual(queue.shift(), undefined);
+}
+
+{
+  // Make a max heap with a custom sort function.
+  const queue = new PriorityQueue((a, b) => b - a);
+  for (let i = 1; i < 17; i++)
+    queue.insert(i);
+
+  for (let i = 16; i > 0; i--) {
+    assert.strictEqual(queue.shift(), i);
+  }
+
+  assert.strictEqual(queue.shift(), undefined);
+}
+
+{
+  // Make a min heap that accepts objects as values, which necessitates
+  // a custom sorting function. In addition, add a setPosition function
+  // as 2nd param which provides a reference to the node in the heap
+  // and allows speedy deletions.
+  const queue = new PriorityQueue((a, b) => {
+    return a.value - b.value;
+  }, (node, pos) => (node.position = pos));
+  for (let i = 1; i < 17; i++)
+    queue.insert({ value: i, position: null });
+
+  for (let i = 1; i < 17; i++) {
+    assert.strictEqual(queue.peek().value, i);
+    queue.removeAt(queue.peek().position);
+  }
+
+  assert.strictEqual(queue.peek(), undefined);
+}