gofront 1.5.1 → 1.6.0

package/src/compiler.js

@@ -18,9 +18,13 @@
 // de-qualifies it because the dependency is inlined.
 
 import { readdirSync, readFileSync, statSync } from "node:fs";
-import { basename, dirname, join, relative, resolve } from "node:path";
-import { buildSourceMap, CodeGen } from "./backend/js/index.js";
-import { compileWasmModule } from "./backend/wasm/index.js";
+import { basename, dirname, join, relative, resolve, sep } from "node:path";
+import {
+	buildSourceMap,
+	buildWasmSourceMap,
+	CodeGen,
+} from "./backend/js/index.js";
+import { compileWasmModule, optimizeWasm } from "./backend/wasm/index.js";
 import { log } from "./colors.js";
 import { parseDts } from "./dts-parser.js";
 import { Lexer } from "./lexer.js";
@@ -67,6 +71,9 @@ function parseGoFrontFile(filePath) {
 		? new TemplParser(tokens, filename, source).parse()
 		: new Parser(tokens, filename, source).parse();
 	ast._source = source;
+	ast._path = resolve(filePath);
+	// compileFiles overwrites `target` with the resolved one; keep the directive.
+	ast._declaredTarget = ast.target;
 	_parseCache.set(filePath, { mtime, ast });
 	return ast;
 }
@@ -163,6 +170,48 @@ function attachSourceMap(merged, outputDir) {
 	merged.js += `\n//# sourceMappingURL=data:application/json;base64,${b64}`;
 }
 
+const WASM_MAP_URL = "app.wasm.map";
+
+// Function-level `app.wasm.map` (sources relative to the output dir).
+function buildWasmMap(info, outputDir) {
+	if (!info || info.mappings.length === 0) return null;
+	return buildWasmSourceMap(
+		info.sources.map((f) => relative(outputDir, f)),
+		info.mappings,
+	);
+}
+
+// Shared tail of compileFiles / compileSingleFile: link `app.wasm`, run
+// wasm-opt when requested and produce the wasm source map.
+function finishWasm(options, merged, pkgTarget, outputDir) {
+	const wantMap = Boolean(options.sourceMap && !options.isDependency);
+	let { wasm, wat, sourceMapInfo } = linkIntoBundle(
+		options,
+		merged,
+		pkgTarget,
+		{
+			sourceMappingURL: wantMap ? WASM_MAP_URL : null,
+		},
+	);
+	let wasmOptInfo = null;
+	let wasmSourceMap =
+		wantMap && wasm ? buildWasmMap(sourceMapInfo, outputDir) : null;
+
+	if (wasm && options.wasmOpt) {
+		const opt = optimizeWasm(wasm, {
+			sourceMap: Boolean(wasmSourceMap),
+			inputSourceMap: wasmSourceMap,
+			sourceMapUrl: WASM_MAP_URL,
+			emitWat: Boolean(options.emitWat),
+		});
+		wasm = opt.wasm;
+		if (opt.wat) wat = opt.wat;
+		wasmOptInfo = opt;
+		if (opt.sourceMap) wasmSourceMap = opt.sourceMap;
+	}
+	return { wasm, wat, wasmOptInfo, wasmSourceMap };
+}
+
 // ── WASM linking ──────────────────────────────────────────────
 //
 // Every `wasm`/`both` package reached during a build registers a unit in the
@@ -205,7 +254,10 @@ function registerPackageUnit(options, unit) {
 	return emitTarget;
 }
 
-function linkWasmUnits(units, { emitWat = false, rootTarget = "js" } = {}) {
+function linkWasmUnits(
+	units,
+	{ emitWat = false, rootTarget = "js", sourceMappingURL = null } = {},
+) {
 	if (!units.some((u) => u.target === "wasm")) return null;
 	checkLinkCollisions(units);
 
@@ -224,6 +276,7 @@ function linkWasmUnits(units, { emitWat = false, rootTarget = "js" } = {}) {
 		bundledPackages,
 		emitWat,
 		callMain: rootTarget === "wasm",
+		sourceMappingURL,
 	});
 }
 
