perf: traverse layout with an iterator stack

packages/ui/src/layout.test.ts

@@ -0,0 +1,73 @@
+import { close, open, text } from '@bomb.sh/tty';
+import { describe, expect, it } from 'vitest';
+import { Document } from './dom.ts';
+import { Layout } from './layout.ts';
+
+describe('Layout', () => {
+	it('groups adjacent text around nested and empty elements and keeps IDs stable', () => {
+		const document = new Document(() => {});
+		const root = document.documentElement;
+		const child = document.createElement('document');
+		const empty = document.createElement('document');
+		root.appendChild(document.createTextNode('before'));
+		root.appendChild(document.createTextNode(' child'));
+		root.appendChild(child);
+		child.appendChild(document.createTextNode('inside'));
+		root.appendChild(document.createTextNode('after'));
+		root.appendChild(empty);
+		const layout = new Layout(root);
+		const expected = [
+			open('DOCUMENT: 0'),
+			text('before child'),
+			open('DOCUMENT: 1'),
+			text('inside'),
+			close(),
+			text('after'),
+			open('DOCUMENT: 2'),
+			close(),
+			close(),
+		];
+
+		expect([...layout.layout()]).toEqual(expected);
+		expect([...layout.layout()]).toEqual(expected);
+		expect([...layout.layout(child)]).toEqual([open('DOCUMENT: 1'), text('inside'), close()]);
+	});
+
+	it('lays out a deeply nested tree without recursive generator delegation', () => {
+		const document = new Document(() => {});
+		const root = document.documentElement;
+		const depth = 20_000;
+		let current = root;
+		for (let i = 1; i < depth; i++) {
+			const child = document.createElement('document');
+			current.appendChild(child);
+			current = child;
+		}
+		current.appendChild(document.createTextNode('leaf'));
+		const iterator = new Layout(root).layout();
+		for (let i = 0; i < depth; i++) {
+			expect(iterator.next().value).toEqual(open(`DOCUMENT: ${i}`));
+		}
+		expect(iterator.next().value).toEqual(text('leaf'));
+		for (let i = 0; i < depth; i++) {
+			expect(iterator.next().value).toEqual(close());
+		}
+		expect(iterator.next().done).toBe(true);
+	});
+
+	it('closes active elements when iteration is returned early', () => {
+		const document = new Document(() => {});
+		const root = document.documentElement;
+		const child = document.createElement('document');
+		root.appendChild(child);
+		child.appendChild(document.createTextNode('inside'));
+		const iterator = new Layout(root).layout();
+		iterator.next();
+		iterator.next();
+		iterator.next();
+
+		expect(iterator.return!()).toEqual({ value: close(), done: false });
+		expect(iterator.next()).toEqual({ value: close(), done: false });
+		expect(iterator.next().done).toBe(true);
+	});
+});

packages/ui/src/layout.ts

@@ -12,31 +12,48 @@ export class Layout {
 
 	/** Serialize an element to TTY Ops. */
 	*layout(element: Element = this.#root): IterableIterator<Op> {
-		let stableId = this.#stableIds.get(element);
-		if (!stableId) {
-			stableId = `${element.nodeName}: ${this.#lastIdNumber++}`;
-			this.#stableIds.set(element, stableId);
-		}
-		yield open(stableId);
+		yield this.#open(element);
 
+		const stack = [element.childNodes[Symbol.iterator]()];
 		try {
 			let content = '';
-			for (const child of element.childNodes) {
+			while (stack.length > 0) {
+				const next = stack[stack.length - 1]!.next();
+				if (next.done) {
+					if (content !== '') {
+						yield text(content);
+						content = '';
+					}
+					stack.pop();
+					yield close();
+					continue;
+				}
+				const child = next.value;
 				if (child instanceof Text) {
 					content += child.data;
 				} else if (child instanceof Element) {
 					if (content !== '') {
 						yield text(content);
 						content = '';
 					}
-					yield* this.layout(child);
+					yield this.#open(child);
+					stack.push(child.childNodes[Symbol.iterator]());
 				}
 			}
-			if (content !== '') {
-				yield text(content);
-			}
 		} finally {
-			yield close();
+			while (stack.length > 0) {
+				stack.pop();
+				yield close();
+			}
+		}
+	}
+
+	#open(element: Element): Op {
+		let stableId = this.#stableIds.get(element);
+		if (!stableId) {
+			stableId = `${element.nodeName}: ${this.#lastIdNumber++}`;
+			this.#stableIds.set(element, stableId);
 		}
+		return open(stableId);
 	}
 }