@antv/layout/src/algorithm/antv-dagre/index.ts

graph layout algorithm

import { isNumber } from '@antv/util';
import type { NodeData, PointObject } from '../../types';
import { parsePoint } from '../../util';
import { formatNodeSizeFn, formatNumberFn } from '../../util/format';
import { BaseLayout } from '../base-layout';
import { DagreGraph, GraphNode } from './graph';
import { layout } from './layout';
import { AntVDagreLayoutOptions } from './types';

export type { AntVDagreLayoutOptions };

const DEFAULTS_LAYOUT_OPTIONS: Partial<AntVDagreLayoutOptions> = {
  nodeSize: 10,
  nodeSpacing: 0,
  rankdir: 'TB',
  nodesep: 50, // 节点水平间距(px)
  ranksep: 50, // 每一层节点之间间距
  edgeLabelSpace: true,
  ranker: 'tight-tree',
  controlPoints: false, // 是否保留布局连线的控制点
  radial: false, // 是否基于 dagre 进行辐射布局
  focusNode: null, // radial 为 true 时生效，关注的节点
};

/**
 * <zh/> AntV 实现的 Dagre 布局
 *
 * <en/> AntV implementation of Dagre layout
 */
export class AntVDagreLayout extends BaseLayout<AntVDagreLayoutOptions> {
  id = 'antv-dagre';

  protected getDefaultOptions(): AntVDagreLayoutOptions {
    return DEFAULTS_LAYOUT_OPTIONS;
  }

