Chris Okasaki's Lazy Persistent Queue

queue.ts

type Assert = (condition: boolean, message: string) => asserts condition;

const assert: Assert = (condition, message) => {
  if (!condition) throw new Error(message);
};

type List<A> = Cons<A> | Empty;

interface Cons<out A> {
  type: 'Cons';
  head: A;
  tail: List<A>;
}

interface Empty {
  type: 'Empty';
}

interface RotateThunk<out A> {
  type: 'RotateThunk';
  front: List<A> | LazyCons<A>;
  back: Cons<A>;
  tail: Cons<A>;
}

interface LazyCons<out A> {
  type: 'LazyCons';
  head: A;
  tail: Cons<A> | LazyCons<A> | RotateThunk<A>;
}

const Cons = <A>(head: A, tail: List<A>): Cons<A> => ({
  type: 'Cons',
  head,
  tail,
});

const Empty: Empty = { type: 'Empty' };

const rotate = <A>(
  front: List<A> | LazyCons<A>,
  back: Cons<A>,
  tail: List<A>
): Cons<A> | LazyCons<A> => {
  const cons: Cons<A> = Cons(back.head, tail);

  if (front.type === 'Empty') {
    assert(
      back.tail.type === 'Empty',
      'Expected back to have one more item than front.'
    );

    return cons;
  }

  assert(
    front.tail.type !== 'RotateThunk',
    'Expected front to be pre-evaluated.'
  );

  assert(
    back.tail.type === 'Cons',
    'Expected back to have one more item than front.'
  );

  return {
    type: 'LazyCons',
    head: front.head,
    tail: {
      type: 'RotateThunk',
      front: front.tail,
      back: back.tail,
      tail: cons,
    },
  };
};

const getTail = <A>(cons: Cons<A> | LazyCons<A>): List<A> | LazyCons<A> => {
  if (cons.tail.type === 'RotateThunk') {
    assert(cons.type === 'LazyCons', 'Expected cons to be lazy.');
    const { front, back, tail } = cons.tail;
    cons.tail = rotate(front, back, tail);
  }

  return cons.tail;
};

const makeQueue = <A>(
  front: List<A> | LazyCons<A>,
  back: List<A>,
  tail: List<A> | LazyCons<A>
): Queue<A> => {
  if (tail.type === 'Empty') {
    assert(
      back.type === 'Cons',
      'Expected back to have one more item than front.'
    );

    const tail = rotate(front, back, Empty);
    return new NonEmptyQueue(tail, Empty, tail);
  }

  if (front.type === 'Empty') {
    assert(back.type === 'Empty', 'Expected back to be empty.');
    return new Queue();
  }

  return new NonEmptyQueue(front, back, getTail(tail));
};

export class Queue<out A> {
  protected readonly front: List<A> | LazyCons<A> = Empty;
  protected readonly back: List<A> = Empty;
  protected readonly tail: List<A> | LazyCons<A> = Empty;

  public isNonEmpty(this: Queue<A>): this is NonEmptyQueue<A> {
    return this instanceof NonEmptyQueue;
  }

  public insert(this: Queue<A>, item: A): NonEmptyQueue<A> {
    const queue = makeQueue(this.front, Cons(item, this.back), this.tail);
    assert(queue.isNonEmpty(), 'Expected queue to be non-empty.');
    return queue;
  }

  public tryRemove(): [A, Queue<A>] | undefined {
    return this.isNonEmpty() ? this.remove() : undefined;
  }
}

class NonEmptyQueue<out A> extends Queue<A> {
  public constructor(
    protected override readonly front: Cons<A> | LazyCons<A>,
    protected override readonly back: List<A>,
    protected override readonly tail: List<A> | LazyCons<A>
  ) {
    super();
  }

  public remove(): [A, Queue<A>] {
    return [
      this.front.head,
      makeQueue(getTail(this.front), this.back, this.tail),
    ];
  }
}

export type { NonEmptyQueue };