shalu-system-designer/apps/designer/src/utils/hierarchy/treeDiagram.ts

import { TopicType } from "@/enum";
import { MindMapProjectInfo, TopicItem, HierarchyResult } from "@/types";

const DISTANCE = 20; // 子节点相对于父节点的偏差
function firstWalk(node: TopicItem): HierarchyResult {
  const root: HierarchyResult = {
    x: 0,
    y: 0,
    totalHeight: 0,
    totalWidth: 0,
    width: node.width,
    height: node.height,
    id: node.id,
    data: node,
    children: [],
  };

  root.children = node.collapsed ? [] : node.children?.map((child) => firstWalk(child)) || [];

  return root;
}

// 设置宽高
function secondWalk(
  node: HierarchyResult,
  pageSetting: MindMapProjectInfo["pageSetting"]
) {
  let totalHeight = node.height;
  let totalWidth = node.data.type === TopicType.main ? 0 : node.width;
  if (node.children?.length) {
    node.children.forEach((c: HierarchyResult, index) => {
      const childSize = secondWalk(c, pageSetting);
      // const offsetY = node.data.type === TopicType.main ? pageSetting.branchY :  pageSetting.subTopicY;
      // 总的高度 = 自身高度 + 子节点的高度 + 节点之间的距离
      totalHeight += childSize.height + pageSetting.branchY;
      
      if (node.data.type === TopicType.main) {
        const offsetX = index < node.children.length - 1 ? pageSetting.branchX : 0;
        totalWidth += childSize.width + offsetX;
      } else {
        // 总的宽度 = 子节点的宽度与自身宽度取最大值
        totalWidth = Math.max(DISTANCE + childSize.width, totalWidth);
      }
    });
  }
  node.totalHeight = totalHeight;
  node.totalWidth = totalWidth;

  return {
    height: totalHeight,
    width: Math.max(totalWidth, node.width),
  };
}

function thirdWalk(node: HierarchyResult, pageSetting: MindMapProjectInfo['pageSetting']) {
  let startX = node.width / 2 - node.totalWidth / 2;
  node.children.forEach((child) => {
    child.y = node.y + node.height + pageSetting.branchY;
    child.x = startX;
    startX += child.totalWidth + pageSetting.branchX;
    fourthWalk(child, pageSetting);
  });
}

function fourthWalk(node: HierarchyResult, pageSetting: MindMapProjectInfo['pageSetting']) {
  let startY = node.y + node.height + pageSetting.branchY;
  node.children.forEach((child) => {
    child.y = startY;
    child.x = node.x + DISTANCE;
    startY += child.totalHeight + pageSetting.branchY;
    if(child.children.length) {
      fourthWalk(child, pageSetting);
    }
  });
}

export const getTreeDiagramHierarchy = <K>(
  topic: TopicItem,
  pageSetting: MindMapProjectInfo["pageSetting"]
): K => {
  const root = firstWalk(topic);
  secondWalk(root, pageSetting);
  thirdWalk(root, pageSetting);
  console.log(root);
  return root as K;
};