SCCommon/tree/src/tree.ts

import { NodeContext } from './nodeCtx';

export interface TreeRaw {
  ID: string;
  parentID: string;
  seq: number;
  [props: string]: any;
}
interface Node<T extends TreeRaw = TreeRaw> extends TreeRaw {
  getCtx: () => NodeContext<T>; // ctx对象改为getCtx()，因为handsontable数据循环引用会报错
}

export type TreeNode<T extends TreeRaw = TreeRaw> = Node<T> & T;

export interface TreeIDMap<T extends TreeRaw = TreeRaw> {
  [propName: string]: TreeNode<T>;
}

export interface CtxIDMap<T extends TreeRaw = TreeRaw> {
  [propName: string]: NodeContext<T>;
}

export interface TreeParentMap<T extends TreeRaw = TreeRaw> {
  [propName: string]: TreeNode<T>[];
}

export interface ParentMap<T extends TreeRaw = TreeRaw> {
  [propName: string]: (TreeRaw | TreeNode<T>)[];
}

export interface UpdateItem {
  parentID?: string;
  seq?: number;
  [propName: string]: any;
}

export interface UpdateData {
  ID: string;
  update: UpdateItem;
}

export type None = null | undefined;

export class Tree<T extends TreeRaw = TreeRaw> {
  // 原始数据，不经过排序的数据
  rawData: TreeNode<T>[];

  // 按照树结构拼装好、排好序的数据。实际只是原始数据进行排序，内部元素跟原始数据内部元素的引用是一致的
  data: TreeNode<T>[];

  rootID: string;

  // 默认初始顺序号。在对树结构进行操作后，没必要保证seq连号，只需保证正确排序就行，以减少需要更新的节点。
  readonly seqStartIndex = 0;

  // ID与原始数据条目映射
  IDMap: TreeIDMap<T>;

  // parentID与多条同parentID的原始数据条目映射，parentMap内部数组要始终保持正确排序
  parentMap: TreeParentMap<T>;

  // 节点上下文与节点ID映射
  ctxMap: CtxIDMap<T>;

  constructor(rawData: TreeRaw[], rootID = '-1', expanded?: boolean) {
    this.rootID = rootID;
    this.rawData = this.genNodeContext(rawData);
    this.rawData = this.sort(this.rawData);
    this.data = [];
    this.IDMap = {};
    this.parentMap = {};
    this.ctxMap = {};
    this.genMap(this.rawData, expanded);
    this.genData();
  }

  // 根据seq进行排序
  sort(nodes: TreeNode<T>[]): TreeNode<T>[] {
    return nodes.sort((a, b) => a.seq - b.seq);
  }

  // 获取节点的所有parentID
  getParentIDList(nodes: TreeNode<T>[]): Set<string> {
    const parentIDList: Set<string> = new Set();
    nodes.forEach(node => parentIDList.add(node.parentID));
    return parentIDList;
  }

  // 生成节点getCtx方法（由于handsontable循环引用bug，将ctx对象改为getCtx方法），用于获取节点上下文，挂载相关节点方法，TreeRaw转换为TreeNode
  private genNodeContext(rawData: TreeRaw[]): TreeNode<T>[] {
    rawData.forEach(item => {
      item.getCtx = () => this.ctxMap[item.ID];
    });
    return rawData as TreeNode<T>[];
  }

  // 生成映射表
  private genMap(data: TreeNode<T>[], expanded?: boolean): void {
    data.forEach(item => {
      this.IDMap[item.ID] = item;
      this.ctxMap[item.ID] = new NodeContext(item, this, expanded);
      (
        this.parentMap[item.parentID] || (this.parentMap[item.parentID] = [])
      ).push(item);
    });
  }

  // 获取顶层原始数据
  getRoots(): TreeNode<T>[] {
    return this.parentMap[this.rootID] || [];
  }

  // 生成按照树结构排好序的数据
  private genData(): void {
    // genData时不需要排序，因为rawData已经排好序
    const roots = this.getRoots();
    const pushNodesToData = (nodes: TreeNode<T>[]): void => {
      nodes.forEach(node => {
        this.data.push(node);
        const children = this.parentMap[node.ID];
        if (children && children.length) {
          pushNodesToData(children);
        }
      });
    };
    pushNodesToData(roots);
  }

