digraph

directed_graph.ts

class DirectedGraph<T> {
  readonly edges: Map<T, Set<T>> = new Map();

  addAll(from: T, ...to: T[]) {
    let dependencies = this.edges.get(from);
    if (dependencies == null) {
      dependencies = new Set();
      this.edges.set(from, dependencies);
    }
    Sets.addAll(dependencies, to);
  }

  isCyclic(): boolean {
    const seenOnAllWalks = new Set<T>();
    for (const node of this.edges.keys()) {
      if (seenOnAllWalks.has(node)) {
        continue;
      }

      const seenOnThisWalk = new Set<T>();
      const toVisit = [...this.edges.get(node)!];
      while (toVisit.length > 0) {
        const nextNode = toVisit.shift()!;
        if (seenOnThisWalk.has(nextNode)) {
          return true; // cyclic
        }
        seenOnThisWalk.add(nextNode);
        const nextNodeChildren = this.edges.get(nextNode);
        nextNodeChildren && toVisit.push(...nextNodeChildren);
      }

      Sets.addAll(seenOnAllWalks, seenOnThisWalk);
    }

    return false;
  }

  indegrees() {
    const inDegrees = new Map<T, number>();
    for (const [node, neighbours] of this.edges.entries()) {
      if (!inDegrees.has(node)) {
        inDegrees.set(node, 0);
      }

      for (const neighbour of neighbours) {
        const count = inDegrees.get(neighbour) || 0;
        inDegrees.set(neighbour, count + 1);
      }
    }
    return inDegrees;
  }

  topoSort(): readonly T[] {
    const inDegrees = this.indegrees();
    const sources: T[] = [];
    for (const [node, count] of inDegrees.entries()) {
      if (count === 0) {
        sources.push(node);
      }
    }

    assert.strict(
      sources.length > 0,
      `a DAG must have at least one source (a node with an in-degree of 0)`
    );

    const topologicalOrdering = [];
    while (sources.length > 0) {
      const node = sources.pop()!;
      topologicalOrdering.push(node);
      const neighbours = this.edges.get(node) || new Set();
      for (const neighbour of neighbours) {
        const neighbourIndegree = inDegrees.get(neighbour)! - 1;
        inDegrees.set(neighbour, neighbourIndegree);
        if (neighbourIndegree === 0) {
          sources.push(neighbour);
        }
      }
    }

    assert.strict(
      topologicalOrdering.length === this.edges.size,
      `Graph has a cycle! No topological ordering exists.`
    );

    return topologicalOrdering;
  }

  invert(): DirectedGraph<T> {
    const inverted = new DirectedGraph<T>();
    for (const [edge, deps] of this.edges) {
      inverted.addAll(edge);
      for (const dep of deps) {
        inverted.addAll(dep, edge);
      }
    }
    return inverted;
  }

  walk(start: T): Set<T> {
    const toVisit = [start];
    const seen = new Set<T>();
    while (toVisit.length > 0) {
      const next = toVisit.shift()!;
      for (const dep of this.edges.get(next)!) {
        if (seen.has(dep)) {
          continue;
        }
        toVisit.push(dep);
      }
      seen.add(next);
    }
    return seen;
  }

  subgraph(keep: Set<T>): DirectedGraph<T> {
    const subgraph = new DirectedGraph<T>();
    for (const [node, deps] of this.edges) {
      if (!keep.has(node)) {
        continue;
      }
      subgraph.addAll(node, ...deps);
    }
    return subgraph;
  }
}

class Sets {
  static addAll<T>(s: Set<T>, xs: Iterable<T>) {
    for (const x of xs) {
      s.add(x);
    }
  }
}