From d695c5bd415e5a8bfb1e427fae3979db8a779cb0 Mon Sep 17 00:00:00 2001
From: Cody Bennett <hi@codyb.co>
Date: Mon, 5 Sep 2022 04:23:37 -0500
Subject: [PATCH] fix(traverseFiber): start traverse at head, accept any Fiber
 type (#7)

---
 README.md            | 12 ++++------
 src/index.tsx        | 54 +++++++++++++-------------------------------
 tests/index.test.tsx | 13 +++++++----
 3 files changed, 29 insertions(+), 50 deletions(-)

diff --git a/README.md b/README.md
index df22f88..a7487d6 100644
--- a/README.md
+++ b/README.md
@@ -169,19 +169,17 @@ Additional exported utility functions for raw handling of Fibers.
 
 ### traverseFiber
 
-Traverses up or down through a `Fiber`, return `true` to stop and select a node.
+Traverses up or down a `Fiber`, return `true` to stop and select a node.
 
 ```ts
 import { type Fiber, traverseFiber } from 'its-fine'
 
-// Whether to ascend and walk up the tree. Will walk down if `false`
-const ascending: boolean = true
-
 // Traverses through the Fiber tree, returns the current node when `true` is passed via selector
 const parentDiv: Fiber<HTMLDivElement> | undefined = traverseFiber<HTMLDivElement>(
-  // A composite component Fiber from `useFiber` or a Fiber handle from a reconciler
-  fiber as Fiber<null>,
-  ascending,
+  // Input Fiber to traverse
+  fiber as Fiber,
+  // Whether to ascend and walk up the tree. Will walk down if `false`
+  true,
   // A Fiber node selector, returns the first <div /> element in JSX
   (node: Fiber<HTMLDivElement | null>) => node.type === 'div',
 )
diff --git a/src/index.tsx b/src/index.tsx
index a0ff0de..b0b65b0 100644
--- a/src/index.tsx
+++ b/src/index.tsx
@@ -12,39 +12,22 @@ export type Fiber<T = any> = Omit<ReactReconciler.Fiber, 'stateNode'> & { stateN
 export type FiberSelector<T = any> = (node: Fiber<T | null>) => boolean | void
 
 /**
- * Traverses up or down through a {@link Fiber}, return `true` to stop and select a node.
+ * Traverses up or down a {@link Fiber}, return `true` to stop and select a node.
  */
 export function traverseFiber<T = any>(
-  fiber: Fiber<null>,
+  fiber: Fiber,
   ascending: boolean,
   selector: FiberSelector<T>,
 ): Fiber<T> | undefined {
-  let halted = false
-  let selected: Fiber<T> | undefined
-
-  let node = ascending ? fiber.return : fiber.child
-  let sibling = fiber.sibling
-  while (node) {
-    while (sibling) {
-      halted ||= selector(sibling) === true
-      if (halted) {
-        selected = sibling
-        break
-      }
-
-      sibling = sibling.sibling
-    }
-
-    halted ||= selector(node) === true
-    if (halted) {
-      selected = node
-      break
-    }
-
-    node = ascending ? node.return : node.child
-  }
+  if (selector(fiber) === true) return fiber
+
+  let child = ascending ? fiber.return : fiber.child
+  while (child) {
+    const match = traverseFiber(child, ascending, selector)
+    if (match) return match
 
-  return selected
+    child = child.sibling
+  }
 }
 
 interface ReactInternal {
@@ -77,17 +60,12 @@ export interface ContainerInstance<T = any> {
  */
 export function useContainer<T = any>(): T {
   const fiber = useFiber()
-  const container = React.useMemo(
-    () =>
-      traverseFiber<ContainerInstance<T>>(
-        fiber,
-        true,
-        (node) => node.type == null && node.stateNode?.containerInfo != null,
-      )!.stateNode.containerInfo,
+  const root = React.useMemo(
+    () => traverseFiber<ContainerInstance<T>>(fiber, true, (node) => node.stateNode?.containerInfo != null)!,
     [fiber],
   )
 
-  return container
+  return root!.stateNode.containerInfo
 }
 
 /**
@@ -135,14 +113,14 @@ export type ContextBridge = React.FC<React.PropsWithChildren<{}>>
 export function useContextBridge(): ContextBridge {
   const fiber = useFiber()
   const contexts = React.useMemo(() => {
-    const unique = new Set<React.Context<any>>()
+    const unique: React.Context<any>[] = []
 
     traverseFiber(fiber, true, (node) => {
       const context = node.type?._context
-      if (context && !unique.has(context)) unique.add(context)
+      if (context && !unique.includes(context)) unique.push(context)
     })
 
-    return Array.from(unique)
+    return unique
   }, [fiber])
 
   return contexts.reduce(
diff --git a/tests/index.test.tsx b/tests/index.test.tsx
index bc0b6a3..b2ce57c 100644
--- a/tests/index.test.tsx
+++ b/tests/index.test.tsx
@@ -94,12 +94,13 @@ describe('traverseFiber', () => {
       )
     })
 
-    const traversed = [] as unknown as [child: Fiber<Primitive>]
+    const traversed = [] as unknown as [self: Fiber<null>, child: Fiber<Primitive>]
     traverseFiber(fiber, false, (node) => void traversed.push(node))
 
-    expect(traversed.length).toBe(1)
+    expect(traversed.length).toBe(2)
 
-    const [child] = traversed
+    const [self, child] = traversed
+    expect(self.type).toBe(Test)
     expect(child.stateNode.props.name).toBe('child')
   })
 
@@ -120,14 +121,16 @@ describe('traverseFiber', () => {
     })
 
     const traversed = [] as unknown as [
+      self: Fiber<null>,
       parent: Fiber<Primitive>,
       rootContainer: Fiber<ContainerInstance<HostContainer>>,
     ]
     traverseFiber(fiber, true, (node) => void traversed.push(node))
 
-    expect(traversed.length).toBe(2)
+    expect(traversed.length).toBe(3)
 
-    const [parent, rootContainer] = traversed
+    const [self, parent, rootContainer] = traversed
+    expect(self.type).toBe(Test)
     expect(parent.stateNode.props.name).toBe('parent')
     expect(rootContainer.stateNode.containerInfo).toBe(container)
   })