refactor!: only halt traverseFiber on boolean, harden Fiber types (#5)

README.md

@@ -60,9 +60,9 @@ function Component() {
 
 ### useContainer
 
-Returns the current react-reconciler `Container` or the Fiber created from `Reconciler.createContainer`.
+Returns the current react-reconciler container info passed to `Reconciler.createContainer`.
 
-In react-dom, `container.containerInfo` will point to the root DOM element; in react-three-fiber, it will point to the root Zustand store.
+In react-dom, a container will point to the root DOM element; in react-three-fiber, it will point to the root Zustand store.
 
 ```tsx
 import * as React from 'react'
@@ -73,7 +73,7 @@ function Component() {
 
   React.useLayoutEffect(() => {
     // <div> (e.g. react-dom)
-    console.log(container.containerInfo)
+    console.log(container)
   }, [container])
 }
 ```
@@ -166,5 +166,10 @@ Traverses up or down through a `Fiber`, return `true` to stop and select a node.
 import { type Fiber, traverseFiber } from 'its-fine'
 
 const ascending = true
-const prevElement: Fiber = traverseFiber(fiber, ascending, (node: Fiber) => node.type === 'element')
+
+const parentDiv: Fiber<HTMLDivElement> | undefined = traverseFiber<HTMLDivElement>(
+  fiber as Fiber<null>,
+  ascending,
+  (node: Fiber<HTMLDivElement | null>) => node.type === 'div',
+)
 ```

src/index.tsx

@@ -4,25 +4,29 @@ import type ReactReconciler from 'react-reconciler'
 /**
  * Represents a react-internal Fiber node.
  */
-export type Fiber = ReactReconciler.Fiber
+export type Fiber<T = any> = Omit<ReactReconciler.Fiber, 'stateNode'> & { stateNode: T }
 
 /**
  * Represents a {@link Fiber} node selector for traversal.
  */
-export type FiberSelector = (node: Fiber) => boolean | void
+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.
  */
-export function traverseFiber(fiber: Fiber, ascending: boolean, selector: FiberSelector): Fiber | undefined {
+export function traverseFiber<T = any>(
+  fiber: Fiber<null>,
+  ascending: boolean,
+  selector: FiberSelector<T>,
+): Fiber<T> | undefined {
   let halted = false
-  let selected: Fiber | undefined
+  let selected: Fiber<T> | undefined
 
-  let node = fiber[ascending ? 'return' : 'child']
+  let node = ascending ? fiber.return : fiber.child
   let sibling = fiber.sibling
   while (node) {
     while (sibling) {
-      halted ||= !!selector(sibling)
+      halted ||= selector(sibling) === true
       if (halted) {
         selected = sibling
         break
@@ -31,13 +35,13 @@ export function traverseFiber(fiber: Fiber, ascending: boolean, selector: FiberS
       sibling = sibling.sibling
     }
 
-    halted ||= !!selector(node)
+    halted ||= selector(node) === true
     if (halted) {
       selected = node
       break
     }
 
-    node = node[ascending ? 'return' : 'child']
+    node = ascending ? node.return : node.child
   }
 
   return selected
@@ -54,28 +58,32 @@ const { ReactCurrentOwner } = (React as unknown as ReactInternal).__SECRET_INTER
 /**
  * Returns the current react-internal {@link Fiber}. This is an implementation detail of [react-reconciler](https://github.com/facebook/react/tree/main/packages/react-reconciler).
  */
-export function useFiber(): Fiber {
-  const [fiber] = React.useState<Fiber>(() => ReactCurrentOwner.current!)
+export function useFiber(): Fiber<null> {
+  const [fiber] = React.useState<Fiber<null>>(() => ReactCurrentOwner.current!)
   return fiber
 }
 
 /**
- * Represents a reconciler container.
+ * Represents a react-reconciler container instance.
  */
-export interface Container<T = any> extends Fiber {
-  /** Represents container state passed to {@link ReactReconciler.Reconciler.createContainer}. */
+export interface ContainerInstance<T = any> {
   containerInfo: T
 }
 
 /**
- * Returns the current react-reconciler {@link Container} or the Fiber created from {@link ReactReconciler.Reconciler.createContainer}.
+ * Returns the current react-reconciler container info passed to {@link ReactReconciler.Reconciler.createContainer}.
  *
- * In react-dom, {@link Container.containerInfo} will point to the root DOM element; in react-three-fiber, it will point to the root Zustand store.
+ * In react-dom, a container will point to the root DOM element; in react-three-fiber, it will point to the root Zustand store.
  */
-export function useContainer<T = any>(): Container<T> {
+export function useContainer<T = any>(): T {
   const fiber = useFiber()
   const container = React.useMemo(
-    () => traverseFiber(fiber, true, (node) => node.type == null && node.stateNode.containerInfo != null)!.stateNode,
+    () =>
+      traverseFiber<ContainerInstance<T>>(
+        fiber,
+        true,
+        (node) => node.type == null && node.stateNode?.containerInfo != null,
+      )!.stateNode.containerInfo,
     [fiber],
   )
 
@@ -89,13 +97,13 @@ export function useContainer<T = any>(): Container<T> {
  */
 export function useNearestChild<T = any>(): React.MutableRefObject<T | undefined> {
   const fiber = useFiber()
-  const instance = React.useRef<T>()
+  const childRef = React.useRef<T>()
 
   React.useLayoutEffect(() => {
-    instance.current = traverseFiber(fiber, false, (node) => typeof node.type === 'string')?.stateNode
+    childRef.current = traverseFiber<T>(fiber, false, (node) => typeof node.type === 'string')?.stateNode
   }, [fiber])
 
-  return instance
+  return childRef
 }
 
 /**
@@ -105,13 +113,13 @@ export function useNearestChild<T = any>(): React.MutableRefObject<T | undefined
  */
 export function useNearestParent<T = any>(): React.MutableRefObject<T | undefined> {
   const fiber = useFiber()
-  const instance = React.useRef<T>()
+  const parentRef = React.useRef<T>()
 
   React.useLayoutEffect(() => {
-    instance.current = traverseFiber(fiber, true, (node) => typeof node.type === 'string')?.stateNode
+    parentRef.current = traverseFiber<T>(fiber, true, (node) => typeof node.type === 'string')?.stateNode
   }, [fiber])
 
-  return instance
+  return parentRef
 }
 
 /**

tests/index.test.tsx

@@ -1,11 +1,12 @@
 import * as React from 'react'
 import { describe, expect, it } from 'vitest'
-import { act, render, type HostContainer, type NilNode } from 'react-nil'
+import { type NilNode, type HostContainer, act, render } from 'react-nil'
 import { create } from 'react-test-renderer'
 import {
   type Fiber,
+  type ContainerInstance,
+  traverseFiber,
   useFiber,
-  type Container,
   useContainer,
   useNearestChild,
   useNearestParent,
@@ -22,6 +23,8 @@ interface PrimitiveProps {
   name?: string
 }
 
+type Primitive = NilNode<PrimitiveProps>
+
 declare global {
   namespace JSX {
     interface IntrinsicElements {
@@ -75,27 +78,96 @@ describe('useFiber', () => {
   })
 })
 
+describe('traverseFiber', () => {
+  it('iterates descending through a fiber', async () => {
+    let fiber!: Fiber
+
+    function Test() {
+      fiber = useFiber()
+      return <primitive name="child" />
+    }
+    await act(async () => {
+      render(
+        <primitive name="parent">
+          <Test />
+        </primitive>,
+      )
+    })
+
+    const traversed = [] as unknown as [child: Fiber<Primitive>]
+    traverseFiber(fiber, false, (node) => void traversed.push(node))
+
+    expect(traversed.length).toBe(1)
+
+    const [child] = traversed
+    expect(child.stateNode.props.name).toBe('child')
+  })
+
+  it('iterates ascending through a fiber', async () => {
+    let fiber!: Fiber
+    let container!: HostContainer
+
+    function Test() {
+      fiber = useFiber()
+      return <primitive name="child" />
+    }
+    await act(async () => {
+      container = render(
+        <primitive name="parent">
+          <Test />
+        </primitive>,
+      )
+    })
+
+    const traversed = [] as unknown as [
+      parent: Fiber<Primitive>,
+      rootContainer: Fiber<ContainerInstance<HostContainer>>,
+    ]
+    traverseFiber(fiber, true, (node) => void traversed.push(node))
+
+    expect(traversed.length).toBe(2)
+
+    const [parent, rootContainer] = traversed
+    expect(parent.stateNode.props.name).toBe('parent')
+    expect(rootContainer.stateNode.containerInfo).toBe(container)
+  })
+
+  it('returns the active node when halted', async () => {
+    let fiber!: Fiber
+    let container!: HostContainer
+
+    function Test() {
+      fiber = useFiber()
+      return <primitive name="child" />
+    }
+    await act(async () => (container = render(<Test />)))
+
+    const child = traverseFiber<Primitive>(fiber, false, (node) => node.stateNode === container.head)
+    expect(child!.stateNode.props.name).toBe('child')
+  })
+})
+
 describe('useContainer', () => {
   it('gets the current react-reconciler container', async () => {
-    let currentContainer!: Container<HostContainer>
+    let rootContainer!: HostContainer
     let container!: HostContainer
 
     function Test() {
-      currentContainer = useContainer()
+      rootContainer = useContainer<HostContainer>()
       return null
     }
     await act(async () => (container = render(<Test />)))
 
-    expect(currentContainer.containerInfo).toBe(container)
+    expect(rootContainer).toBe(container)
   })
 })
 
 describe('useNearestChild', () => {
   it('gets the nearest child instance', async () => {
-    const instances: React.MutableRefObject<NilNode<PrimitiveProps> | undefined>[] = []
+    const instances: React.MutableRefObject<Primitive | undefined>[] = []
 
     function Test(props: React.PropsWithChildren) {
-      instances.push(useNearestChild())
+      instances.push(useNearestChild<Primitive>())
       return <>{props.children}</>
     }
 
@@ -130,10 +202,10 @@ describe('useNearestChild', () => {
 
 describe('useNearestParent', () => {
   it('gets the nearest parent instance', async () => {
-    const instances: React.MutableRefObject<NilNode<PrimitiveProps> | undefined>[] = []
+    const instances: React.MutableRefObject<Primitive | undefined>[] = []
 
     function Test(props: React.PropsWithChildren) {
-      instances.push(useNearestParent())
+      instances.push(useNearestParent<Primitive>())
       return <>{props.children}</>
     }