  protected async layout(options: AntVDagreLayoutOptions): Promise<void> {
    const {
      nodeSize,
      nodeSpacing,
      align,
      rankdir = 'TB',
      ranksep,
      nodesep,
      edgeLabelSpace,
      ranker = 'tight-tree',
      nodeOrder,
      begin,
      controlPoints,
      radial,
      sortByCombo,
      // focusNode,
      preset,
      ranksepFunc,
      nodesepFunc,
    } = options;

    const ranksepfunc = formatNumberFn(ranksepFunc, ranksep ?? 50, 'node');
    const nodesepfunc = formatNumberFn(nodesepFunc, nodesep ?? 50, 'node');
    let horisep: (node: NodeData) => number = nodesepfunc;
    let vertisep: (node: NodeData) => number = ranksepfunc;

    if (rankdir === 'LR' || rankdir === 'RL') {
      horisep = ranksepfunc;
      vertisep = nodesepfunc;
    }

    // Create internal graph
    const g = new DagreGraph<NodeData, any>({ tree: [] });

    // copy graph to g
    const nodes = this.model.nodes();
    const edges = this.model.edges();

    const sizeFn = formatNodeSizeFn(
      nodeSize,
      nodeSpacing,
      DEFAULTS_LAYOUT_OPTIONS.nodeSize as number,
      DEFAULTS_LAYOUT_OPTIONS.nodeSpacing as number,
    );

    nodes.forEach((node) => {
      const raw = node._original;
      const size = sizeFn(raw);
      const verti = vertisep(raw);
      const hori = horisep(raw);
      const width = size[0] + 2 * hori;
      const height = size[1] + 2 * verti;
      const layer = node.data?.layer;
      if (isNumber(layer)) {
        // 如果有layer属性，加入到node的label中
        g.addNode({
          id: node.id,
          data: {
            width,
            height,
            layer,
            originalWidth: size[0],
            originalHeight: size[1],
          },
        });
      } else {
        g.addNode({
          id: node.id,
          data: {
            width,
            height,
            originalWidth: size[0],
            originalHeight: size[1],
          },
        });
      }
    });

    edges.forEach((edge) => {
      g.addEdge({
        id: edge.id,
        source: edge.source,
        target: edge.target,
        data: {},
      });
    });

    if (sortByCombo) {
      g.attachTreeStructure('combo');
      nodes.forEach((node) => {
        const parentId = node?.parentId;
        if (parentId === undefined) return;
        if (g.hasNode(parentId as any)) {
          g.setParent(node.id, parentId as any, 'combo');
        }
      });
    }

    let prevGraph: DagreGraph | null = null;
    if (preset?.length) {
      prevGraph = new DagreGraph<NodeData, any>();
      preset.forEach((node) => {
        prevGraph!.addNode({
          id: node.id,
          data: node.data,
        });
      });
    }

    layout(g, {
      prevGraph,
      edgeLabelSpace,
      keepNodeOrder: !!nodeOrder,
      nodeOrder: nodeOrder || [],
      acyclicer: 'greedy',
      ranker,
      rankdir,
      nodesep,
      align,
    });

    const layoutTopLeft = [0, 0];
    if (begin) {
      let minX = Infinity;
      let minY = Infinity;
      g.getAllNodes().forEach((node) => {
        if (minX > node.data.x!) minX = node.data.x!;
        if (minY > node.data.y!) minY = node.data.y!;
      });
      g.getAllEdges().forEach((edge) => {
        edge.data.points?.forEach((point: PointObject) => {
          if (minX > point.x) minX = point.x;
          if (minY > point.y) minY = point.y;
        });
      });
      layoutTopLeft[0] = begin[0] - minX;
      layoutTopLeft[1] = begin[1] - minY;
    }

    const isHorizontal = rankdir === 'LR' || rankdir === 'RL';
    if (radial) {
      // const focusId = (isString(focusNode) ? focusNode : focusNode?.id) as ID;
      // const focusLayer = focusId ? g.getNode(focusId)?.data._rank as number : 0;
      // const layers: any[] = [];
      // const dim = isHorizontal ? "y" : "x";
      // const sizeDim = isHorizontal ? "height" : "width";
      // // 找到整个图作为环的坐标维度（dim）的最大、最小值，考虑节点宽度
      // let min = Infinity;
      // let max = -Infinity;
      // g.getAllNodes().forEach((node) => {
      //   const currentNodesep = nodesepfunc(node);
      //   if (focusLayer === 0) {
      //     if (!layers[node.data._rank!]) {
      //       layers[node.data._rank!] = {
      //         nodes: [],
      //         totalWidth: 0,
      //         maxSize: -Infinity,
      //       };
      //     }
      //     layers[node.data._rank!].nodes.push(node);
      //     layers[node.data._rank!].totalWidth += currentNodesep * 2 + node.data[sizeDim]!;
      //     if (
      //       layers[node.data._rank!].maxSize < Math.max(node.data.width!, node.data.height!)
      //     ) {
      //       layers[node.data._rank!].maxSize = Math.max(node.data.width!, node.data.height!);
      //     }
      //   } else {
      //     const diffLayer = node.data._rank! - focusLayer!;
      //     if (diffLayer === 0) {
      //       if (!layers[diffLayer]) {
      //         layers[diffLayer] = {
      //           nodes: [],
      //           totalWidth: 0,
      //           maxSize: -Infinity,
      //         };
      //       }
      //       layers[diffLayer].nodes.push(node);
      //       layers[diffLayer].totalWidth += currentNodesep * 2 + node.data[sizeDim]!;
      //       if (
      //         layers[diffLayer].maxSize < Math.max(node.data.width!, node.data.height!)
      //       ) {
      //         layers[diffLayer].maxSize = Math.max(node.data.width!, node.data.height!);
      //       }
      //     } else {
      //       const diffLayerAbs = Math.abs(diffLayer);
      //       if (!layers[diffLayerAbs]) {
      //         layers[diffLayerAbs] = {
      //           left: [],
      //           right: [],
      //           totalWidth: 0,
      //           maxSize: -Infinity,
      //         };
      //       }
      //       layers[diffLayerAbs].totalWidth +=
      //         currentNodesep * 2 + node.data[sizeDim]!;
      //       if (
      //         layers[diffLayerAbs].maxSize < Math.max(node.data.width!, node.data.height!)
      //       ) {
      //         layers[diffLayerAbs].maxSize = Math.max(
      //           node.data.width!,
      //           node.data.height!
      //         );
      //       }
      //       if (diffLayer < 0) {
      //         layers[diffLayerAbs].left.push(node);
      //       } else {
      //         layers[diffLayerAbs].right.push(node);
      //       }
      //     }
      //   }
      //   const leftPos = node.data[dim]! - node.data[sizeDim]! / 2 - currentNodesep;
      //   const rightPos = node.data[dim]! + node.data[sizeDim]! / 2 + currentNodesep;
      //   if (leftPos < min) min = leftPos;
      //   if (rightPos > max) max = rightPos;
      // });
      // // const padding = (max - min) * 0.1; // TODO
      // // 初始化为第一圈的半径，后面根据每层 ranksep 叠加
      // let radius = ranksep || 50; // TODO;
      // const radiusMap: any = {};
      // // 扩大最大最小值范围，以便为环上留出接缝处的空隙
      // const rangeLength = (max - min) / 0.9;
      // const range = [
      //   (min + max - rangeLength) * 0.5,
      //   (min + max + rangeLength) * 0.5,
      // ];
      // // 根据半径、分布比例，计算节点在环上的位置，并返回该组节点中最大的 ranksep 值
      // const processNodes = (
      //   layerNodes: any,
      //   radius: number,
      //   propsMaxRanksep = -Infinity,
      //   arcRange = [0, 1]
      // ) => {
      //   let maxRanksep = propsMaxRanksep;
      //   layerNodes.forEach((node: any) => {
      //     const coord = g.node(node);
      //     radiusMap[node] = radius;
      //     // 获取变形为 radial 后的直角坐标系坐标
      //     const { x: newX, y: newY } = getRadialPos(
      //       coord![dim]!,
      //       range,
      //       rangeLength,
      //       radius,
      //       arcRange
      //     );
      //     // 将新坐标写入源数据
      //     const i = nodes.findIndex((it) => it.id === node);
      //     if (!nodes[i]) return;
      //     nodes[i].x = newX + dBegin[0];
      //     nodes[i].y = newY + dBegin[1];
      //     // @ts-ignore: pass layer order to data for increment layout use
      //     nodes[i]._order = coord._order;
      //     // 找到本层最大的一个 ranksep，作为下一层与本层的间隙，叠加到下一层的半径上
      //     const currentNodeRanksep = ranksepfunc(nodes[i]);
      //     if (maxRanksep < currentNodeRanksep) maxRanksep = currentNodeRanksep;
      //   });
      //   return maxRanksep;
      // };
      // let isFirstLevel = true;
      // const lastLayerMaxNodeSize = 0;
      // layers.forEach((layerNodes) => {
      //   if (
      //     !layerNodes?.nodes?.length &&
      //     !layerNodes?.left?.length &&
      //     !layerNodes?.right?.length
      //   ) {
      //     return;
      //   }
      //   // 第一层只有一个节点，直接放在圆心，初始半径设定为 0
      //   if (isFirstLevel && layerNodes.nodes.length === 1) {
      //     // 将新坐标写入源数据
      //     const i = nodes.findIndex((it) => it.id === layerNodes.nodes[0]);
      //     if (i <= -1) return;
      //     nodes[i].x = dBegin[0];
      //     nodes[i].y = dBegin[1];
      //     radiusMap[layerNodes.nodes[0]] = 0;
      //     radius = ranksepfunc(nodes[i]);
      //     isFirstLevel = false;
      //     return;
      //   }
      //   // 为接缝留出空隙，半径也需要扩大
      //   radius = Math.max(radius, layerNodes.totalWidth / (2 * Math.PI)); // / 0.9;
      //   let maxRanksep = -Infinity;
      //   if (focusLayer === 0 || layerNodes.nodes?.length) {
      //     maxRanksep = processNodes(
      //       layerNodes.nodes,
      //       radius,
      //       maxRanksep,
      //       [0, 1]
      //     ); // 0.8
      //   } else {
      //     const leftRatio =
      //       layerNodes.left?.length /
      //       (layerNodes.left?.length + layerNodes.right?.length);
      //     maxRanksep = processNodes(layerNodes.left, radius, maxRanksep, [
      //       0,
      //       leftRatio,
      //     ]); // 接缝留出 0.05 的缝隙
      //     maxRanksep = processNodes(layerNodes.right, radius, maxRanksep, [
      //       leftRatio + 0.05,
      //       1,
      //     ]); // 接缝留出 0.05 的缝隙
      //   }
      //   radius += maxRanksep;
      //   isFirstLevel = false;
      //   lastLayerMaxNodeSize - layerNodes.maxSize;
      // });
      // g.edges().forEach((edge: any) => {
      //   const coord = g.edge(edge);
      //   const i = edges.findIndex((it) => {
      //     const source = getEdgeTerminal(it, "source");
      //     const target = getEdgeTerminal(it, "target");
      //     return source === edge.v && target === edge.w;
      //   });
      //   if (i <= -1) return;
      //   if (
      //     self.edgeLabelSpace &&
      //     self.controlPoints &&
      //     edges[i].type !== "loop"
      //   ) {
      //     const otherDim = dim === "x" ? "y" : "x";
      //     const controlPoints = coord?.points?.slice(
      //       1,
      //       coord.points.length - 1
      //     );
      //     const newControlPoints: PointObject[] = [];
      //     const sourceOtherDimValue = g.node(edge.v)?.[otherDim]!;
      //     const otherDimDist =
      //       sourceOtherDimValue - g.node(edge.w)?.[otherDim]!;
      //     const sourceRadius = radiusMap[edge.v];
      //     const radiusDist = sourceRadius - radiusMap[edge.w];
      //     controlPoints?.forEach((point: any) => {
      //       // 根据该边的起点、终点半径，及起点、终点、控制点位置关系，确定该控制点的半径
      //       const cRadius =
      //         ((point[otherDim] - sourceOtherDimValue) / otherDimDist) *
      //           radiusDist +
      //         sourceRadius;
      //       // 获取变形为 radial 后的直角坐标系坐标
      //       const newPos = getRadialPos(
      //         point[dim],
      //         range,
      //         rangeLength,
      //         cRadius
      //       );
      //       newControlPoints.push({
      //         x: newPos.x + dBegin[0],
      //         y: newPos.y + dBegin[1],
      //       });
      //     });
      //     edges[i].controlPoints = newControlPoints;
      //   }
      // });
    } else {
      const layerCoords: Set<number> = new Set();
      const isInvert = rankdir === 'BT' || rankdir === 'RL';
      const layerCoordSort = isInvert
        ? (a: number, b: number) => b - a
        : (a: number, b: number) => a - b;
      g.getAllNodes().forEach((node) => {
        // let ndata: any = this.nodeMap[node];
        // if (!ndata) {
        //   ndata = combos?.find((it) => it.id === node);
        // }
        // if (!ndata) return;
        // ndata.x = node.data.x! + dBegin[0];
        // ndata.y = node.data.y! + dBegin[1];
        // //  pass layer order to data for increment layout use
        // ndata._order = node.data._order;
        // layerCoords.add(isHorizontal ? ndata.x : ndata.y);

        node.data.x = node.data.x! + layoutTopLeft[0];
        node.data.y = node.data.y! + layoutTopLeft[1];
        layerCoords.add(isHorizontal ? node.data.x : node.data.y);
      });
      const layerCoordsArr = Array.from(layerCoords).sort(layerCoordSort);

      // pre-define the isHorizontal related functions to avoid redundant calc in interations
      const isDifferentLayer = isHorizontal
        ? (point1: PointObject, point2: PointObject) => point1.x !== point2.x
        : (point1: PointObject, point2: PointObject) => point1.y !== point2.y;
      const filterControlPointsOutOfBoundary = isHorizontal
        ? (ps: PointObject[], point1: PointObject, point2: PointObject) => {
            const max = Math.max(point1.y, point2.y);
            const min = Math.min(point1.y, point2.y);
            return ps.filter((point) => point.y <= max && point.y >= min);
          }
        : (ps: PointObject[], point1: PointObject, point2: PointObject) => {
            const max = Math.max(point1.x, point2.x);
            const min = Math.min(point1.x, point2.x);
            return ps.filter((point) => point.x <= max && point.x >= min);
          };

      g.getAllEdges().forEach((edge, i) => {
        // const i = edges.findIndex((it) => {
        //   return it.source === edge.source && it.target === edge.target;
        // });
        // if (i <= -1) return;
        if (edgeLabelSpace && controlPoints && edge.data.type !== 'loop') {
          edge.data.controlPoints = getControlPoints(
            edge.data.points?.map(({ x, y }: PointObject) => ({
              x: x + layoutTopLeft[0],
              y: y + layoutTopLeft[1],
            })),
            g.getNode(edge.source),
            g.getNode(edge.target),
            layerCoordsArr,
            isHorizontal,
            isDifferentLayer,
            filterControlPointsOutOfBoundary,
          );
        }
      });
    }

    this.model.forEachNode((node) => {
      const layoutNode = g.getNode(node.id);

      if (layoutNode) {
        const { x, y, width, height, originalWidth, originalHeight } =
          layoutNode.data;

        const children = sortByCombo
          ? g.getChildren(node.id, 'combo')
          : g.getChildren(node.id);
        const hasChildren = children.length > 0;

        if (hasChildren) {
          let minX = Infinity,
            maxX = -Infinity;
          let minY = Infinity,
            maxY = -Infinity;

          children.forEach((child) => {
            const childId = child.id;
            const childNode = g.getNode(childId);
            if (childNode?.data) {
              const childX = childNode.data.x!;
              const childY = childNode.data.y!;
              const childWidth =
                childNode.data.originalWidth || childNode.data.width || 0;
              const childHeight =
                childNode.data.originalHeight || childNode.data.height || 0;

              minX = Math.min(minX, childX - childWidth / 2);
              maxX = Math.max(maxX, childX + childWidth / 2);
              minY = Math.min(minY, childY - childHeight / 2);
              maxY = Math.max(maxY, childY + childHeight / 2);
            }
          });

          const padding = 20;
          const groupWidth = (maxX - minX || 0) + padding * 2;
          const groupHeight = (maxY - minY || 0) + padding * 2;

          node.x = (minX + maxX) / 2;
          node.y = (minY + maxY) / 2;
          node.size = [groupWidth, groupHeight];
        } else {
          node.x = x;
          node.y = y;
          node.size = [originalWidth, originalHeight];
        }
      }
    });

    this.model.forEachEdge((edge) => {
      const layoutEdge = g.getEdge(edge.id);
      if (layoutEdge && layoutEdge.data.controlPoints) {
        edge.points = layoutEdge.data.controlPoints.map(parsePoint);
      }
    });
  }
}