  // 重新生成排序序好的数据，不改变原引用
  private reGenData(): void {
    this.data.splice(0, this.data.length);
    this.genData();
  }

  // 将相关parentMap内的数据、data进行重新排序。新增、删除等操作进行完后，数据需要重新排序
  reSortData(nodes: TreeNode<T>[]): void {
    const toSortList = this.getParentIDList(nodes);
    toSortList.forEach(parentID => {
      if (this.parentMap[parentID] && this.parentMap[parentID].length) {
        this.sort(this.parentMap[parentID]);
      }
    });
    this.reGenData();
  }

  // 根据parentID对parentMap进行重新排序，并重新排序生成data
  resortDataByID(parentIDList: string[]): void {
    parentIDList.forEach(parentID => {
      if (this.parentMap[parentID] && this.parentMap[parentID].length) {
        this.sort(this.parentMap[parentID]);
      }
    });
    this.reGenData();
  }

  // 查找ID节点
  find(ID: string): TreeNode<T> | null {
    return this.IDMap[ID] || null;
  }

  // 查找ID节点的父节点
  findParent(ID: string): TreeNode<T> | null {
    const node = this.find(ID);
    if (!node) {
      return null;
    }
    return this.find(node.parentID);
  }

  // 查找ID节点的下一个节点
  findNext(ID: string): TreeNode<T> | null {
    const node = this.find(ID);
    if (!node) {
      return null;
    }
    const nodes = this.parentMap[node.parentID];
    const nodeIndex = nodes.findIndex(item => item.ID === node.ID);
    if (nodeIndex < 0) {
      return null;
    }
    return nodes[nodeIndex + 1] || null;
  }

  // 查找ID节点的上一个节点
  findPrev(ID: string): TreeNode<T> | null {
    const node = this.find(ID);
    if (!node) {
      return null;
    }
    const nodes = this.parentMap[node.parentID];
    const nodeIndex = nodes.findIndex(item => item.ID === node.ID);
    if (nodeIndex < 0) {
      return null;
    }
    return nodes[nodeIndex - 1] || null;
  }

  // 查询ID节点的子节点
  findChildren(ID: string): TreeNode<T>[] {
    return this.parentMap[ID] || [];
  }

  updateValue(updateData: UpdateData[]): void {
    if (updateData.length) {
      updateData.forEach(updateItem => {
        const node = this.find(updateItem.ID);
        if (node) {
          Object.assign(node, updateItem.update);
        }
      });
    }
  }

  // 递归获取节点
  getNodesPosterity(
    treeNodes: TreeNode<T>[],
    includeSelf = true
  ): TreeNode<T>[] {
    const rst: TreeNode<T>[] = [];
    const pushNodes = (nodes: TreeNode<T>[]): void => {
      nodes.forEach(node => {
        if (includeSelf || !treeNodes.includes(node)) {
          rst.push(node);
        }
        const children = this.findChildren(node.ID);
        if (children.length) {
          pushNodes(children);
        }
      });
    };
    pushNodes(treeNodes);
    return rst;
  }

  // 节选中点块是否是可操作的节点块，选中一批节点进行升降、上下移的时候可能需要用到这个判断。
  // 以第一个节点为基准，如果后续的节点深度均大于于于等于第一个节点的深度，且节点间是连续的，则此为可操作节点块
  isOperable(nodes: TreeNode<T>[]): boolean {
    const baseDepth = nodes[0].getCtx().depth();
    for (let i = 0; i < nodes.length; i++) {
      const node = nodes[i];
      const depth = node.getCtx().depth();
      if (depth < baseDepth) {
        return false;
      }
      const nextRowNode = nodes[i + 1];
      if (
        nextRowNode &&
        node.getCtx().row() + 1 !== nextRowNode.getCtx().row()
      ) {
        return false;
      }
    }
    return true;
  }

  // 从节点块中获取相同深度的节点
  sameDepthNodes(nodes: TreeNode<T>[], depth?: number): TreeNode<T>[] {
    if (!depth) {
      depth = nodes[0].getCtx().depth();
    }
    return nodes.filter(node => node.getCtx().depth() === depth);
  }

