diff --git a/benchmark/util/priority-queue.js b/benchmark/util/priority-queue.js new file mode 100644 index 00000000000000..51a696439a2864 --- /dev/null +++ b/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); +} diff --git a/lib/internal/priority_queue.js b/lib/internal/priority_queue.js new file mode 100644 index 00000000000000..cb046507a667d9 --- /dev/null +++ b/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; + } +}; diff --git a/node.gyp b/node.gyp index bf768c1995cd20..038e5219bcc0fe 100644 --- a/node.gyp +++ b/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', diff --git a/test/parallel/test-priority-queue.js b/test/parallel/test-priority-queue.js new file mode 100644 index 00000000000000..5b8f53a1766bb5 --- /dev/null +++ b/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); +}