diff --git a/src/data_manager/priority_queue.ts b/src/data_manager/priority_queue.ts new file mode 100644 index 00000000..e512dc70 --- /dev/null +++ b/src/data_manager/priority_queue.ts @@ -0,0 +1,128 @@ +type Entry = { + priority: P; + value: V; + heapIndex: number; +}; + +const entryToTuple = (entry?: Entry): [P, V] | undefined => entry && [entry.priority, entry.value]; + +const parentIndexOf = (i: number) => Math.ceil(i / 2) - 1; +const childIndexOf = (i: number) => i * 2 + 1; + +export class PriorityQueue { + private keys: Map> = new Map(); + private heap: Entry[] = []; + + constructor(private compare: (a: P, b: P) => number) {} + + get length() { + return this.heap.length; + } + + /** Useful common primitive for `siftUp` and `siftDown`: compares two entries and swaps them if out of order. */ + private trySwap(entry: Entry, swapEntry?: Entry, mul = 1): boolean { + if (swapEntry !== undefined && this.compare(entry.priority, swapEntry.priority) * mul < 0) { + this.heap[entry.heapIndex] = swapEntry; + const newHeapIndex = swapEntry.heapIndex; + swapEntry.heapIndex = entry.heapIndex; + entry.heapIndex = newHeapIndex; + return true; + } + + return false; + } + + /** Moves `entry` *up* to the correct position within the heap. */ + private siftUp(entry: Entry) { + let parentEntry: Entry | undefined; + + do { + parentEntry = this.heap[parentIndexOf(entry.heapIndex)]; + } while (this.trySwap(entry, parentEntry)); + + this.heap[entry.heapIndex] = entry; + } + + /** Moves `entry` *down* to the correct position within the heap. */ + private siftDown(entry: Entry) { + let childEntry: Entry | undefined; + + do { + const childHeapIndex = childIndexOf(entry.heapIndex); + const cEnt1 = this.heap[childHeapIndex]; + const cEnt2 = this.heap[childHeapIndex + 1]; + const use2 = cEnt1 !== undefined && cEnt2 !== undefined && this.compare(cEnt2.priority, cEnt1.priority) < 0; + childEntry = use2 ? cEnt2 : cEnt1; + } while (this.trySwap(entry, childEntry, -1)); + + this.heap[entry.heapIndex] = entry; + } + + private removeEntry(entry?: Entry) { + if (entry === undefined) { + return; + } + + this.keys.delete(entry.value); + const siftEntry = this.heap.pop(); + + if (siftEntry !== undefined && siftEntry !== entry) { + this.heap[entry.heapIndex] = siftEntry; + siftEntry.heapIndex = entry.heapIndex; + if (this.compare(entry.priority, siftEntry.priority) < 0) { + this.siftDown(siftEntry); + } else { + this.siftUp(siftEntry); + } + } + } + + /** Updates the priority of value `value` to `newPriority`, or returns `false` if no such value is in the queue. */ + update(value: V, newPriority: P): boolean { + const entry = this.keys.get(value); + if (entry === undefined) { + return false; + } + + const priorityIncreased = this.compare(newPriority, entry.priority) < 0; + entry.priority = newPriority; + if (priorityIncreased) { + this.siftUp(entry); + } else { + this.siftDown(entry); + } + + return true; + } + + /** Adds a new `value` to the queue with the given `priority`. */ + insert(value: V, priority: P) { + if (this.update(value, priority)) { + return; + } + + const entry = { priority, value, heapIndex: this.heap.length }; + this.keys.set(value, entry); + this.heap.push(entry); + this.siftUp(entry); + } + + /** Returns the highest-priority item in the queue, without removing it. */ + peek(): [P, V] | undefined { + return entryToTuple(this.heap[0]); + } + + /** Removes and returns the highest-priority item in the queue. */ + pop(): [P, V] | undefined { + const entry = this.heap[0]; + this.removeEntry(entry); + return entryToTuple(entry); + } + + /** Removes the item with value `value` from the queue. */ + remove(value: V): P | undefined { + const entry = this.keys.get(value); + this.removeEntry(entry); + return entry?.priority; + } +} diff --git a/src/test/priority_queue.test.ts b/src/test/priority_queue.test.ts new file mode 100644 index 00000000..cafa0ce1 --- /dev/null +++ b/src/test/priority_queue.test.ts @@ -0,0 +1,71 @@ +import { PriorityQueue } from "../data_manager/priority_queue.js"; + +const ENTRY_COUNT = 100; + +const setup = (): [PriorityQueue, Map, number[]] => { + const queue = new PriorityQueue((i, j) => j - i); + const entryMap = new Map(); + const priorities: number[] = []; + + for (let i = 0; i < ENTRY_COUNT; i++) { + const priority = Math.random(); + const key = i.toString(); + queue.insert(key, priority); + entryMap.set(priority, key); + priorities.push(priority); + } + + return [queue, entryMap, priorities]; +}; + +const sort = (priorities: number[]) => priorities.sort((a, b) => b - a); + +describe("PriorityQueue", () => { + it("pops items in priority order", () => { + const [queue, entryMap, priorities] = setup(); + sort(priorities); + + for (const priority of priorities) { + expect(queue.pop()?.[1]).to.equal(entryMap.get(priority)); + } + expect(queue.pop()).to.equal(undefined); + }); + + it("keeps items in order after some have been removed", () => { + const REMOVE_COUNT = 10; + const [queue, entryMap, priorities] = setup(); + + const removed = priorities.slice(0, REMOVE_COUNT); + const rest = priorities.slice(REMOVE_COUNT); + + for (const p of removed) { + const key = entryMap.get(p) as string; + const removedPriority = queue.remove(key); + expect(removedPriority).to.equal(p); + } + + sort(rest); + for (const p of rest) { + expect(queue.pop()?.[1]).to.equal(entryMap.get(p)); + } + expect(queue.pop()).to.equal(undefined); + }); + + it("sorts items correctly when their priorities change", () => { + const [queue, entryMap, priorities] = setup(); + + const newPriorities = priorities.map((_, i) => { + const key = i.toString(); + const newPriority = Math.random(); + queue.update(key, newPriority); + entryMap.set(newPriority, key); + return newPriority; + }); + + sort(newPriorities); + for (const p of newPriorities) { + expect(queue.pop()?.[1]).to.equal(entryMap.get(p)); + } + expect(queue.pop()).to.equal(undefined); + }); +});