  // 准备插入节点，插入节点前，计算出需要更新的数据。（可能会直接覆盖插入数据的seq）
  // 可调用完此方法后，将需要更新、插入的数据提交至数据库，成功响应后调用插入节点更新缓存的方法
  prepareInsert(rawData: TreeRaw[]): UpdateData[] {
    const updateData: UpdateData[] = [];
    const insertParentMap: ParentMap<T> = {};
    // 将相同父项的插入数据和已存在数据进行合并
    rawData.forEach(item => {
      if (typeof item.seq === 'undefined') {
        item.seq = this.seqStartIndex;
      }
      (
        insertParentMap[item.parentID] || (insertParentMap[item.parentID] = [])
      ).push(item);
    });
    const combineMap: ParentMap<T> = {};
    Object.entries(insertParentMap).forEach(([parentID, insertItems]) => {
      const items = this.parentMap[parentID];
      combineMap[parentID] = [...insertParentMap[parentID]];
      if (items) {
        combineMap[parentID].push(...items);
      }
      // 重新排序
      const combineItems = combineMap[parentID];
      this.sort(combineItems as any);
      combineItems.forEach((item, index) => {
        // 插入数据重新赋值
        if (insertItems.includes(item)) {
          item.seq = index;
        } else if (item.seq !== index) {
          // 需要更新的原数据
          updateData.push({
            ID: item.ID,
            update: {
              seq: index,
            },
          });
        }
      });
    });
    return updateData;
  }

  // 插入节点数据
  insert(items: TreeRaw[], updateData: UpdateData[] = []): TreeNode<T>[] {
    // 更新需要更新的节点(一般为更新seq)
    this.updateValue(updateData);
    // 建立映射、插入数据
    const nodes = this.genNodeContext(items);
    this.genMap(nodes);
    this.rawData.push(...nodes);
    // 排序
    this.reSortData(nodes);
    return nodes;
  }

  // 准备删除，返回所有需要删除的节点，包括嵌套节点
  prepareDelete(deleteNodes: TreeNode<T>[]): TreeNode<T>[] {
    return this.getNodesPosterity(deleteNodes);
  }

  /**
   * 删除节点
   * @param treeNodes - 要删除的节点，不需要包含嵌套节点
   */
  delete(treeNodes: TreeNode<T>[]): TreeNode<T>[] {
    const allDeletedNodes: TreeNode<T>[] = [];
    // 递归删除节点
    const deleteNodes = (nodes: TreeNode<T>[]): void => {
      // 删除映射、删除数据
      const toDels: { nodes: TreeNode<T>[]; delNode: TreeNode<T> }[] = [];
      nodes.forEach(node => {
        allDeletedNodes.push(node);
        delete this.IDMap[node.ID];
        delete this.ctxMap[node.ID];
        const children = this.parentMap[node.ID];
        delete this.parentMap[node.ID];
        const nodesInParentMap = this.parentMap[node.parentID];
        if (nodesInParentMap && nodesInParentMap.length) {
          const nIndex = nodesInParentMap.findIndex(
            item => item.ID === node.ID
          );
          if (nIndex >= 0) {
            toDels.push({ nodes: nodesInParentMap, delNode: node });
          }
        }
        const index = this.rawData.findIndex(item => item.ID === node.ID);
        if (index >= 0) {
          this.rawData.splice(index, 1);
        }
        if (children && children.length) {
          deleteNodes(children);
        }
      });
      // 删除parentMap的数据
      toDels.forEach(delItem => {
        const delIndex = delItem.nodes.findIndex(
          item => item.ID === delItem.delNode.ID
        );
        if (delIndex >= 0) {
          delItem.nodes.splice(delIndex, 1);
        }
      });
    };
    deleteNodes(treeNodes);
    // 排序
    this.reSortData(allDeletedNodes);
    return allDeletedNodes;
  }

  // IDList 返回所有需要删除的节点，包括嵌套节点
  prepareDeleteByID(IDList: string[]): TreeNode<T>[] {
    const deleteNodes: TreeNode<T>[] = [];
    IDList.forEach(ID => {
      const node = this.find(ID);
      if (node) {
        deleteNodes.push(node);
      }
    });
    return this.prepareDelete(deleteNodes);
  }

  /**
   * 根据ID删除节点
   * @param IDList - 要删除的节点的ID列表（不包含嵌套节点ID）
   */
  deleteByID(IDList: string[]): TreeNode<T>[] {
    const deleteNodes: TreeNode<T>[] = [];
    IDList.forEach(ID => {
      const node = this.find(ID);
      if (node) {
        deleteNodes.push(node);
      }
    });
    this.delete(deleteNodes);
    return deleteNodes;
  }

