feat(TreeData): 新增 removeNode、removeNodes、count

fjc0k · Dec 12, 2020 · 7c009f3 · 7c009f3
1 parent 5ca95e0
commit 7c009f3
Show file tree

Hide file tree

Showing 4 changed files with 1,101 additions and 384 deletions.
diff --git a/jest.config.js b/jest.config.js
@@ -7,5 +7,5 @@ module.exports = require('haoma').getJestConfig({
     '@babel/runtime',
     '@tarojs/runtime',
   ],
-  transformer: 'swc',
+  transformer: 'babel',
 })
diff --git a/src/utils/TreeData.test.ts b/src/utils/TreeData.test.ts
@@ -59,7 +59,7 @@ describe('TreeData', () => {
       expect(ids).toMatchSnapshot()
     })
 
-    test('removeNode', () => {
+    test('payload.removeNode', () => {
       const ids: string[] = []
       new TreeData(data).traverse([
         payload => {
@@ -134,6 +134,27 @@ describe('TreeData', () => {
     ).toMatchSnapshot()
   })
 
+  test('removeNode', () => {
+    const tree = new TreeData(data)
+    const removed = tree.removeNode(_ => _.node.id.startsWith('1.'))
+    expect(removed).toMatchSnapshot()
+    expect(tree.export()).toMatchSnapshot()
+  })
+
+  test('removeNodes', () => {
+    const tree = new TreeData(data)
+    const removed = tree.removeNodes(_ => _.node.id.startsWith('1.'))
+    expect(removed).toMatchSnapshot()
+    expect(tree.export()).toMatchSnapshot()
+  })
+
+  test('count', () => {
+    expect(new TreeData(data).count()).toMatchSnapshot()
+    expect(
+      new TreeData(data).count(_ => _.node.id.startsWith('1.')),
+    ).toMatchSnapshot()
+  })
+
   test('clone', () => {
     expect(new TreeData(data).clone().export()).toMatchSnapshot()
   })

diff --git a/src/utils/TreeData.ts b/src/utils/TreeData.ts
@@ -108,23 +108,27 @@ export class TreeData<TNode extends TreeDataNode> {
     path: TNode[],
   ) {
     if (searchMethod === 'DFS') {
+      const postActions: Array<() => void> = []
       for (const node of data) {
-        const postActions: Array<() => void> = []
+        let isRemove = false
         let isExit = false
         fn({
           node: node,
           parentNode: parentNode,
           depth: depth,
           path: path,
-          removeNode: () =>
-            postActions.push(() => data.splice(data.indexOf(node), 1)),
+          removeNode: () => {
+            isRemove = true
+            postActions.push(() => data.splice(data.indexOf(node), 1))
+          },
           exit: () => (isExit = true),
         })
-        for (const action of postActions) {
-          action()
-        }
         if (isExit) return
-        if (node[childrenPropName] && Array.isArray(node[childrenPropName])) {
+        if (
+          !isRemove &&
+          node[childrenPropName] &&
+          Array.isArray(node[childrenPropName])
+        ) {
           TreeData.traverse(
             node[childrenPropName],
             childrenPropName,
@@ -136,17 +140,23 @@ export class TreeData<TNode extends TreeDataNode> {
           )
         }
       }
+      for (const action of postActions) {
+        action()
+      }
     } else {
+      const removeNodes: TNode[] = []
+      const postActions: Array<() => void> = []
       for (const node of data) {
-        const postActions: Array<() => void> = []
         let isExit = false
         fn({
           node: node,
           parentNode: parentNode,
           depth: depth,
           path: path,
-          removeNode: () =>
-            postActions.push(() => data.splice(data.indexOf(node), 1)),
+          removeNode: () => {
+            removeNodes.push(node)
+            postActions.push(() => data.splice(data.indexOf(node), 1))
+          },
           exit: () => (isExit = true),
         })
         for (const action of postActions) {
@@ -155,7 +165,11 @@ export class TreeData<TNode extends TreeDataNode> {
         if (isExit) return
       }
       for (const node of data) {
-        if (node[childrenPropName] && Array.isArray(node[childrenPropName])) {
+        if (
+          removeNodes.indexOf(node) !== -1 &&
+          node[childrenPropName] &&
+          Array.isArray(node[childrenPropName])
+        ) {
           TreeData.traverse(
             node[childrenPropName],
             childrenPropName,
@@ -167,6 +181,9 @@ export class TreeData<TNode extends TreeDataNode> {
           )
         }
       }
+      for (const action of postActions) {
+        action()
+      }
     }
   }
 
@@ -354,6 +371,60 @@ export class TreeData<TNode extends TreeDataNode> {
     return paths
   }
 
+  /**
+   * 移除符合条件的第一个节点。返回被移除的节点。
+   *
+   * @param predicate 条件
+   */
+  removeNode(
+    predicate: (payload: TreeDataTraverseFnPayload<TNode>) => boolean,
+  ): TNode | undefined {
+    let node: TNode | undefined
+    this.traverse(payload => {
+      if (predicate(payload)) {
+        payload.removeNode()
+        node = payload.node
+        payload.exit()
+      }
+    })
+    return node
+  }
+
+  /**
+   * 移除符合条件的所有节点。返回被移除的节点组成的数组。
+   *
+   * @param predicate 条件
+   */
+  removeNodes(
+    predicate: (payload: TreeDataTraverseFnPayload<TNode>) => boolean,
+  ): TNode[] {
+    const nodes: TNode[] = []
+    this.traverse(payload => {
+      if (predicate(payload)) {
+        payload.removeNode()
+        nodes.push(payload.node)
+      }
+    })
+    return nodes
+  }
+
+  /**
+   * 计算符合条件的节点个数。不给出条件则计算所有节点的个数。
+   *
+   * @param predicate 条件
+   */
+  count(
+    predicate?: (payload: TreeDataTraverseFnPayload<TNode>) => boolean,
+  ): number {
+    let counter = 0
+    this.traverse(payload => {
+      if (predicate ? predicate(payload) : true) {
+        counter++
+      }
+    })
+    return counter
+  }
+
   /**
    * 克隆实例。
    */