flowgraph/src/graph/construction.ts

import { DirectedGraph } from "graphology";
import type { TSchema, Static } from "@alkdev/typebox";
import { willCreateCycle } from "graphology-dag";
import {
  DuplicateNodeError,
  DuplicateEdgeError,
  NodeNotFoundError,
  CycleError,
} from "../error/index.js";

export interface FlowGraphOptions {
  type?: "directed";
  multi?: false;
  allowSelfLoops?: false;
}

type Attrs = Record<string, unknown>;

export class FlowGraph<
  NodeAttrs extends TSchema = TSchema,
  EdgeAttrs extends TSchema = TSchema,
> {
  private _graph: DirectedGraph;

  constructor(_options?: FlowGraphOptions) {
    this._graph = new DirectedGraph({
      type: "directed",
      multi: false,
      allowSelfLoops: false,
    });
  }

  get graph(): DirectedGraph {
    return this._graph;
  }

  _edgeKey(source: string, target: string): string {
    return `${source}->${target}`;
  }

  addNode(key: string, attrs: Static<NodeAttrs>): void {
    if (this._graph.hasNode(key)) {
      throw new DuplicateNodeError(key);
    }
    this._graph.addNode(key, attrs as Attrs);
  }

  removeNode(key: string): void {
    if (!this._graph.hasNode(key)) {
      throw new NodeNotFoundError(key);
    }
    this._graph.dropNode(key);
  }

  updateNode(key: string, attrs: Partial<Static<NodeAttrs>>): void {
    if (!this._graph.hasNode(key)) {
      throw new NodeNotFoundError(key);
    }
    this._graph.mergeNodeAttributes(key, attrs as Attrs);
  }

  hasNode(key: string): boolean {
    return this._graph.hasNode(key);
  }

  getNodeAttributes(key: string): Static<NodeAttrs> {
    if (!this._graph.hasNode(key)) {
      throw new NodeNotFoundError(key);
    }
    return this._graph.getNodeAttributes(key) as Static<NodeAttrs>;
  }

  addEdge(source: string, target: string, attrs?: Static<EdgeAttrs>): void {
    if (!this._graph.hasNode(source)) {
      throw new NodeNotFoundError(source);
    }
    if (!this._graph.hasNode(target)) {
      throw new NodeNotFoundError(target);
    }
    if (this._graph.hasDirectedEdge(source, target)) {
      throw new DuplicateEdgeError(source, target);
    }
    if (willCreateCycle(this._graph, source, target)) {
      const path = this._findPath(target, source);
      const cycle = [source, ...path, source];
      throw new CycleError([cycle]);
    }
    const key = this._edgeKey(source, target);
    this._graph.addEdgeWithKey(key, source, target, (attrs ?? {}) as Attrs);
  }

  removeEdge(source: string, target: string): void {
    if (this._graph.hasDirectedEdge(source, target)) {
      this._graph.dropDirectedEdge(source, target);
    }
  }

  hasEdge(source: string, target: string): boolean {
    return this._graph.hasDirectedEdge(source, target);
  }

  getEdgeAttributes(source: string, target: string): Static<EdgeAttrs> {
    if (!this._graph.hasDirectedEdge(source, target)) {
      throw new Error(
        `Edge "${this._edgeKey(source, target)}" not found`,
      );
    }
    const edgeKey = this._graph.edge(source, target);
    return this._graph.getEdgeAttributes(edgeKey) as Static<EdgeAttrs>;
  }

  nodes(): string[] {
    return this._graph.nodes();
  }

  edges(): string[] {
    return this._graph.edges();
  }

  get order(): number {
    return this._graph.order;
  }

  get size(): number {
    return this._graph.size;
  }

  forEachNode(
    callback: (key: string, attrs: Static<NodeAttrs>) => void,
  ): void {
    this._graph.forEachNode((key, attrs) => {
      callback(key, attrs as Static<NodeAttrs>);
    });
  }

  forEachEdge(
    callback: (
      key: string,
      attrs: Static<EdgeAttrs>,
      source: string,
      target: string,
    ) => void,
  ): void {
    this._graph.forEachEdge((key, attrs, source, target) => {
      callback(key, attrs as Static<EdgeAttrs>, source!, target!);
    });
  }

  predecessors(nodeId: string): string[] {
    return this._graph.inNeighbors(nodeId) ?? [];
  }

  successors(nodeId: string): string[] {
    return this._graph.outNeighbors(nodeId) ?? [];
  }

  static fromSpecs(
    _specs: unknown[],
  ): FlowGraph<TSchema, TSchema> {
    throw new Error("not implemented");
  }

  static fromCallEvents(
    _events: unknown[],
  ): FlowGraph<TSchema, TSchema> {
    throw new Error("not implemented");
  }

  static fromJSON(
    _data: unknown,
  ): FlowGraph<TSchema, TSchema> {
    throw new Error("not implemented");
  }

  private _findPath(from: string, to: string): string[] {
    const visited = new Set<string>();
    const queue: Array<{ node: string; path: string[] }> = [
      { node: from, path: [from] },
    ];
    visited.add(from);
    while (queue.length > 0) {
      const current = queue.shift()!;
      if (current.node === to) {
        return current.path;
      }
      const neighbors = this._graph.outNeighbors(current.node) ?? [];
      for (const neighbor of neighbors) {
        if (!visited.has(neighbor)) {
          visited.add(neighbor);
          queue.push({
            node: neighbor,
            path: [...current.path, neighbor],
          });
        }
      }
    }
    return [];
  }
}