  // 准备上移节点块（连续的兄弟节点），注意节点的seq可能不连号
  prepareUpMove(nodes: TreeNode<T>[]): UpdateData[] {
    const updateData: UpdateData[] = [];
    const firstNode = nodes[0];
    const firstNodePrev = this.findPrev(firstNode.ID);
    if (!firstNodePrev) {
      return [];
    }
    let tempSeq = firstNodePrev.seq;
    nodes.forEach(node => {
      const orgSeq = node.seq;
      updateData.push({
        ID: node.ID,
        update: { seq: tempSeq },
      });
      tempSeq = orgSeq;
    });
    updateData.push({
      ID: firstNodePrev.ID,
      update: { seq: tempSeq },
    });
    return updateData;
  }

  // 准备下移节点块（连续的兄弟节点），注意节点的seq可能不连号
  prepareDownMove(nodes: TreeNode<T>[]): UpdateData[] {
    const updateData: UpdateData[] = [];
    const lastNode = nodes[nodes.length - 1];
    const lastNodeNext = this.findNext(lastNode.ID);
    if (!lastNodeNext) {
      return [];
    }
    let tempSeq = lastNodeNext.seq;
    for (let i = nodes.length - 1; i >= 0; i--) {
      const node = nodes[i];
      const orgSeq = node.seq;
      updateData.push({
        ID: node.ID,
        update: { seq: tempSeq },
      });
      tempSeq = orgSeq;
    }
    updateData.push({
      ID: lastNodeNext.ID,
      update: { seq: tempSeq },
    });
    return updateData;
  }

  // 上下移
  move(nodes: TreeNode<T>[], updateData: UpdateData[]): void {
    this.updateValue(updateData);
    this.reSortData(nodes);
  }

  // 准备升级节点块（连续的兄弟节点），不维护seq连号
  prepareUpLevel(nodes: TreeNode<T>[]): UpdateData[] {
    const updateData: UpdateData[] = [];
    const firstNode = nodes[0];
    const lastNode = nodes[nodes.length - 1];
    const parent = this.findParent(firstNode.ID);
    if (!parent) {
      return [];
    }
    const baseSeq = parent.seq + 1;
    nodes.forEach((node, index) => {
      updateData.push({
        ID: node.ID,
        update: { parentID: parent.parentID, seq: baseSeq + index },
      });
    });
    const parentNextBrothers = parent.getCtx().nextBrothers();
    // 因为seq可能是不连号的，如果升级块的最末节点seq，小于下一节点的seq，那就不更新所有下兄弟节点的seq，减少更新的数据量
    const lastNodeCurSeq = updateData[updateData.length - 1].update.seq;
    const firstBrother = parentNextBrothers[0];
    if (lastNodeCurSeq && firstBrother && lastNodeCurSeq >= firstBrother.seq) {
      parentNextBrothers.forEach((node, index) => {
        updateData.push({
          ID: node.ID,
          update: { seq: baseSeq + index + nodes.length },
        });
      });
    }
    // 最末节点的所有后兄弟节点，成为最末节点的子节点
    const lastNodeNextBrothers = lastNode.getCtx().nextBrothers();
    const lastNodeLastChild = lastNode.getCtx().lastChild();
    const lastNodeLastChildSeq = lastNodeLastChild ? lastNodeLastChild.seq : 0;
    lastNodeNextBrothers.forEach((node, index) => {
      updateData.push({
        ID: node.ID,
        update: { parentID: lastNode.ID, seq: index + lastNodeLastChildSeq },
      });
    });
    return updateData;
  }

  upLevel(nodes: TreeNode<T>[], updateData: UpdateData[]): void {
    const firstNode = nodes[0];
    const lastNode = nodes[nodes.length - 1];
    if (!firstNode.getCtx().canUpLevel()) {
      return;
    }
    const orgParentID = firstNode.parentID;
    const orgBrothers = this.parentMap[orgParentID];
    const lastNodeNextBrothers = lastNode.getCtx().nextBrothers();
    const index = orgBrothers.findIndex(item => item.ID === firstNode.ID);
    orgBrothers.splice(index, nodes.length + lastNodeNextBrothers.length);
    (this.parentMap[lastNode.ID] || (this.parentMap[lastNode.ID] = [])).push(
      ...lastNodeNextBrothers
    );
    this.updateValue(updateData);
    const newParentID = firstNode.parentID;
    this.parentMap[newParentID].push(...nodes);
    this.resortDataByID([orgParentID, newParentID, lastNode.ID]);
  }

