return result from context.next()

README.md

@@ -123,7 +123,7 @@ walk(node, state, visitors);
 
 Each visitor receives a second argument, `context`, which is an object with the following properties and methods:
 
-- `next(state?: State): void` — a function that allows you to control when child nodes are visited, and which state they are visited with
+- `next(state?: State): void` — a function that allows you to control when child nodes are visited, and which state they are visited with. If child visitors transform their inputs, this will return the transformed node (if not, returns `undefined`)
 - `path: Node[]` — an array of parent nodes. For example, to get the root node you would do `path.at(0)`; to get the current node's immediate parent you would do `path.at(-1)`
 - `state: State` — an object of the same type as the second argument to `walk`. Visitors can pass new state objects to their children with `next(childState)` or `visit(node, childState)`
 - `stop(): void` — prevents any subsequent traversal

src/types.d.ts

@@ -13,7 +13,7 @@ export type Visitors<T extends BaseNode, U> = T['type'] extends '_'
 	: SpecialisedVisitors<T, U> & { _?: Visitor<T, U, T> };
 
 export interface Context<T, U> {
-	next: (state?: U) => void;
+	next: (state?: U) => T | void;
 	path: T[];
 	state: U;
 	stop: () => void;

src/walk.js

@@ -73,6 +73,10 @@ export function walk(node, state, visitors) {
 					}
 				}
 				path.pop();
+
+				if (Object.keys(mutations).length > 0) {
+					return { ...node, ...mutations };
+				}
 			},
 			stop: () => {
 				stopped = true;
@@ -103,6 +107,8 @@ export function walk(node, state, visitors) {
 						...context,
 						state: next_state
 					});
+
+					return inner_result;
 				}
 			});
 

test/transformation.js

@@ -81,3 +81,74 @@ test('respects individual visitors if universal visitor calls next()', () => {
 		children: [{ type: 'TransformedA' }, { type: 'B' }, { type: 'C' }]
 	});
 });
+
+test('returns the result of child transforms when calling next', () => {
+	/** @type {import('./types').TestNode} */
+	const tree = {
+		type: 'Root',
+		children: [{ type: 'A' }, { type: 'B' }, { type: 'C' }]
+	};
+
+	let count = 0;
+	let children;
+
+	const transformed = /** @type {import('./types').TestNode} */ (
+		walk(/** @type {import('./types').TestNode} */ (tree), null, {
+			Root: (node, { next }) => {
+				const result = /** @type {import('./types').Root} */ (next());
+				children = result.children;
+				return node;
+			},
+			A: (node) => {
+				count += 1;
+				return {
+					type: 'TransformedA'
+				};
+			},
+			C: (node) => {
+				count += 1;
+				return {
+					type: 'TransformedC'
+				};
+			}
+		})
+	);
+
+	expect(count).toBe(2);
+
+	// check that `tree` wasn't mutated
+	expect(tree).toEqual({
+		type: 'Root',
+		children: [{ type: 'A' }, { type: 'B' }, { type: 'C' }]
+	});
+
+	expect(transformed).toBe(tree);
+
+	expect(children).toEqual([
+		{ type: 'TransformedA' },
+		{ type: 'B' },
+		{ type: 'TransformedC' }
+	]);
+});
+
+test('returns undefined if there are no child transformations', () => {
+	/** @type {import('./types').TestNode} */
+	const tree = {
+		type: 'Root',
+		children: [{ type: 'A' }, { type: 'B' }, { type: 'C' }]
+	};
+
+	let result;
+
+	const transformed = /** @type {import('./types').TestNode} */ (
+		walk(/** @type {import('./types').TestNode} */ (tree), null, {
+			Root: (node, { next }) => {
+				result = next();
+			}
+		})
+	);
+
+	expect(transformed).toBe(tree);
+
+	expect(result).toBe(undefined);
+});