4pc/src/view/step-tree-v2/tree-helper.ts

export type Step<T> = {
  raw: T;
  parents: Step<T>[];
  children: Step<T>[];
  box: StepCtx<T>;
};
export type Steps<T> = Step<T>[];

export type StepCtx<T> = {
  level: number;
  step: Step<T>;
  // 自身大小
  size: { w: number; h: number };
  // 如果是多对一则存储多宽度
  parallelWidth?: number;
  // 树宽高
  treeSize?: { w: number; h: number };
  // 相对于兄弟
  offset: { x: number; y: number };
  treeOffset: { x: number; y: number };
  lines: number[][][];
};

export type StepsCtx<T> = {
  levelHeight: number[];
  boxs: StepCtx<T>[];
  roots: Step<T>[];
  offset: { x: number; y: number };
  size: { w: number; h: number };
};

type traversalStepsProps<T> = {
  steps: Steps<T>;
  ctx: StepsCtx<T>;
  oper: (data: {
    currents: Step<T>[];
    prev: Step<T> | null;
    next: Step<T> | null;
    levelSteps: Step<T>[];
  }) => void;
  cs?: Step<T>[];
  checkeds?: Steps<T>;
  level?: number;
  reverse?: boolean;
};

const setStepsLevel = <T>(
  steps: Steps<T>,
  ctx: StepsCtx<T>,
  level = 0,
  cs: Step<T>[] = ctx.roots
) => {
  if (cs && cs.length === 0) return;
  const cSteps = cs;

  for (let i = 0; i < cSteps.length; i++) {
    const current = cSteps[i];
    if ("level" in current.box) {
      current.box.level = Math.max(level, current.box.level);
    } else {
      (current.box as any).level = level;
    }
    ctx.levelHeight[level] = 0;
    setStepsLevel(steps, ctx, level + 1, current.children);
  }
};

export const traversalSteps = <T>({
  steps,
  ctx,
  oper,
  cs = ctx.roots,
  level = 0,
  reverse = false,
  checkeds = [],
}: traversalStepsProps<T>) => {
  if (cs.length === 0) return;
  const cSteps = cs;

  for (let i = 0; i < cSteps.length; i++) {
    const current = cSteps[i];
    if (checkeds.includes(current)) continue;
    const children = current.children.filter(
      (child) => child.box.level === level + 1
    );

    // 查看是否是多对一的情况，如果是这current为所有多的step
    let currents = [current];
    if (children.length === 1 && children[0].parents.length > 1) {
      const steps = children[0].parents;
      currents = steps.filter((step) => step.box.level === level);
    }
    checkeds.push(...currents);

    const props = {
      currents: currents,
      prev: cSteps[i - 1] || null,
      next: cSteps[i + 1] || null,
      level: level,
      levelSteps: cSteps,
    };

    reverse || oper(props);
    traversalSteps({
      steps,
      ctx,
      oper,
      cs: current.children,
      level: level + 1,
      reverse,
      checkeds,
    });
    reverse && oper(props);
  }
};

const setStepsBound = <T>(
  steps: Steps<T>,
  ctx: StepsCtx<T>,
  getStepSize: (step: T) => { w: number; h: number }
) => {
  // 注入levelHeights, box size offset
  traversalSteps({
    steps,
    ctx,
    oper: ({ currents }) => {
      currents.forEach((current) => {
        const size = getStepSize(current.raw);
        current.box.size = size;
        ctx.levelHeight[current.box.level] = Math.max(
          size.h,
          ctx.levelHeight[current.box.level] || 0
        );
      });
    },
  });
};

const setStepsTreeSize = <T>(steps: Steps<T>, ctx: StepsCtx<T>) => {
  // 从低到顶分别计算树大小
  traversalSteps({
    steps,
    ctx,
    oper: ({ currents }) => {
      let levelHeight = 0;
      const level = currents[0].box.level;
      for (let i = 0; i <= level; i++) {
        levelHeight += ctx.levelHeight[i];
      }
      const current = currents[0];
      const treeSize = { w: 0, h: 0 };
      for (const child of current.children) {
        if (child.box.treeSize) {
          treeSize.w += child.box.treeSize.w;
          treeSize.h = Math.max(child.box.treeSize.w, treeSize.w);
        }
      }
      treeSize.h += levelHeight;

      if (currents.length === 1) {
        // 一对一  一对多情况
        treeSize.w = Math.max(current.box.size.w, treeSize.w);
        current.box.treeSize = treeSize;
      } else {
        // 多对一情况
        let parallelWidth = 0;
        for (const parallel of currents) {
          parallelWidth += parallel.box.size.w;
        }
        treeSize.w = Math.max(parallelWidth, treeSize.w);
        currents[0].box.parallelWidth = parallelWidth;
        currents[0].box.treeSize = treeSize;
      }
    },
    reverse: true,
  });
};