@@ -291,17 +344,23 @@ function spliceFacade(merged, facade) {
 }
 
 // Root compiles only: link every registered unit into app.wasm and splice the
-// facade into the merged bundle.  Returns `{ wasm, wat }` (nulls when there
-// is nothing to link).
-function linkIntoBundle(options, merged, pkgTarget) {
-	if (options.isDependency) return { wasm: null, wat: null };
+// facade into the merged bundle.  Returns `{ wasm, wat, sourceMapInfo }`
+// (nulls when there is nothing to link).
+function linkIntoBundle(options, merged, pkgTarget, { sourceMappingURL } = {}) {
+	const none = { wasm: null, wat: null, sourceMapInfo: null };
+	if (options.isDependency) return none;
 	const linked = linkWasmUnits(options.wasmUnits, {
 		emitWat: options.emitWat,
 		rootTarget: pkgTarget,
+		sourceMappingURL,
 	});
-	if (!linked) return { wasm: null, wat: null };
+	if (!linked) return none;
 	spliceFacade(merged, linked.facade);
-	return { wasm: linked.wasm, wat: linked.wat ?? null };
+	return {
+		wasm: linked.wasm,
+		wat: linked.wat ?? null,
+		sourceMapInfo: linked.sourceMapInfo ?? null,
+	};
 }
 
 // ── Import resolution ─────────────────────────────────────────
@@ -309,7 +368,7 @@ function linkIntoBundle(options, merged, pkgTarget) {
 // Shared by compileFiles (multi-file) and the single-file path in index.js.
 // Mutates checker, jsImports, bundledPackages, and preambles in place.
 
-export function resolveImports(
+function resolveImports(
 	programs,
 	fromFile,
 	checker,
@@ -350,6 +409,11 @@ export function resolveImports(
 							"package 'gom' is not available in wasm packages",
 							impNode,
 						);
+					} else if (path === "gofront/shared" && pkgTarget === "both") {
+						checker.err(
+							"'gofront/shared' requires a wasm-only package (//gofront:target wasm): a both package has no linear memory on the JS side",
+							impNode,
+						);
 					} else if (!isLocalPath(path)) {
 						if (!isBuiltinPackage(path)) {
 							checker.recordBlocker("js: import", path);
@@ -472,6 +536,7 @@ export function resolveImports(
 					nameUsed,
 					dep.exportedSymbols,
 					dep.exportedTypes,
+					dep.target,
 				);
 				checker.trackImport(nameUsed, { _line }, p._filename, p._source);
 			}
@@ -481,6 +546,29 @@ export function resolveImports(
 
 // ── Main entry points ─────────────────────────────────────────
 
+// Package target after build-time overrides (Task H6.3): `forceTarget`
+// (`--js-only`) compiles every package to JS; `targetOverrides` maps package
+// directories relative to the root source dir (`"engine/physics"`, `"."` for
+// the root) to a target, replacing the `//gofront:target` directive.
+const PACKAGE_TARGETS = new Set(["js", "wasm", "both"]);
+
+export function resolvePackageTarget(declared, pkgDir, options = {}) {
+	if (options.forceTarget) return options.forceTarget;
+	const overrides = options.targetOverrides;
+	if (!overrides || !options.rootDir) return declared;
+	const rel = relative(resolve(options.rootDir), resolve(pkgDir))
+		.split(sep)
+		.join("/");
+	const key = rel === "" ? "." : rel;
+	const t = overrides[key] ?? (key === "." ? overrides[""] : undefined);
+	if (t === undefined) return declared;
+	if (!PACKAGE_TARGETS.has(t))
+		throw new Error(
+			`gofront.json: invalid target '${t}' for package '${key}' (expected js, wasm or both)`,
+		);
+	return t;
+}
+
 export function compileSingleFile(inputPath, options = {}) {
 	const {
 		sourceMap = false,
@@ -502,6 +590,7 @@ export function compileSingleFile(inputPath, options = {}) {
 
 	const ast = new Parser(tokens, basename(inputPath), source).parse();
 	ast._source = source;
+	ast._path = resolve(inputPath);
 	if (dumpAst) return { ast };
 
 	const outputDir =
@@ -513,7 +602,11 @@ export function compileSingleFile(inputPath, options = {}) {
 	const bundledPackages = new Set();
 	const preambles = [];
 
-	const pkgTarget = ast.target ?? options.target ?? "js";
+	const pkgTarget = resolvePackageTarget(
+		ast.target ?? options.target ?? "js",
+		dirname(resolve(inputPath)),
+		options,
+	);
 	ast.target = pkgTarget;
 	checker.target = pkgTarget;
 	checker.pkgName = ast.pkg?.name ?? "main";
@@ -534,6 +627,7 @@ export function compileSingleFile(inputPath, options = {}) {
 	if (errors.length > 0) {
 		throw new Error(errors.map((e) => e.message).join("\n"));
 	}
+	for (const w of checker.warnings) log.warn(w);
 
 	if (pkgTarget === "wasm" || pkgTarget === "both") {
 		registerWasmUnit(options, {
@@ -562,13 +656,26 @@ export function compileSingleFile(inputPath, options = {}) {
 	};
 
 	const merged = mergeCompilationChunks([...preambles, mainChunk]);
-	const { wasm, wat } = linkIntoBundle(options, merged, pkgTarget);
+	const { wasm, wat, wasmOptInfo, wasmSourceMap } = finishWasm(
+		{ ...options, sourceMap },
+		merged,
+		pkgTarget,
+		outputDir,
+	);
 
 	if (sourceMap && !options.isDependency) {
 		attachSourceMap(merged, outputDir);
 	}
 
-	return { js: merged.js, css: cg.getCss(), target: pkgTarget, wasm, wat };
+	return {
+		js: merged.js,
+		css: cg.getCss(),
+		target: pkgTarget,
+		wasm,
+		wat,
+		wasmOptInfo,
+		wasmSourceMap,
+	};
 }
 
 export function compileDir(dir, options = {}) {
@@ -591,8 +698,12 @@ export function compilePackageTestsWasm(dir, options = {}) {
 }
 
 export function compileFiles(files, options = {}) {
-	options = { ...options, wasmUnits: options.wasmUnits ?? [] };
 	const fromDir = options.fromDir ?? dirname(resolve(files[0]));
+	options = {
+		...options,
+		wasmUnits: options.wasmUnits ?? [],
+		rootDir: options.rootDir ?? fromDir,
+	};
 	const outputDir =
 		options.outputDir ??
 		(options.outputFile ? dirname(resolve(options.outputFile)) : fromDir);
@@ -622,9 +733,10 @@ export function compileFiles(files, options = {}) {
 	// Validate target directives across files in this package
 	const declaredTargets = new Map();
 	for (const p of programs) {
-		if (p.target) {
-			if (!declaredTargets.has(p.target)) declaredTargets.set(p.target, []);
-			declaredTargets.get(p.target).push(p._filename);
+		const declared = p._declaredTarget;
+		if (declared) {
+			if (!declaredTargets.has(declared)) declaredTargets.set(declared, []);
+			declaredTargets.get(declared).push(p._filename);
 		}
 	}
 	if (declaredTargets.size > 1) {
@@ -635,10 +747,13 @@ export function compileFiles(files, options = {}) {
 			`Conflicting //gofront:target directives in package '${pkgName}': ${list}`,
 		);
 	}
-	const pkgTarget =
+	const pkgTarget = resolvePackageTarget(
 		declaredTargets.size === 1
 			? declaredTargets.keys().next().value
-			: (options.target ?? "js");
+			: (options.target ?? "js"),
+		fromDir,
+		options,
+	);
 	for (const p of programs) p.target = pkgTarget;
 
 	// ── 2. Resolve imports ────────────────────────────────────────
@@ -668,6 +783,7 @@ export function compileFiles(files, options = {}) {
 		const msgs = errors.map((e) => e.message).join("\n");
 		throw new Error(msgs);
 	}
+	for (const w of checker.warnings) log.warn(w);
 
 	// ── 4. Code generation ────────────────────────────────────────
 	const emitTarget = registerPackageUnit(options, {
@@ -699,7 +815,12 @@ export function compileFiles(files, options = {}) {
 	};
 
 	const merged = mergeCompilationChunks([...preambles, mainChunk]);
-	const { wasm, wat } = linkIntoBundle(options, merged, pkgTarget);
+	const { wasm, wat, wasmOptInfo, wasmSourceMap } = finishWasm(
+		options,
+		merged,
+		pkgTarget,
+		outputDir,
+	);
 
 	if (options.sourceMap && !options.isDependency) {
 		attachSourceMap(merged, outputDir);
@@ -717,6 +838,8 @@ export function compileFiles(files, options = {}) {
 		css: allCss,
 		wasm,
 		wat,
+		wasmOptInfo,
+		wasmSourceMap,
 		programs,
 		exportedSymbols: checker.getExportedSymbols(),
 		exportedTypes: checker.getExportedTypes(),

package/src/index.js

@@ -18,7 +18,7 @@ import { createRequire } from "node:module";
 const _require = createRequire(import.meta.url);
 const { version } = _require("../package.json");
 
-import { statSync, watch } from "node:fs";
+import { statSync } from "node:fs";
 import { dirname, relative, resolve } from "node:path";
 import { copyAssets } from "./asset-manager.js";
 import {
@@ -35,13 +35,13 @@ import {
 	parseBuildArgs,
 	parseCheckArgs,
 	parseDevArgs,
+	parseLegacyArgs,
 	parsePrepArgs,
 	parseTestArgs,
 	runCompile,
 	writeCompileOutput,
 } from "./cli-core.js";
 import { colors, formatDiagnostic, log, ms, stamp } from "./colors.js";
-import { createDevServer } from "./dev-server.js";
 
 // ── Parse CLI args ───────────────────────────────────────────
 
@@ -54,11 +54,11 @@ if (args[0] === "--version" || args[0] === "-v") {
 if (args.length === 0 || args[0] === "--help" || args[0] === "-h") {
 	console.log(
 		`
-GoFront — a Go-inspired language that compiles to JavaScript
+GoFront — a Go-inspired language that compiles to JavaScript and WebAssembly
 
 Usage:
   gofront dev [dir] [options]    Start dev server with live reload (default port 3000)
-  gofront build [dir] [options]  Build production bundle (-o <dir>, --pwa, --minify, --emit-wat)
+  gofront build [dir] [options]  Build production bundle (-o <dir>, --pwa, --minify, --emit-wat, --release, --wasm-opt, --js-only)
   gofront check [dir|dir/...]    Type-check only (dir/... recurses into every package)
   gofront test [dir|dir/...] [--dom]  Run unit tests (-v verbose, -run <regex>)
   gofront prep [dir] [--minify]  Copy static assets and bundle vendor dependencies
@@ -66,6 +66,8 @@ Usage:
   gofront <file.go>              Compile single file and print to stdout
   gofront <dir>  (or gofront .)  Compile all *.go in directory as one bundle
   gofront <input> -o out.js      Compile and write to file (wasm packages also emit app.wasm)
+  gofront <input> --release      Production build with minification and Binaryen wasm optimization
+  gofront <input> --wasm-opt     Optimize wasm output with Binaryen (-O3 + GUFA)
   gofront <input> --check        Type-check only
   gofront <input> --watch        Watch for changes and recompile
   gofront <input> -o out.js --serve          Watch + serve with live reload (default port 3000)
@@ -131,9 +133,9 @@ if (args[0] === "test") {
 // ── check subcommand ──────────────────────────────────────────
 
 if (args[0] === "check") {
-	const { targetDir } = parseCheckArgs(args.slice(1));
+	const { targetDir, forceTarget } = parseCheckArgs(args.slice(1));
 	try {
-		const { elapsedMs, packages } = handleCheck(targetDir);
+		const { elapsedMs, packages } = handleCheck(targetDir, { forceTarget });
 		const okLine = (label, t) =>
 			log.info(`${colors.cyan(label)} — ${colors.green("OK")} ${ms(t)}`);
 		if (packages) {
@@ -202,20 +204,22 @@ if (args[0] === "dev") {
 	}
 }
 
-const inputArg = args[0];
-const outputFlag = args.indexOf("-o");
-const outputFile = outputFlag !== -1 ? args[outputFlag + 1] : null;
-const checkOnly = args.includes("--check");
-const dumpAst = args.includes("--ast");
-const dumpTokens = args.includes("--tokens");
-const sourceMap = args.includes("--source-map");
-const serveMode = args.includes("--serve");
-const watchMode = args.includes("--watch") || serveMode;
-const copyAssetsFlag = args.includes("--copy-assets");
-const minifyOutput = args.includes("--minify");
-const mangleOutput = args.includes("--mangle");
-const portFlag = args.indexOf("--port");
-const servePort = portFlag !== -1 ? parseInt(args[portFlag + 1], 10) : 3000;
+const {
+	inputArg,
+	outputFile,
+	checkOnly,
+	dumpAst,
+	dumpTokens,
+	sourceMap,
+	serve: serveMode,
+	watch: watchMode,
+	copyAssets: copyAssetsFlag,
+	wasmOpt,
+	minify: minifyOutput,
+	mangle: mangleOutput,
+	port: servePort,
+	forceTarget,
+} = parseLegacyArgs(args);
 
 // ── Determine input mode ─────────────────────────────────────
 
@@ -228,6 +232,17 @@ try {
 	process.exit(1);
 }
 
+function copyProjectAssets(verbose) {
+	try {
+		const { copied, skipped } = copyAssets(resolve("."));
+		if (verbose && (copied > 0 || skipped > 0)) {
+			log.info(`copied ${copied} assets (${skipped} skipped)`);
+		}
+	} catch (e) {
+		log.warn(`asset copy failed: ${e.message}`);
+	}
+}
+
 // ── Single-shot mode ─────────────────────────────────────────
 
 if (!watchMode) {
@@ -239,6 +254,8 @@ if (!watchMode) {
 			outputFile,
 			dumpTokens,
 			dumpAst,
+			wasmOpt,
+			forceTarget,
 		});
 	} catch (e) {
 		log.fail(e.message);
@@ -282,6 +299,13 @@ if (!watchMode) {
 			log.info(
 				`wrote ${formatWrittenDesc(written, outputFile)} ${ms(elapsedMs)}`,
 			);
+			if (result.wasmOptInfo && result.wasmOptInfo.originalSize > 0) {
+				const origKb = (result.wasmOptInfo.originalSize / 1024).toFixed(1);
+				const optKb = (result.wasmOptInfo.optimizedSize / 1024).toFixed(1);
+				log.info(
+					`wasm-opt: app.wasm ${origKb} kB → ${optKb} kB (-${result.wasmOptInfo.percentSaved}%) [${result.wasmOptInfo.engine}]`,
+				);
+			}
 		} catch (e) {
 			log.fail(`cannot write '${outputFile}': ${e.message}`);
 			process.exit(1);
@@ -295,110 +319,57 @@ if (!watchMode) {
 		}
 	}
 
-	if (copyAssetsFlag) {
-		try {
-			const projectDir = resolve(".");
-			const { copied, skipped } = copyAssets(projectDir);
-			if (copied > 0 || skipped > 0) {
-				log.info(`copied ${copied} assets (${skipped} skipped)`);
-			}
-		} catch (e) {
-			log.warn(`asset copy failed: ${e.message}`);
-		}
-	}
+	if (copyAssetsFlag) copyProjectAssets(true);
 
 	process.exit(0);
 }
 
 // ── Watch mode ───────────────────────────────────────────────
 
-// Start dev server before first build so the browser can connect immediately
-let devServer = null;
-if (serveMode) {
-	if (!outputFile) {
-		log.fail("--serve requires -o <output file>");
-		process.exit(1);
-	}
-	const serveDir = dirname(resolve(outputFile));
-	devServer = createDevServer(serveDir, servePort);
+if (serveMode && !outputFile) {
+	log.fail("--serve requires -o <output file>");
+	process.exit(1);
 }
 
-function buildOnce(changedFile = null) {
-	try {
-		const startMs = performance.now();
-		const result = runCompile(inputPath, isDir, { sourceMap, outputFile });
-		const js = maybeMinify(result.js, {
-			minify: minifyOutput,
-			mangle: mangleOutput,
-			sourceMap,
-		});
-		const elapsedMs = (performance.now() - startMs).toFixed(0);
-		const changeNote = changedFile ? ` — ${changedFile} changed` : "";
-		const timing = colors.dim(`(${elapsedMs}ms${changeNote})`);
-		if (outputFile) {
-			const written = writeCompileOutput(resolve(outputFile), result, js);
-			console.error(
-				`${stamp()} ${colors.bold("gofront:")} ${colors.green("OK")} — wrote ${formatWrittenDesc(written, outputFile)} ${timing}`,
-			);
-		} else {
-			// Clear screen then print
-			process.stdout.write("\x1Bc");
-			console.log(js);
-			console.error(
-				`${stamp()} ${colors.bold("gofront:")} ${colors.green("OK")} ${timing}`,
-			);
-		}
+const prefix = () => `${stamp()} ${colors.bold("gofront:")}`;
 
-		if (copyAssetsFlag) {
-			try {
-				copyAssets(resolve("."));
-			} catch (e) {
-				log.warn(`asset copy failed: ${e.message}`);
+try {
+	await handleDev(".", {
+		srcDir: inputPath,
+		outputFile,
+		serve: serveMode,
+		serveDir: outputFile ? dirname(resolve(outputFile)) : undefined,
+		port: servePort,
+		silent: false,
+		sourceMap,
+		minify: minifyOutput,
+		mangle: mangleOutput,
+		copyAssets: copyAssetsFlag,
+		forceTarget,
+		onBuild: ({ js, written, elapsedMs, changedFile }) => {
+			const changeNote = changedFile ? ` — ${changedFile} changed` : "";
+			const timing = colors.dim(`(${elapsedMs}ms${changeNote})`);
+			if (outputFile) {
+				console.error(
+					`${prefix()} ${colors.green("OK")} — wrote ${formatWrittenDesc(written, outputFile)} ${timing}`,
+				);
+			} else {
+				// Clear screen then print
+				process.stdout.write("\x1Bc");
+				console.log(js);
+				console.error(`${prefix()} ${colors.green("OK")} ${timing}`);
 			}
-		}
-
-		devServer?.notify();
-	} catch (e) {
-		console.error(
-			`${stamp()} ${colors.bold("gofront:")} ${colors.bold(colors.red("ERROR"))}`,
-		);
-		for (const line of formatDiagnostic(e.message).split("\n"))
-			console.error(`  ${line}`);
-		devServer?.notifyError?.(e);
-	}
-}
-
-// Initial build
-buildOnce();
-
-// Determine what to watch
-const watchTarget = isDir ? inputPath : dirname(inputPath);
-
-function handleCssWatch(filename) {
-	if (copyAssetsFlag) {
-		try {
-			copyAssets(resolve("."));
-		} catch (e) {
-			log.warn(`asset copy failed: ${e.message}`);
-		}
-	}
-	devServer?.notifyCss?.(filename);
+			if (copyAssetsFlag) copyProjectAssets(false);
+		},
+		onError: (e) => {
+			console.error(`${prefix()} ${colors.bold(colors.red("ERROR"))}`);
+			for (const line of formatDiagnostic(e.message).split("\n"))
+				console.error(`  ${line}`);
+		},
+	});
+} catch (e) {
+	log.fail(e.message);
+	process.exit(1);
 }
 
-let debounce = null;
-let cssDebounce = null;
-watch(watchTarget, { recursive: true }, (_event, filename) => {
-	if (!filename) return;
-	if (filename.endsWith(".css")) {
-		clearTimeout(cssDebounce);
-		cssDebounce = setTimeout(() => handleCssWatch(filename), 50);
-		return;
-	}
-	if (!filename.endsWith(".go") && !filename.endsWith(".templ")) return;
-	clearTimeout(debounce);
-	debounce = setTimeout(() => buildOnce(filename), 80);
-});
-
-console.error(
-	`${stamp()} ${colors.bold("gofront:")} watching ${colors.cyan(inputArg)} ...`,
-);
+console.error(`${prefix()} watching ${colors.cyan(inputArg)} ...`);

package/src/lower/boxing.js

@@ -1,6 +1,8 @@
 // src/lower/boxing.js
 // Identifies address-taken scalar variables that require heap boxing ({ value: x }).
 
+import { forEachChild } from "./walk.js";
+
 export function isReferenceType(t) {
 	if (!t) return false;
 	const base = t.kind === "named" ? t.underlying : t;
@@ -28,9 +30,6 @@ export function scanAddressTaken(node, boxedVars = new Set()) {
 		}
 	}
 	// Recurse into FuncLit too — closures may take address of outer vars
-	for (const key of Object.keys(node)) {
-		if (key.startsWith("_")) continue;
-		scanAddressTaken(node[key], boxedVars);
-	}
+	forEachChild(node, (child) => scanAddressTaken(child, boxedVars));
 	return boxedVars;
 }

package/src/lower/captures.js

@@ -3,6 +3,7 @@
 // are mutated and thus require heap-allocated environment cells for WASM closures.
 
 import { collectMutatedVars, nodeMutatesVar } from "./ownership.js";
+import { forEachChild } from "./walk.js";
 
 /**
  * Result of capture analysis on a function or closure.
@@ -41,17 +42,14 @@ export function collectDeclaredNames(node, names = new Set()) {
 	// Do not descend into nested functions when collecting local declarations
 	if (node.kind === "FuncLit") return names;
 
-	for (const key of Object.keys(node)) {
-		if (key.startsWith("_")) continue;
-		collectDeclaredNames(node[key], names);
-	}
+	forEachChild(node, (child) => collectDeclaredNames(child, names));
 	return names;
 }
 
 /**
  * Finds all identifier references in a node that are not shadowed by inner declarations.
  */
-export function collectReferencedNames(node, localDecls, refs = new Set()) {
+function collectReferencedNames(node, localDecls, refs = new Set()) {
 	if (!node || typeof node !== "object") return refs;
 	if (Array.isArray(node)) {
 		for (const item of node) collectReferencedNames(item, localDecls, refs);
@@ -73,10 +71,9 @@ export function collectReferencedNames(node, localDecls, refs = new Set()) {
 			if (r.name && r.name !== "_") innerLocals.add(r.name);
 		}
 		collectDeclaredNames(node.body, innerLocals);
-		for (const key of Object.keys(node)) {
-			if (key.startsWith("_")) continue;
-			collectReferencedNames(node[key], innerLocals, refs);
-		}
+		forEachChild(node, (child) =>
+			collectReferencedNames(child, innerLocals, refs),
+		);
 		return refs;
 	}
 
@@ -89,10 +86,9 @@ export function collectReferencedNames(node, localDecls, refs = new Set()) {
 	}
 
 	// For nested closures, free variables are also referenced
-	for (const key of Object.keys(node)) {
-		if (key.startsWith("_")) continue;
-		collectReferencedNames(node[key], localDecls, refs);
-	}
+	forEachChild(node, (child) =>
+		collectReferencedNames(child, localDecls, refs),
+	);
 	return refs;
 }
 
@@ -168,10 +164,7 @@ export function analyzeCaptures(fnDecl) {
 			return;
 		}
 
-		for (const key of Object.keys(node)) {
-			if (key.startsWith("_")) continue;
-			walkClosures(node[key], scopeStack);
-		}
+		forEachChild(node, (child) => walkClosures(child, scopeStack));
 	};
 
 	walkClosures(fnDecl.body, [rootScope]);