spica/queue.ts

Supervisor, Coroutine, Channel, select, AtomicPromise, Cancellation, Cache, List, Queue, Stack, and some utils.

import { Heap } from './heap';
import { memoize } from './memoize';
import { IterableDict } from './dict';

const size = 2048;
const initsize = 16;

export class Queue<T> {
  private head = new FixedQueue<T>(initsize);
  private tail = this.head;
  private count = 0;
  public get length(): number {
    return this.head === this.tail
      ? this.head.length
      : this.head.length + this.count + this.tail.length;
  }
  // Faster than queue.length > 0.
  public isEmpty(): boolean {
    return this.head.isEmpty();
  }
  public peek(index: 0 | -1 = 0): T | undefined {
    return index === 0
      ? this.head.peek(0)
      : this.tail.peek(-1);
  }
  public push(value: T): void {
    const tail = this.tail;
    if (tail.isFull()) {
      if (tail.next.isEmpty()) {
        this.tail = tail.next;
      }
      else {
        this.tail = tail.next = new FixedQueue(size, tail.next);
      }
      assert(this.tail.isEmpty());
      if (this.head !== tail) {
        assert(this.head.next !== this.tail);
        this.count += tail.size;
      }
    }
    this.tail.push(value);
  }
  public pop(): T | undefined {
    const head = this.head;
    const value = head.pop();
    if (head.isEmpty() && !head.next.isEmpty()) {
      const tail = this.tail;
      // 空になるごとの削除と再作成を避ける
      if (tail.next !== head) {
        assert(tail.next.next === head);
        // 初期サイズの方を消す
        tail.next.next = tail.next;
        tail.next = head;
        assert(head.size === size);
      }
      assert(this.tail.next.isEmpty());
      this.head = head.next;
      if (this.head !== tail) {
        assert(this.head !== this.tail);
        this.count -= head.next.size;
      }
    }
    return value;
  }
  public clear(): void {
    this.head = this.tail = new FixedQueue(initsize);
    this.count = 0;
  }
  public *[Symbol.iterator](): Iterator<T, undefined, undefined> {
    while (!this.isEmpty()) {
      yield this.pop()!;
    }
  }
}

class FixedQueue<T> {
  constructor(
    public readonly size: number,
    next?: FixedQueue<T>,
  ) {
    assert((this.array.length & this.array.length - 1) === 0);
    this.next = next ?? this;
  }
  private readonly array = Array<T | undefined>(this.size);
  private readonly mask = this.array.length - 1;
  private head = 0;
  private tail = 0;
  // 1要素無駄にしフラグを使用しない場合と有意差がないため可読性とテスト性を優先しフラグを使用。
  private empty = true;
  public next: FixedQueue<T>;
  public get length(): number {
    return this.tail >= this.head
      ? this.empty ? 0 : this.tail - this.head || this.size
      : this.array.length - this.head + this.tail;
  }
  public isEmpty(): boolean {
    return this.empty;
  }
  public isFull(): boolean {
    return this.tail === this.head && !this.empty;
  }
  public peek(index: number = 0): T | undefined {
    return index >= 0
      ? this.array[this.head + index & this.mask]
      : this.array[this.tail + index & this.mask];
  }
  public push(value: T): void {
    this.array[this.tail] = value;
    this.tail = this.tail + 1 & this.mask;
    this.empty = false;
  }
  public pop(): T | undefined {
    if (this.empty) return;
    const value = this.array[this.head];
    this.array[this.head] = undefined;
    this.head = this.head + 1 & this.mask;
    // isEmptyの前倒し
    this.empty = this.tail === this.head;
    return value;
  }
}

export class PriorityQueue<T, P = T> {
  private static readonly priority = Symbol('priority');
  public static readonly max = Heap.max;
  public static readonly min = Heap.min;
  constructor(
    cmp: (a: P, b: P) => number = PriorityQueue.max,
    private clean = true,
  ) {
    this.heap = new Heap(cmp);
  }
  private readonly heap: Heap<Queue<T>, P>;
  private readonly dict = new Map<P, Queue<T>>();
  private readonly queue = memoize<P, Queue<T>>(priority => {
    const queue = new Queue<T>();
    queue[PriorityQueue.priority] = priority;
    this.heap.insert(queue, priority);
    return queue;
  }, this.dict);
  private $length = 0;
  public get length(): number {
    return this.$length;
  }
  public isEmpty(): boolean {
    return this.$length === 0;
  }
  public peek(priority?: P): T | undefined {
    return arguments.length === 0
      ? this.heap.peek()?.value.peek()
      : this.dict.get(priority!)?.peek();
  }
  public push(priority: P, value: T): void {
    ++this.$length;
    this.queue(priority).push(value);
  }
  public pop(priority?: P): T | undefined {
    if (this.$length === 0) return;
    --this.$length;
    const queue = arguments.length === 0
      ? this.heap.peek()!.value
      : this.dict.get(priority!);
    const value = queue?.pop();
    if (queue?.isEmpty()) {
      this.heap.extract();
      this.clean && this.dict.delete(queue[PriorityQueue.priority]);
    }
    return value;
  }
  public clear(): void {
    this.heap.clear();
    this.dict.clear();
    this.$length = 0;
  }
  public *[Symbol.iterator](): Iterator<T, undefined, undefined> {
    while (!this.isEmpty()) {
      yield this.pop()!;
    }
  }
}

export class MultiQueue<K, V> implements IterableDict<K, V> {
  constructor(
    entries?: Iterable<readonly [K, V]>,
  ) {
    if (entries) for (const { 0: k, 1: v } of entries) {
      this.set(k, v);
    }
  }
  private dict = new Map<K, Queue<V>>();
  public get length(): number {
    return this.dict.size;
  }
  public isEmpty(): boolean {
    return this.dict.size === 0;
  }
  public peek(key: K): V | undefined {
    return this.dict.get(key)?.peek();
  }
  public push(key: K, value: V): void {
    let vs = this.dict.get(key);
    if (vs) return void vs.push(value);
    vs = new Queue();
    vs.push(value);
    this.dict.set(key, vs);
  }
  public pop(key: K): V | undefined {
    return this.dict.get(key)?.pop();
  }
  public clear(): void {
    this.dict = new Map();
  }
  public take(key: K): V | undefined;
  public take(key: K, count: number): V[];
  public take(key: K, count?: number): V | undefined | V[] {
    if (count === undefined) return this.pop(key);
    const vs = this.dict.get(key);
    const acc: V[] = [];
    while (vs && !vs.isEmpty() && count--) {
      acc.push(vs.pop()!);
    }
    return acc;
  }
  public ref(key: K): Queue<V> {
    let vs = this.dict.get(key);
    if (vs) return vs;
    vs = new Queue();
    this.dict.set(key, vs);
    return vs;
  }
  public get size(): number {
    return this.length;
  }
  public get(key: K): V | undefined {
    return this.peek(key);
  }
  public set(key: K, value: V): this {
    this.push(key, value);
    return this;
  }
  public has(key: K): boolean {
    return this.dict.has(key);
  }
  public delete(key: K): boolean {
    return this.dict.delete(key);
  }
  public *[Symbol.iterator](): Iterator<[K, V], undefined, undefined> {
    for (const { 0: k, 1: vs } of this.dict) {
      while (!vs.isEmpty()) {
        yield [k, vs.pop()!];
      }
    }
  }
}