fix traversers

fern-api · Oct 10, 2024 · fd66e56 · fd66e56
1 parent 4c1c299
commit fd66e56
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 24 deletions.
diff --git a/packages/fdr-sdk/src/navigation/utils/pruneNavigationTree.ts b/packages/fdr-sdk/src/navigation/utils/pruneNavigationTree.ts
@@ -1,10 +1,14 @@
 import structuredClone from "@ungap/structured-clone";
 import { FernNavigation } from "../..";
-import { bfs } from "../../utils/traversers/bfs";
 import { prunetree } from "../../utils/traversers/prunetree";
 import { mutableDeleteChild } from "./deleteChild";
 import { mutableUpdatePointsTo } from "./updatePointsTo";
 
+type Predicate<T extends FernNavigation.NavigationNode = FernNavigation.NavigationNode> = (
+    node: T,
+    parents: readonly FernNavigation.NavigationNodeParent[],
+) => boolean;
+
 export class Pruner<ROOT extends FernNavigation.NavigationNode> {
     public static from<ROOT extends FernNavigation.NavigationNode>(tree: ROOT): Pruner<ROOT> {
         return new Pruner(tree);
@@ -15,7 +19,7 @@ export class Pruner<ROOT extends FernNavigation.NavigationNode> {
         this.tree = structuredClone(tree) as ROOT;
     }
 
-    public keep(predicate: (node: FernNavigation.NavigationNode) => boolean): this {
+    public keep(predicate: Predicate): this {
         if (this.tree == null) {
             return this;
         }
@@ -29,35 +33,31 @@ export class Pruner<ROOT extends FernNavigation.NavigationNode> {
         return this;
     }
 
-    public hide(predicate: (node: FernNavigation.NavigationNodeWithMetadata) => boolean): this {
+    public remove(predicate: Predicate): this {
+        return this.keep((node, parents) => !predicate(node, parents));
+    }
+
+    public hide(predicate: Predicate<FernNavigation.NavigationNodeWithMetadata>): this {
         if (this.tree == null) {
             return this;
         }
-        bfs(
-            this.tree,
-            (node) => {
-                if (FernNavigation.hasMarkdown(node) && predicate(node)) {
-                    node.hidden = true;
-                }
-            },
-            FernNavigation.getChildren,
-        );
+        FernNavigation.traverseBF(this.tree, (node, parents) => {
+            if (FernNavigation.hasMetadata(node)) {
+                node.hidden = predicate(node, parents) ? true : undefined;
+            }
+        });
         return this;
     }
 
-    public authed(predicate: (node: FernNavigation.NavigationNodeWithMetadata) => boolean): this {
+    public authed(predicate: Predicate<FernNavigation.NavigationNodeWithMetadata>): this {
         if (this.tree == null) {
             return this;
         }
-        bfs(
-            this.tree,
-            (node) => {
-                if (FernNavigation.hasMarkdown(node) && predicate(node)) {
-                    node.authed = true;
-                }
-            },
-            FernNavigation.getChildren,
-        );
+        FernNavigation.traverseBF(this.tree, (node, parents) => {
+            if (FernNavigation.hasMetadata(node)) {
+                node.authed = predicate(node, parents) ? true : undefined;
+            }
+        });
         return this;
     }
 

diff --git a/packages/fdr-sdk/src/utils/traversers/prunetree.ts b/packages/fdr-sdk/src/utils/traversers/prunetree.ts
@@ -6,7 +6,7 @@ interface PruneTreeOptions<NODE, PARENT extends NODE = NODE, POINTER = NODE> {
      * @param node the node to check
      * @returns **false** if the node SHOULD be deleted
      */
-    predicate: (node: NODE) => boolean;
+    predicate: (node: NODE, parents: readonly PARENT[]) => boolean;
     getChildren: (node: PARENT) => readonly NODE[];
 
     /**
@@ -53,7 +53,7 @@ export function prunetree<NODE, ROOT extends NODE = NODE, PARENT extends NODE =
         }
 
         // continue traversal if the node is not to be deleted
-        if (predicate(node)) {
+        if (predicate(node, parents)) {
             return;
         }
         const ancestors = [...parents];