  // 准备降级节点块（连续的兄弟节点），不维护seq连号
  prepareDownLevel(nodes: TreeNode<T>[]): UpdateData[] {
    const updateData: UpdateData[] = [];
    const firstNode = nodes[0];
    const prevNode = this.findPrev(firstNode.ID);
    if (!prevNode) {
      return [];
    }
    // 节点块成为前节点的子节点
    const prevNodeLastChild = prevNode.getCtx().lastChild();
    const baseSeq = prevNodeLastChild
      ? prevNodeLastChild.seq + 1
      : this.seqStartIndex;
    nodes.forEach((node, index) => {
      updateData.push({
        ID: node.ID,
        update: { parentID: prevNode.ID, seq: baseSeq + index },
      });
    });
    return updateData;
  }

  downLevel(nodes: TreeNode<T>[], updateData: UpdateData[]): void {
    const firstNode = nodes[0];
    if (!firstNode.getCtx().canDownLevel()) {
      return;
    }
    const prevNode = this.findPrev(firstNode.ID);
    if (!prevNode) {
      return;
    }
    const orgBrothers = this.parentMap[firstNode.parentID];
    const index = orgBrothers.findIndex(item => item.ID === firstNode.ID);
    orgBrothers.splice(index, nodes.length);
    (this.parentMap[prevNode.ID] || (this.parentMap[prevNode.ID] = [])).push(
      ...nodes
    );
    this.updateValue(updateData);
    this.reGenData();
  }

  // 准备移动节点块（连续的兄弟节点），不维护seq连号
  prepareMoveTo(
    nodes: TreeNode<T>[],
    parent: TreeNode<T> | None,
    next: TreeNode<T> | None
  ): UpdateData[] {
    const updateData: UpdateData[] = [];
    let prev: TreeNode<T> | null;
    if (next) {
      prev = next.getCtx().prev();
    } else {
      const roots = this.getRoots();
      prev = parent
        ? parent.getCtx().lastChild()
        : roots[roots.length - 1] || null;
    }
    const baseSeq = prev ? prev.seq + 1 : this.seqStartIndex;
    updateData.push(
      ...nodes.map((node, index) => ({
        ID: node.ID,
        update: {
          parentID: (parent && parent.ID) || this.rootID,
          seq: baseSeq + index,
        },
      }))
    );
    const curBaseSeq = baseSeq + nodes.length;
    // eslint-disable-next-line no-nested-ternary
    const nextBrothers = prev
      ? prev.getCtx().nextBrothers()
      : next
      ? [next, ...next.getCtx().nextBrothers()]
      : [];
    updateData.push(
      ...nextBrothers.map((node, index) => ({
        ID: node.ID,
        update: {
          seq: curBaseSeq + index,
        },
      }))
    );
    return updateData;
  }

  moveTo(
    nodes: TreeNode<T>[],
    parent: TreeNode<T> | None,
    updateData: UpdateData[]
  ): void {
    const firstNode = nodes[0];
    const newParentID = parent ? parent.ID : this.rootID;
    const orgParentID = firstNode.parentID;
    const orgBrothers = this.parentMap[orgParentID];
    const index = orgBrothers.findIndex(item => item.ID === firstNode.ID);
    orgBrothers.splice(index, nodes.length);
    (this.parentMap[newParentID] || (this.parentMap[newParentID] = [])).push(
      ...nodes
    );
    this.updateValue(updateData);
    this.resortDataByID([orgParentID, newParentID]);
  }

  // 清除所有节点（不变更data引用）
  clear(): void {
    this.rawData.splice(0, this.rawData.length);
    this.data.splice(0, this.data.length);
    Object.keys(this.parentMap).forEach(key => delete this.parentMap[key]);
    Object.keys(this.IDMap).forEach(key => delete this.IDMap[key]);
    Object.keys(this.ctxMap).forEach(key => delete this.ctxMap[key]);
  }
}