/**
 * Format controlPoints to avoid polylines crossing nodes
 * @param points
 * @param sourceNode
 * @param targetNode
 * @param layerCoordsArr
 * @param isHorizontal
 * @returns
 */
const getControlPoints = (
  points: PointObject[] | undefined,
  sourceNode: GraphNode,
  targetNode: GraphNode,
  layerCoordsArr: number[],
  isHorizontal: boolean,
  isDifferentLayer: (point1: PointObject, point2: PointObject) => boolean,
  filterControlPointsOutOfBoundary: (
    ps: PointObject[],
    point1: PointObject,
    point2: PointObject,
  ) => PointObject[],
) => {
  let controlPoints = points?.slice(1, points.length - 1) || []; // 去掉头尾
  // 酌情增加控制点，使折线不穿过跨层的节点
  if (sourceNode && targetNode) {
    let { x: sourceX, y: sourceY } = sourceNode.data;
    let { x: targetX, y: targetY } = targetNode.data;
    if (isHorizontal) {
      sourceX = sourceNode.data.y;
      sourceY = sourceNode.data.x;
      targetX = targetNode.data.y;
      targetY = targetNode.data.x;
    }
    // 为跨层级的边增加第一个控制点。忽略垂直的/横向的边。
    // 新控制点 = {
    //   x: 终点x,
    //   y: (起点y + 下一层y) / 2,   #下一层y可能不等于终点y
    // }
    if (targetY !== sourceY && sourceX !== targetX) {
      const sourceLayer = layerCoordsArr.indexOf(sourceY!);
      const sourceNextLayerCoord = layerCoordsArr[sourceLayer + 1];
      if (sourceNextLayerCoord) {
        const firstControlPoint = controlPoints[0];
        const insertStartControlPoint = (
          isHorizontal
            ? {
                x: (sourceY! + sourceNextLayerCoord) / 2,
                y: firstControlPoint?.y || targetX,
              }
            : {
                x: firstControlPoint?.x || targetX,
                y: (sourceY! + sourceNextLayerCoord) / 2,
              }
        ) as PointObject;
        // 当新增的控制点不存在（!=当前第一个控制点）时添加
        if (
          !firstControlPoint ||
          isDifferentLayer(firstControlPoint, insertStartControlPoint)
        ) {
          controlPoints.unshift(insertStartControlPoint);
        }
      }

      const targetLayer = layerCoordsArr.indexOf(targetY!);
      const layerDiff = Math.abs(targetLayer - sourceLayer);
      if (layerDiff === 1) {
        controlPoints = filterControlPointsOutOfBoundary(
          controlPoints,
          sourceNode.data as PointObject,
          targetNode.data as PointObject,
        );
        // one controlPoint at least
        if (!controlPoints.length) {
          controlPoints.push(
            (isHorizontal
              ? {
                  x: (sourceY! + targetY!) / 2,
                  y: sourceX,
                }
              : {
                  x: sourceX,
                  y: (sourceY! + targetY!) / 2,
                }) as PointObject,
          );
        }
      } else if (layerDiff > 1) {
        const targetLastLayerCoord = layerCoordsArr[targetLayer - 1];
        if (targetLastLayerCoord) {
          const lastControlPoints = controlPoints[controlPoints.length - 1];
          const insertEndControlPoint = (
            isHorizontal
              ? {
                  x: (targetY! + targetLastLayerCoord) / 2,
                  y: lastControlPoints?.y || targetX,
                }
              : {
                  x: lastControlPoints?.x || sourceX,
                  y: (targetY! + targetLastLayerCoord) / 2,
                }
          ) as PointObject;
          // 当新增的控制点不存在（!=当前最后一个控制点）时添加
          if (
            !lastControlPoints ||
            isDifferentLayer(lastControlPoints, insertEndControlPoint)
          ) {
            controlPoints.push(insertEndControlPoint);
          }
        }
      }
    }
  }
  return controlPoints;
};