add Tree.findDepthFirst and Tree.findBreadthFirst and Tree.filter

src/Tree.ts

@@ -32,6 +32,7 @@ import { Predicate } from './Predicate'
 import * as RA from './ReadonlyArray'
 import type { Show } from './Show'
 import type { Traversable1 } from './Traversable'
+import * as O from './Option'
 
 // -------------------------------------------------------------------------------------
 // model
@@ -531,6 +532,58 @@ export const exists = <A>(predicate: Predicate<A>) => (ma: Tree<A>): boolean =>
  */
 export const elem = <A>(E: Eq<A>) => (a: A): ((fa: Tree<A>) => boolean) => exists(E.equals(a))
 
+/**
+ * Deep first search
+ * The predicate judgment acts on the tree itself, not the value, so that the predicate can access the forest
+ * @since:3.0.0
+ */
+export const findDepthFirst = <A>(predicate: (a: Tree<A>) => boolean) => (tree: Tree<A>): O.Option<Tree<A>> => {
+  if (predicate(tree)) {
+    return O.some(tree)
+  }
+
+  const todo: Array<Tree<A>> = [...tree.forest]
+
+  while (todo.length > 0) {
+    const current = todo.shift()!
+    if (predicate(current)) {
+      return O.some(current)
+    }
+    todo.unshift(...current.forest)
+  }
+
+  return O.none
+}
+
+/**
+ * Breadth first search
+ * The predicate judgment acts on the tree itself, not the value, so that the predicate can access the forest
+ * @since:3.0.0
+ */
+export const findBreadthFirst = <A>(predicate: (a: Tree<A>) => boolean) => (tree: Tree<A>): O.Option<Tree<A>> => {
+  if (predicate(tree)) {
+    return O.some(tree)
+  }
+
+  const todo: Array<Tree<A>> = [...tree.forest]
+
+  while (todo.length > 0) {
+    const result = pipe(
+      todo,
+      RA.findFirst((tree) => predicate(tree))
+    )
+
+    if (O.isSome(result)) {
+      return result
+    }
+
+    const todoCopy: ReadonlyArray<Tree<A>> = [...todo].reverse()
+    todoCopy.forEach((tree) => todo.unshift(...tree.forest))
+  }
+
+  return O.none
+}
+
 const draw = (indentation: string, forest: Forest<string>): string => {
   let r = ''
   const len = forest.length
@@ -575,6 +628,27 @@ export const drawForest = (forest: Forest<string>): string => draw('\n', forest)
  */
 export const drawTree = (tree: Tree<string>): string => tree.value + drawForest(tree.forest)
 
+/**
+ * Filter the tree, if the root does not pass the predicate, it returns none
+ * @since:3.0.0
+ */
+export const filter = <A>(predicate: (a: Tree<A>) => boolean) => (t: Tree<A>): O.Option<Tree<A>> => {
+  if (predicate(t)) {
+    const forest = pipe(t.forest, RA.map(filter(predicate)), RA.compact)
+    return O.some(tree(t.value, forest))
+  }
+  return O.none
+}
+
+/**
+ * Like the `map`, but give the chance to change the forest
+ * @since:3.0.0
+ */
+export const modify = <A>(f: (ta: Tree<A>) => Tree<A>) => (ta: Tree<A>): Tree<A> => {
+  const tb = f(ta)
+  return tree(tb.value, tb.forest.map(modify(f)))
+}
+
 // -------------------------------------------------------------------------------------
 // do notation
 // -------------------------------------------------------------------------------------

test/Tree.ts

@@ -1,4 +1,5 @@
 import * as Eq from '../src/Eq'
+import * as Ra from '../src/ReadonlyArray'
 import { flow, identity, pipe } from '../src/function'
 import * as S from '../src/string'
 import * as O from '../src/Option'
@@ -258,4 +259,80 @@ describe('Tree', () => {
     U.deepStrictEqual(_.exists((user: User) => user.id === 4)(users), true)
     U.deepStrictEqual(_.exists((user: User) => user.id === 5)(users), false)
   })
+
+  it('findDepthFirst', () => {
+    interface User {
+      readonly id: number
+      readonly name?: string
+    }
+    const users: _.Tree<User> = _.tree({ id: 1 }, [
+      _.tree({ id: 1 }, [
+        _.tree({ id: 3, name: '' }, [_.tree({ id: 2, name: 'deep first [id]=2 node' })]),
+        _.tree({ id: 4 })
+      ]),
+      _.tree({ id: 2 })
+    ])
+
+    U.deepStrictEqual(
+      _.findDepthFirst((tree: _.Tree<User>) => tree.value.id === 2)(users),
+      O.some(_.tree({ id: 2, name: 'deep first [id]=2 node' }))
+    )
+    U.deepStrictEqual(_.findDepthFirst((tree: _.Tree<User>) => tree.value.id === 5)(users), O.none)
+  })
+
+  it('findBreadthFirst', () => {
+    interface User {
+      readonly id: number
+      readonly name?: string
+    }
+    const users: _.Tree<User> = _.tree({ id: 1 }, [
+      _.tree({ id: 1 }, [_.tree({ id: 3 }, [_.tree({ id: 2 })]), _.tree({ id: 4 })]),
+      _.tree({ id: 2, name: 'breadth first [id]=2 node' })
+    ])
+
+    U.deepStrictEqual(
+      _.findBreadthFirst((tree: _.Tree<User>) => tree.value.id === 2)(users),
+      O.some(_.tree({ id: 2, name: 'breadth first [id]=2 node' }))
+    )
+    U.deepStrictEqual(_.findDepthFirst((tree: _.Tree<User>) => tree.value.id === 5)(users), O.none)
+  })
+
+  it('filter', () => {
+    interface User {
+      readonly id: number
+      readonly name?: string
+    }
+    const users: _.Tree<User> = _.tree({ id: 1 }, [
+      _.tree({ id: 1 }, [_.tree({ id: 3 }, [_.tree({ id: 2 })]), _.tree({ id: 4 })]),
+      _.tree({ id: 2, name: 'breadth first [id]=2 node' })
+    ])
+
+    U.deepStrictEqual(_.filter((tree: _.Tree<User>) => tree.value.id < 0)(users), O.none)
+    U.deepStrictEqual(
+      _.filter((tree: _.Tree<User>) => tree.value.id !== 2)(users),
+      O.some(_.tree({ id: 1 }, [_.tree({ id: 1 }, [_.tree({ id: 3 }), _.tree({ id: 4 })])]))
+    )
+  })
+
+  it('modify', () => {
+    interface User {
+      readonly id: number
+      readonly name?: string
+    }
+    const users: _.Tree<User> = _.tree({ id: 1 }, [
+      _.tree({ id: 1 }, [_.tree({ id: 3 }, [_.tree({ id: 2 })]), _.tree({ id: 4 })]),
+      _.tree({ id: 2, name: 'breadth first [id]=2 node' })
+    ])
+
+    U.deepStrictEqual(
+      _.modify((tree: _.Tree<User>) => ({
+        value: tree.value,
+        forest: Ra.reverse(tree.forest)
+      }))(users),
+      _.tree({ id: 1 }, [
+        _.tree({ id: 2, name: 'breadth first [id]=2 node' }),
+        _.tree({ id: 1 }, [_.tree({ id: 4 }), _.tree({ id: 3 }, [_.tree({ id: 2 })])])
+      ])
+    )
+  })
 })