From 103664ba3ba84236ab1d8bd6e71e0ec16c1e0ea5 Mon Sep 17 00:00:00 2001
From: w2xi <43wangxi@gmail.com>
Date: Tue, 12 Mar 2024 23:25:22 +0800
Subject: [PATCH] fix: ignore some when using bfs (fix #2)

---
 src/toTree.ts       | 46 ++++++++++++++++++++++++---------------------
 src/utils.ts        | 13 +++++++++++++
 test/toTree.spec.ts | 21 +++++++++++----------
 3 files changed, 49 insertions(+), 31 deletions(-)

diff --git a/src/toTree.ts b/src/toTree.ts
index b563953..5b1e013 100644
--- a/src/toTree.ts
+++ b/src/toTree.ts
@@ -45,7 +45,6 @@ function dfs(path: string, options: Options): TreeNode {
 }
 
 function bfs(path: string, options: Options) {
-  let deep = 0
   const { ignore, onlyFolder, layer } = options
   const root = {
     path,
@@ -55,32 +54,37 @@ function bfs(path: string, options: Options) {
   } as TreeNode
   const queue = [root]
 
-  while (queue.length) {
-    const node = queue.shift()
-    const { path, children } = node!
-    const dir = fs.readdirSync(path!)
+  let deep = 0
 
+  while (queue.length) {
     if (layer && deep >= layer) break
-
     deep++
 
-    for (let i = 0; i < dir.length; i++) {
-      const item = dir[i]
-      if (ignore && ignore.includes(item)) continue
-      const childPath = `${path}/${item}`
-      const isDir = isDirectory(childPath)
-      if (onlyFolder && !isDir) continue
+    let size = queue.length
+
+    while (size--) {
+      const node = queue.shift()
+      const { path, children } = node!
+      const dir = fs.readdirSync(path!)
+
+      for (let i = 0; i < dir.length; i++) {
+        const item = dir[i]
+        if (ignore && ignore.includes(item)) continue
+        const childPath = `${path}/${item}`
+        const isDir = isDirectory(childPath)
+        if (onlyFolder && !isDir) continue
 
-      const childItem = {
-        path: childPath,
-        name: item,
-        type: isDir ? NodeTypes.DIRECTORY : NodeTypes.FILE
-      } as TreeNode
-      if (isDir) {
-        queue.push(childItem)
-        childItem.children = []
+        const childItem = {
+          path: childPath,
+          name: item,
+          type: isDir ? NodeTypes.DIRECTORY : NodeTypes.FILE
+        } as TreeNode
+        if (isDir) {
+          queue.push(childItem)
+          childItem.children = []
+        }
+        children && children.push(childItem)
       }
-      children && children.push(childItem)
     }
   }
 
diff --git a/src/utils.ts b/src/utils.ts
index 6bfcf1c..bc2eb7d 100644
--- a/src/utils.ts
+++ b/src/utils.ts
@@ -1,4 +1,5 @@
 import fs from 'fs'
+import { TreeNode } from './type'
 
 // check if a file or directory exists.
 export function fileExistSync(path: string) {
@@ -14,3 +15,15 @@ export function isDirectory(path: string) {
   const stats = fs.lstatSync(path)
   return stats.isDirectory()
 }
+
+export function getMaxLayer(data: TreeNode) {
+  if (!data.children) return 0
+  let max = 0
+  for (let i = 0; i < data.children.length; i++) {
+    const child = data.children[i]
+    if (child.children && child.children.length) {
+      max = Math.max(max, getMaxLayer(child))
+    }
+  }
+  return max + 1
+}
diff --git a/test/toTree.spec.ts b/test/toTree.spec.ts
index 5e4f347..cd7a424 100644
--- a/test/toTree.spec.ts
+++ b/test/toTree.spec.ts
@@ -2,6 +2,7 @@ import { describe, test, expect } from 'vitest'
 import { toTree } from '../src/toTree'
 import { Options } from '../src/type'
 import { NodeTypes } from '../src/config'
+import { getMaxLayer } from '../src/utils'
 
 describe('toTree', () => {
   const options = {
@@ -24,16 +25,16 @@ describe('toTree', () => {
 
   test('1 layer', () => {
     const result = toTree({ ...options, layer: 1 })
-    const layer = 1
-    let wantedLayer = 1
-    for (let i = 0; i < result.children!.length; i++) {
-      if (result.children![i].children) {
-        wantedLayer = 2
-      } else {
-        break
-      }
-    }
-    expect(wantedLayer).toBe(layer)
+    const wantedLayer = 1
+    const layer = getMaxLayer(result)
+    expect(layer).toBe(wantedLayer)
+  })
+
+  test('2 layer', () => {
+    const result = toTree({ ...options, layer: 2 })
+    const wantedLayer = 2
+    const layer = getMaxLayer(result)
+    expect(layer).toBe(wantedLayer)
   })
 
   test('only folder', () => {