Red-black trees in Typescript (no deletion yet)

tree-map.ts

type Color = 'r' | 'b';
type NullableNode<K, V> = Node<K, V> | null;
type Node<K, V> = {
  value: V;
  color: Color;
  kids: [NullableNode<K, V>, NullableNode<K, V>];
};

export class TreeMap<K, V> {
  private root: NullableNode<K, V> = null;
  private size: number = 0;

  constructor(
    private readonly cmpFn: (a: K, b: K) => number,
    private readonly keyFn: (value: V) => K,
  ) {}

  public find(key: K): V | undefined {
    let p: Node<K, V> | null = this.root;
    while (p) {
      const cmp = this.cmpFn(key, this.keyFn(p.value));
      if (cmp === 0) {
        return p.value;
      }
      p = p.kids[+(cmp > 0)];
    }
    return undefined;
  }

  /** Inserts given value and returns true unless a value with the same key already exists. */
  public insert(value: V): boolean {
    if (this.root === null) {
      this.root = { value, color: 'b', kids: [null, null] };
      this.size = 1;
      return true;
    }

    // Search for insert position. Build a stack of nodes and directions taken.
    const key = this.keyFn(value);
    const stack: [Node<K, V>, number][] = [];
    let lastSearchPtr = null;
    let searchPtr: NullableNode<K, V> = this.root;
    let searchDir: number = NaN;
    do {
      const cmp = this.cmpFn(key, this.keyFn(searchPtr.value));
      if (cmp === 0) {
        return false;
      }
      searchDir = +(cmp > 0);
      stack.push([searchPtr, searchDir]);
      lastSearchPtr = searchPtr;
      searchPtr = searchPtr.kids[searchDir];
    } while (searchPtr);

    // Insert new node.
    let child: Node<K, V> = { value, color: 'r', kids: [null, null] };
    lastSearchPtr.kids[searchDir] = child;
    this.size++;

    // Rebalance.
    while (stack.length >= 2) {
      const [parent, parentDir] = stack.pop()!;
      if (parent.color === 'b') {
        break;
      }
      const [grandparent, grandparentDir] = stack.pop()!;
      const oppositeGrandparentDir = 1 - grandparentDir;
      const parentSibling = grandparent.kids[oppositeGrandparentDir];
      if (parentSibling === null || parentSibling.color === 'b') {
        // Do canonical rotations and return. Swap value references to avoid
        // mutating the great-grandparent's kids, where the root is a special case.
        if (parentDir === grandparentDir) {
          [grandparent.value, parent.value] = [parent.value, grandparent.value];
          grandparent.kids[grandparentDir] = child;
          grandparent.kids[oppositeGrandparentDir] = parent;
          parent.kids[grandparentDir] = parent.kids[oppositeGrandparentDir];
          parent.kids[oppositeGrandparentDir] = parentSibling;
        } else {
          [grandparent.value, child.value] = [child.value, grandparent.value];
          grandparent.kids[oppositeGrandparentDir] = child;
          parent.kids[oppositeGrandparentDir] = child.kids[grandparentDir];
          child.kids[grandparentDir] = child.kids[oppositeGrandparentDir];
          child.kids[oppositeGrandparentDir] = parentSibling;
        }
        break;
      }
      // Do re-coloring and recur for the grandparent.
      parent.color = parentSibling.color = 'b';
      if (grandparent !== this.root) {
        grandparent.color = 'r';
      }
      child = grandparent;
    }
    return true;
  }

NOTE: No parent pointers! Follows the pattern of RB trees in libavl by using a stack. But libavl RB trees have at least one bad bug. This passes pretty thorough unit tests including a R-B rule validator. Using 2-element arrays for child references reduces code bulk for re-balancing, perhaps at a slight runtime penalty for the extra indexing.