From 1c6bd23989c16f0e25fa161f03430e070b650df0 Mon Sep 17 00:00:00 2001
From: Rich Harris <rich.harris@vercel.com>
Date: Sat, 28 Oct 2023 18:09:33 -0400
Subject: [PATCH 1/3] return undefined if no transformations occurred

---
 src/types.d.ts         |  2 +-
 src/walk.js            |  6 ++++
 test/transformation.js | 71 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 78 insertions(+), 1 deletion(-)

diff --git a/src/types.d.ts b/src/types.d.ts
index ecbae2d..7309527 100644
--- a/src/types.d.ts
+++ b/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;
diff --git a/src/walk.js b/src/walk.js
index ab038e4..a258869 100644
--- a/src/walk.js
+++ b/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;
 				}
 			});
 
diff --git a/test/transformation.js b/test/transformation.js
index 91380da..dc3f66e 100644
--- a/test/transformation.js
+++ b/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 = 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);
+});

From 9eae9043422bdf354ce41c65ebd76cbc26965f58 Mon Sep 17 00:00:00 2001
From: Rich Harris <rich.harris@vercel.com>
Date: Sat, 28 Oct 2023 18:41:12 -0400
Subject: [PATCH 2/3] update README

---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index f88634b..44ef2a8 100644
--- a/README.md
+++ b/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

From 9581c8962938d2025dbe739c8a357edb6fffad0e Mon Sep 17 00:00:00 2001
From: Rich Harris <rich.harris@vercel.com>
Date: Sat, 28 Oct 2023 18:44:00 -0400
Subject: [PATCH 3/3] lint

---
 test/transformation.js | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/transformation.js b/test/transformation.js
index dc3f66e..c5daf23 100644
--- a/test/transformation.js
+++ b/test/transformation.js
@@ -95,7 +95,7 @@ test('returns the result of child transforms when calling next', () => {
 	const transformed = /** @type {import('./types').TestNode} */ (
 		walk(/** @type {import('./types').TestNode} */ (tree), null, {
 			Root: (node, { next }) => {
-				const result = next();
+				const result = /** @type {import('./types').Root} */ (next());
 				children = result.children;
 				return node;
 			},