const setStepsOffset = <T>(steps: Steps<T>, ctx: StepsCtx<T>) => {
  // 从顶到底分别计算树偏移量以及step偏移
  traversalSteps({
    steps,
    ctx,
    oper: ({ currents, prev }) => {
      const box = currents[0].box;
      let levelHeight = 0;
      const level = currents[0].box.level;
      for (let i = 0; i < level; i++) {
        levelHeight += ctx.levelHeight[i];
      }

      const treeOffset = { x: 0, y: levelHeight };

      if (prev) {
        // 上一个是普通树
        if (prev.box.treeOffset) {
          treeOffset.x += prev.box.treeOffset.x + prev.box.treeSize.w;
        } else {
          // 上一个是多对一树，需要找到box属性值
          let prevTreeBox: StepCtx<T>;
          for (const prevParent of prev.parents) {
            for (let prevLevel of prevParent.children) {
              if (prevLevel.box.treeOffset) {
                prevTreeBox = prevLevel.box;
                break;
              }
            }
            if (prevTreeBox) break;
          }
          treeOffset.x += prevTreeBox.treeOffset.x + prevTreeBox.parallelWidth;
        }
      } else if (currents[0].parents.length) {
        // 如果是第一个，则需要计算子级，相对父级做偏移
        const parents = currents[0].parents;
        const levels = new Set(parents.flatMap((parent) => parent.children));
        let levelWidth = 0;
        for (const levelStep of levels) {
          if ("parallelWidth" in levelStep.box) {
            levelWidth += levelStep.box.parallelWidth;
          } else if (levelStep.box.treeSize) {
            levelWidth += levelStep.box.treeSize.w || 0;
          }
        }

        let parentWidth = 0;
        for (const parent of parents) {
          parentWidth +=
            "parallelWidth" in parent.box
              ? parent.box.parallelWidth
              : "treeSize" in parent.box
              ? parent.box.treeSize.w
              : 0;
        }

        treeOffset.x +=
          parents[0].box.treeOffset.x + (parentWidth - levelWidth) / 2;
      }

      box.treeOffset = { ...treeOffset };
      if ("parallelWidth" in box) {
        let lOffset = 0;
        for (const current of currents) {
          current.box.offset = {
            x: treeOffset.x + lOffset,
            y: treeOffset.y + (ctx.levelHeight[level] - current.box.size.h) / 2,
          };
          lOffset += current.box.size.w;
        }
      } else {
        treeOffset.x += (box.treeSize.w - box.size.w) / 2;
        treeOffset.y += (ctx.levelHeight[level] - box.size.h) / 2;

        box.offset = treeOffset;
      }
    },
  });
};

const setStepsLines = <T>(
  steps: Steps<T>,
  ctx: StepsCtx<T>,
  margin: number[]
) => {
  traversalSteps({
    steps,
    ctx,
    oper: ({ currents }) => {
      for (const current of currents) {
        current.box.lines = current.children.map((child) =>
          getStepLine(ctx, current, child, margin)
        );
      }
    },
  });
};

export const getStepLine = <T>(
  ctx: StepsCtx<T>,
  topStep: Step<T>,
  bottomStep: Step<T>,
  margin: number[] = [0, 0]
) => {
  const top = topStep.box;
  const bottom = bottomStep.box;
  const bottomHeight = ctx.levelHeight[bottom.level];
  const start = [
    top.offset.x + top.size.w / 2,
    top.offset.y + top.size.h - margin[0],
  ];

  const diffHeight = (bottomHeight - bottom.size.h) / 2;

  const c1 = [top.offset.x + top.size.w / 2, bottom.offset.y - diffHeight];
  const c2 = [
    bottom.offset.x + bottom.size.w / 2,
    bottom.offset.y - diffHeight,
  ];
  const end = [
    bottom.offset.x + bottom.size.w / 2,
    bottom.offset.y + margin[0],
  ];
  return [start, c1, c2, end];
};

export const getStepsTreeCtx = <T>(
  steps: Steps<T>,
  margin: number[],
  getStepSize: (step: T) => { w: number; h: number }
) => {
  const roots = steps.filter((step) => step.parents.length === 0);
  const ctx: StepsCtx<T> = {
    size: { w: 0, h: 0 },
    offset: { x: 0, y: 0 },
    levelHeight: [],
    boxs: steps.map((step) => {
      const box = { step } as any;
      step.box = box;
      step.box.lines = [];
      return box;
    }),
    roots: roots,
  };
  setStepsLevel(steps, ctx);
  setStepsBound(steps, ctx, getStepSize);
  setStepsTreeSize(steps, ctx);
  setStepsOffset(steps, ctx);
  setStepsLines(steps, ctx, margin);

  ctx.size.w = roots.reduce((t, root) => t + root.box.treeSize.w, 0);
  ctx.size.h = ctx.levelHeight.reduce((t, h) => t + h, 0);
  return ctx;
};