dao-governance-framework/forum-network/src/classes/supporting/wdg.js

import { Vertex } from './vertex.js';
import { Edge } from './edge.js';

const graphs = [];

const makeWDGHandler = (graph) => (vertexId) => {
  // We want a document for editing this node, which may be a vertex or an edge
  const editorDoc = graph.scene.getDocument('editorDocument')
    ?? graph.scene.withDocument('editorDocument').lastDocument;
  if (vertexId.startsWith('edge:')) {
    const [, from, to] = vertexId.split(':');
    Edge.prepareEditorDocument(graph, editorDoc, from, to);
  } else {
    Vertex.prepareEditorDocument(graph, editorDoc, vertexId);
  }
};

export class WeightedDirectedGraph {
  constructor(scene, options = {}) {
    this.scene = scene;
    this.vertices = new Map();
    this.edgeTypes = new Map();
    this.nextVertexId = 0;
    this.flowchart = scene?.flowchart;
    this.editable = options.editable;
    this.index = graphs.length;
    graphs.push(this);
    // Mermaid supports a click callback, but we can't customize arguments; we just get the vertex ID.
    // In order to provide the appropriate graph context for each callback, we create a separate callback
    // function for each graph.
    window[`WDGHandler${this.index}`] = makeWDGHandler(this);
  }

  withFlowchart() {
    this.scene?.withAdditionalFlowchart();
    this.flowchart = this.scene?.lastFlowchart();
    return this;
  }

  addVertex(type, id, data, label, options) {
    // Supports simple case of auto-incremented numeric ids
    if (typeof id === 'object') {
      data = id;
      id = this.nextVertexId++;
    }
    id = (typeof id === 'number') ? id.toString(10) : id;
    if (this.vertices.has(id)) {
      throw new Error(`Vertex already exists with id: ${id}`);
    }
    const vertex = new Vertex(this, type, id, data, options);
    this.vertices.set(id, vertex);
    vertex.setDisplayLabel(label ?? id);
    return vertex;
  }

  getVertex(id) {
    id = (typeof id === 'number') ? id.toString(10) : id;
    return this.vertices.get(id);
  }

  getVertexData(id) {
    return this.getVertex(id)?.data;
  }

  getVerticesData() {
    return Array.from(this.vertices.values()).map(({ data }) => data);
  }

  getEdge(type, from, to) {
    from = from instanceof Vertex ? from : this.getVertex(from);
    to = to instanceof Vertex ? to : this.getVertex(to);
    if (!from || !to) {
      return undefined;
    }
    const edges = this.edgeTypes.get(type);
    const edgeKey = Edge.getKey({ from, to, type });
    return edges?.get(edgeKey);
  }

  deleteEdge(type, from, to) {
    from = from instanceof Vertex ? from : this.getVertex(from);
    to = to instanceof Vertex ? to : this.getVertex(to);
    const edges = this.edgeTypes.get(type);
    const edgeKey = Edge.getKey({ type, from, to });
    if (!edges) return;
    const edge = edges.get(edgeKey);
    if (!edge) return;
    to.edges.from.forEach((x, i) => (x === edge) && to.edges.from.splice(i, 1));
    from.edges.to.forEach((x, i) => (x === edge) && from.edges.to.splice(i, 1));
    edges.delete(edgeKey);
  }

  getEdgeWeight(type, from, to) {
    return this.getEdge(type, from, to)?.weight;
  }

  setEdgeWeight(type, from, to, weight, data, options) {
    from = from instanceof Vertex ? from : this.getVertex(from);
    to = to instanceof Vertex ? to : this.getVertex(to);
    const edge = new Edge(this, type, from, to, weight, data, options);
    let edges = this.edgeTypes.get(type);
    if (!edges) {
      edges = new Map();
      this.edgeTypes.set(type, edges);
    }
    const edgeKey = Edge.getKey(edge);
    edges.set(edgeKey, edge);
    edge.displayEdgeNode();
    return edge;
  }

  addEdge(type, from, to, weight, data, options) {
    from = from instanceof Vertex ? from : this.getVertex(from);
    to = to instanceof Vertex ? to : this.getVertex(to);
    if (this.getEdge(type, from, to)) {
      throw new Error(`Edge ${type} from ${from.id} to ${to.id} already exists`);
    }
    const edge = this.setEdgeWeight(type, from, to, weight, data, options);
    from.edges.from.push(edge);
    to.edges.to.push(edge);
    edge.displayEdge();
    return edge;
  }

  getEdges(type, from, to) {
    from = from instanceof Vertex ? from : this.getVertex(from);
    to = to instanceof Vertex ? to : this.getVertex(to);
    const edgeTypes = type ? [type] : Array.from(this.edgeTypes.keys());
    return edgeTypes.flatMap((edgeType) => {
      const edges = this.edgeTypes.get(edgeType);
      return Array.from(edges?.values() || []).filter((edge) => {
        const matchFrom = from === null || from === undefined || from === edge.from;
        const matchTo = to === null || to === undefined || to === edge.to;
        return matchFrom && matchTo;
      });
    });
  }

  countVertices(type) {
    if (!type) {
      return this.vertices.size;
    }
    return Array.from(this.vertices.values()).filter((vertex) => vertex.type === type).length;
  }
}