处理树形结构遍历的工具--Tree

/**
 * 处理树结构遍历工具
 * 所有遍历都是基于递归且深度优先
 */
export default class Tree<T> {
  private _value: T[] = [];

  private childrenKey: string = 'children';

  constructor(tree: T[], { childrenKey = 'children' }: { childrenKey?: string } = {}) {
    this._value = tree;
    this.childrenKey = childrenKey;
  }

  forEach(callback: (item: T, parent?: T) => void) {
    const stack: { node: T; parent?: T }[] = this._value.map(node => ({ node }));
    while (stack.length) {
      const { node, parent } = stack.shift()!;
      callback(node, parent);
      const children: T[] | undefined = (node as any)[this.childrenKey];
      if (children?.length)
        stack.unshift(...children.map(child => ({ node: child, parent: node })));
    }
    return this;
  }

  map<P>(callback: (item: T, parent?: P) => P, newChildrenKey: string = this.childrenKey) {
    const stack: { node: T; parent?: P }[] = this._value.map(node => ({ node }));
    const newTree: P[] = [];
    while (stack.length) {
      const { node, parent } = stack.shift()!;
      const newNode = callback(node, parent);
      if (!parent) newTree.push(newNode);
      else {
        (parent as any)[newChildrenKey] = (parent as any)[newChildrenKey] || [];
        (parent as any)[newChildrenKey].push(newNode);
      }
      const children: T[] | undefined = (node as any)[this.childrenKey];
      if (children?.length) {
        (newNode as any)[newChildrenKey] = [];
        stack.unshift(...children.map(child => ({ node: child, parent: newNode })));
      }
    }
    return new Tree(newTree, { childrenKey: newChildrenKey });
  }

  filter(callback: (item: T, parent?: T) => boolean) {
    const stack: { node: T; parent?: T }[] = this._value.map(node => ({ node }));
    const newTree: T[] = [];
    while (stack.length) {
      const { node, parent } = stack.shift()!;
      if (callback(node, parent)) {
        if (!parent) newTree.push(node);
        else {
          (parent as any)[this.childrenKey] = (parent as any)[this.childrenKey] || [];
          (parent as any)[this.childrenKey].push(node);
        }
      }
      const children: T[] | undefined = (node as any)[this.childrenKey];
      if (children?.length) {
        (node as any)[this.childrenKey] = [];
        stack.unshift(...children.map(child => ({ node: child, parent: node })));
      }
    }
    return new Tree(newTree, { childrenKey: this.childrenKey });
  }

  find(callback: (item: T, parent?: T) => boolean) {
    const stack: { node: T; parent?: T }[] = this._value.map(node => ({ node }));
    while (stack.length) {
      const { node, parent } = stack.shift()!;
      if (callback(node, parent)) return node;
      const children: T[] | undefined = (node as any)[this.childrenKey];
      if (children?.length)
        stack.unshift(...children.map(child => ({ node: child, parent: node })));
    }
  }

  findAll(callback: (item: T, parent?: T) => boolean) {
    const target: T[] = [];
    this.forEach(item => {
      if (callback(item)) target.push(item);
    });
    return new Tree(target, { childrenKey: this.childrenKey });
  }

  get value() {
    return this._value;
  }
}