feat: implement stack-based traverseFiber

Author: Asuka109

packages/bippy/src/core.test.tsx

@@ -2,7 +2,7 @@ import type React from 'react';
 import { render, type RenderOptions } from '@testing-library/react';
 import { describe, expect, it, vi } from 'vitest';
 import { instrument, traverseFiber } from './core.js';
-import type { FiberRoot } from './types.js';
+import type { Fiber, FiberRoot } from './types.js';
 
 describe('traverseFiber', () => {
   const onCommitFiberRoot = vi.fn();
@@ -31,19 +31,108 @@ describe('traverseFiber', () => {
   );
 
   it('should traverse a fiber', () => {
-    const handler = vi.fn();
-    traverseFiber(fiber.current, fiber => handler(fiber.key));
-    const keys = handler.mock.calls.map(call => call[0]).slice(1);
-    const expected = ['root', 'a', 'a1', 'a2', 'b', 'c', 'd', 'd1', 'd11'];
-    expect(keys).toEqual(expected);
+    const order: string[] = [];
+    traverseFiber(fiber.current, fiber => {
+      fiber.key && order.push(fiber.key);
+    });
+    expect(order).toEqual([
+      'root',
+      'a',
+      'a1',
+      'a2',
+      'b',
+      'c',
+      'd',
+      'd1',
+      'd11',
+    ]);
   });
 
   it('should traverse a fiber in reverse', () => {
-    const handler = vi.fn();
+    const order: string[] = [];
     const d11 = traverseFiber(fiber.current, fiber => fiber.key === 'd11');
     expect(d11?.key).toBe('d11');
-    traverseFiber(d11, fiber => handler(fiber.key), true);
-    const keys = handler.mock.calls.map(call => call[0]).slice(0, -1);
-    expect(keys).toEqual(['d11', 'd1', 'd', 'root']);
+
+    traverseFiber(
+      d11,
+      fiber => {
+        fiber.key && order.push(fiber.key);
+      },
+      true,
+    );
+    expect(order).toEqual(['d11', 'd1', 'd', 'root']);
+  });
+
+  it('should traverse a fiber with entry and leave handlers', () => {
+    const enterOrder: string[] = [];
+    const leaveOrder: string[] = [];
+    traverseFiber(fiber.current, {
+      enter: fiber => {
+        fiber.key && enterOrder.push(fiber.key);
+      },
+      leave: fiber => {
+        fiber.key && leaveOrder.push(fiber.key);
+      },
+    });
+    expect(enterOrder).toEqual([
+      'root',
+      'a',
+      'a1',
+      'a2',
+      'b',
+      'c',
+      'd',
+      'd1',
+      'd11',
+    ]);
+    expect(leaveOrder).toEqual([
+      'a1',
+      'a2',
+      'a',
+      'b',
+      'c',
+      'd11',
+      'd1',
+      'd',
+      'root',
+    ]);
+  });
+
+  it('should traverse a fiber with entry and leave handlers in reverse', () => {
+    const d11 = traverseFiber(fiber.current, fiber => fiber.key === 'd11');
+    expect(d11?.key).toBe('d11');
+
+    const enterOrder: string[] = [];
+    const leaveOrder: string[] = [];
+    traverseFiber(d11, {
+      ascending: true,
+      enter: fiber => {
+        fiber.key && enterOrder.push(fiber.key);
+      },
+      leave: fiber => {
+        fiber.key && leaveOrder.push(fiber.key);
+      },
+    });
+    expect(enterOrder).toEqual(['d11', 'd1', 'd', 'root']);
+    expect(leaveOrder).toEqual(['root', 'd', 'd1', 'd11']);
+  });
+
+  it('should traverse a fiber and get stack', () => {
+    const stack: Fiber[] = [];
+    traverseFiber(fiber.current, {
+      enter: fiber => {
+        if (fiber.key === 'd11') {
+          const keys = stack.map(fiber => fiber.key).filter(Boolean);
+          expect(keys).toEqual(['root', 'd', 'd1']);
+        }
+
+        stack.push(fiber);
+      },
+      leave: fiber => {
+        const last = stack.pop();
+        expect(last).toBe(fiber);
+      },
+    });
+    expect(stack).toEqual([]);
   });
 });

packages/bippy/src/core.ts

@@ -436,7 +436,8 @@ export const getNearestHostFibers = (fiber: Fiber): Fiber[] => {
   return hostFibers;
 };
 
-export type FiberSelector = (node: Fiber) => boolean | undefined;
+// biome-ignore lint/suspicious/noConfusingVoidType: <explanation>
+export type FiberSelector = (node: Fiber) => boolean | void;
 
 export interface TraverseFiberOptions {
   /**
@@ -464,30 +465,13 @@ export interface TraverseFiber {
     ascending?: boolean,
   ): Fiber | null;
   (fiber: Fiber | null, options: TraverseFiberOptions): Fiber | null;
+  (
+    fiber: Fiber | null,
+    selectorOrOpts: FiberSelector | TraverseFiberOptions,
+    ascendingOrNever?: boolean,
+  ): Fiber | null;
 }
 
-const traverseFiberImpl = (
-  fiber: Fiber | null,
-  options: TraverseFiberOptions,
-): Fiber | null => {
-  if (!fiber) return null;
-  const { enter, leave, ascending } = options;
-
-  if (enter && enter(fiber) === true) return fiber;
-
-  let child = ascending ? fiber.return : fiber.child;
-  while (child) {
-    const match = traverseFiber(child, options);
-    if (match) return match;
-
-    child = ascending ? null : child.sibling;
-  }
-
-  if (leave && leave(fiber) === true) return fiber;
-
-  return null;
-};
-
 /**
  * Traverses up or down a {@link Fiber}, return `true` to stop and select a node.
  */
@@ -497,13 +481,49 @@ export const traverseFiber: TraverseFiber = (
   ascendingOrNever = false,
 ) => {
   if (!fiber) return null;
+  let enter: FiberSelector | undefined;
+  let leave: FiberSelector | undefined;
+  let ascending = false;
+
   if (typeof selectorOrOpts === 'function') {
-    const opts: TraverseFiberOptions = { enter: selectorOrOpts };
-    if (typeof ascendingOrNever === 'boolean')
-      opts.ascending = ascendingOrNever;
-    return traverseFiberImpl(fiber, opts);
+    enter = selectorOrOpts;
+    if (typeof ascendingOrNever === 'boolean') ascending = ascendingOrNever;
+  } else {
+    enter = selectorOrOpts.enter;
+    leave = selectorOrOpts.leave;
+    ascending = selectorOrOpts.ascending ?? false;
   }
-  return traverseFiberImpl(fiber, selectorOrOpts);
+
+  const stack: Fiber[] = [fiber];
+  const visited = new Set<Fiber>();
+
+  while (stack.length > 0) {
+    const current = stack[stack.length - 1];
+
+    if (!visited.has(current)) {
+      visited.add(current);
+
+      // Trigger enter handler only once per fiber.
+      if (enter && enter(current) === true) return current;
+
+      // Keep going down the tree. We will back up later.
+      const next = ascending ? current.return : current.child;
+      if (next) {
+        stack.push(next);
+        continue;
+      }
+    }
+
+    // Go back to the visited parent fiber and trigger leave handler.
+    stack.pop();
+
+    if (leave && leave(current) === true) return current;
+
+    const sibling = ascending ? null : current.sibling;
+    if (sibling) stack.push(sibling);
+  }
+
+  return null;
 };
 
 /**