Skip to content

Commit ab6b3c6

Browse files
committed
feat: implement stack-based traverseFiber
1 parent d722ac5 commit ab6b3c6

File tree

2 files changed

+147
-38
lines changed

2 files changed

+147
-38
lines changed

packages/bippy/src/core.test.tsx

Lines changed: 99 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import type React from 'react';
22
import { render, type RenderOptions } from '@testing-library/react';
33
import { describe, expect, it, vi } from 'vitest';
44
import { instrument, traverseFiber } from './core.js';
5-
import type { FiberRoot } from './types.js';
5+
import type { Fiber, FiberRoot } from './types.js';
66

77
describe('traverseFiber', () => {
88
const onCommitFiberRoot = vi.fn();
@@ -31,19 +31,108 @@ describe('traverseFiber', () => {
3131
);
3232

3333
it('should traverse a fiber', () => {
34-
const handler = vi.fn();
35-
traverseFiber(fiber.current, fiber => handler(fiber.key));
36-
const keys = handler.mock.calls.map(call => call[0]).slice(1);
37-
const expected = ['root', 'a', 'a1', 'a2', 'b', 'c', 'd', 'd1', 'd11'];
38-
expect(keys).toEqual(expected);
34+
const order: string[] = [];
35+
traverseFiber(fiber.current, fiber => {
36+
fiber.key && order.push(fiber.key);
37+
});
38+
expect(order).toEqual([
39+
'root',
40+
'a',
41+
'a1',
42+
'a2',
43+
'b',
44+
'c',
45+
'd',
46+
'd1',
47+
'd11',
48+
]);
3949
});
4050

4151
it('should traverse a fiber in reverse', () => {
42-
const handler = vi.fn();
52+
const order: string[] = [];
4353
const d11 = traverseFiber(fiber.current, fiber => fiber.key === 'd11');
4454
expect(d11?.key).toBe('d11');
45-
traverseFiber(d11, fiber => handler(fiber.key), true);
46-
const keys = handler.mock.calls.map(call => call[0]).slice(0, -1);
47-
expect(keys).toEqual(['d11', 'd1', 'd', 'root']);
55+
56+
traverseFiber(
57+
d11,
58+
fiber => {
59+
fiber.key && order.push(fiber.key);
60+
},
61+
true,
62+
);
63+
expect(order).toEqual(['d11', 'd1', 'd', 'root']);
64+
});
65+
66+
it('should traverse a fiber with entry and leave handlers', () => {
67+
const enterOrder: string[] = [];
68+
const leaveOrder: string[] = [];
69+
traverseFiber(fiber.current, {
70+
enter: fiber => {
71+
fiber.key && enterOrder.push(fiber.key);
72+
},
73+
leave: fiber => {
74+
fiber.key && leaveOrder.push(fiber.key);
75+
},
76+
});
77+
expect(enterOrder).toEqual([
78+
'root',
79+
'a',
80+
'a1',
81+
'a2',
82+
'b',
83+
'c',
84+
'd',
85+
'd1',
86+
'd11',
87+
]);
88+
expect(leaveOrder).toEqual([
89+
'a1',
90+
'a2',
91+
'a',
92+
'b',
93+
'c',
94+
'd11',
95+
'd1',
96+
'd',
97+
'root',
98+
]);
99+
});
100+
101+
it('should traverse a fiber with entry and leave handlers in reverse', () => {
102+
const d11 = traverseFiber(fiber.current, fiber => fiber.key === 'd11');
103+
expect(d11?.key).toBe('d11');
104+
105+
const enterOrder: string[] = [];
106+
const leaveOrder: string[] = [];
107+
traverseFiber(d11, {
108+
ascending: true,
109+
enter: fiber => {
110+
fiber.key && enterOrder.push(fiber.key);
111+
},
112+
leave: fiber => {
113+
fiber.key && leaveOrder.push(fiber.key);
114+
},
115+
});
116+
expect(enterOrder).toEqual(['d11', 'd1', 'd', 'root']);
117+
expect(leaveOrder).toEqual(['root', 'd', 'd1', 'd11']);
118+
});
119+
120+
it('should traverse a fiber and get stack', () => {
121+
const stack: Fiber[] = [];
122+
traverseFiber(fiber.current, {
123+
enter: fiber => {
124+
if (fiber.key === 'd11') {
125+
const keys = stack.map(fiber => fiber.key).filter(Boolean);
126+
expect(keys).toEqual(['root', 'd', 'd1']);
127+
}
128+
129+
stack.push(fiber);
130+
},
131+
leave: fiber => {
132+
const last = stack.pop();
133+
expect(last).toBe(fiber);
134+
},
135+
});
136+
expect(stack).toEqual([]);
48137
});
49138
});

packages/bippy/src/core.ts

Lines changed: 48 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -436,7 +436,8 @@ export const getNearestHostFibers = (fiber: Fiber): Fiber[] => {
436436
return hostFibers;
437437
};
438438

439-
export type FiberSelector = (node: Fiber) => boolean | undefined;
439+
// biome-ignore lint/suspicious/noConfusingVoidType: <explanation>
440+
export type FiberSelector = (node: Fiber) => boolean | void;
440441

441442
export interface TraverseFiberOptions {
442443
/**
@@ -464,30 +465,13 @@ export interface TraverseFiber {
464465
ascending?: boolean,
465466
): Fiber | null;
466467
(fiber: Fiber | null, options: TraverseFiberOptions): Fiber | null;
468+
(
469+
fiber: Fiber | null,
470+
selectorOrOpts: FiberSelector | TraverseFiberOptions,
471+
ascendingOrNever?: boolean,
472+
): Fiber | null;
467473
}
468474

469-
const traverseFiberImpl = (
470-
fiber: Fiber | null,
471-
options: TraverseFiberOptions,
472-
): Fiber | null => {
473-
if (!fiber) return null;
474-
const { enter, leave, ascending } = options;
475-
476-
if (enter && enter(fiber) === true) return fiber;
477-
478-
let child = ascending ? fiber.return : fiber.child;
479-
while (child) {
480-
const match = traverseFiber(child, options);
481-
if (match) return match;
482-
483-
child = ascending ? null : child.sibling;
484-
}
485-
486-
if (leave && leave(fiber) === true) return fiber;
487-
488-
return null;
489-
};
490-
491475
/**
492476
* Traverses up or down a {@link Fiber}, return `true` to stop and select a node.
493477
*/
@@ -497,13 +481,49 @@ export const traverseFiber: TraverseFiber = (
497481
ascendingOrNever = false,
498482
) => {
499483
if (!fiber) return null;
484+
let enter: FiberSelector | undefined;
485+
let leave: FiberSelector | undefined;
486+
let ascending = false;
487+
500488
if (typeof selectorOrOpts === 'function') {
501-
const opts: TraverseFiberOptions = { enter: selectorOrOpts };
502-
if (typeof ascendingOrNever === 'boolean')
503-
opts.ascending = ascendingOrNever;
504-
return traverseFiberImpl(fiber, opts);
489+
enter = selectorOrOpts;
490+
if (typeof ascendingOrNever === 'boolean') ascending = ascendingOrNever;
491+
} else {
492+
enter = selectorOrOpts.enter;
493+
leave = selectorOrOpts.leave;
494+
ascending = selectorOrOpts.ascending ?? false;
505495
}
506-
return traverseFiberImpl(fiber, selectorOrOpts);
496+
497+
const stack: Fiber[] = [fiber];
498+
const visited = new Set<Fiber>();
499+
500+
while (stack.length > 0) {
501+
const current = stack[stack.length - 1];
502+
503+
if (!visited.has(current)) {
504+
visited.add(current);
505+
506+
// Trigger enter handler only once per fiber.
507+
if (enter && enter(current) === true) return current;
508+
509+
// Keep going down the tree. We will back up later.
510+
const next = ascending ? current.return : current.child;
511+
if (next) {
512+
stack.push(next);
513+
continue;
514+
}
515+
}
516+
517+
// Go back to the visited parent fiber and trigger leave handler.
518+
stack.pop();
519+
520+
if (leave && leave(current) === true) return current;
521+
522+
const sibling = ascending ? null : current.sibling;
523+
if (sibling) stack.push(sibling);
524+
}
525+
526+
return null;
507527
};
508528

509529
/**

0 commit comments

Comments
 (0)