From cb07a93ae996c1ab038da311328187d9ab4f4ec1 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 13 Apr 2026 16:09:28 +0800 Subject: [PATCH 1/2] fix(ir): walk anonymous function args of call declarations Surfaces handler methods hidden inside factory patterns like const X = call(() => { return { a, b } }). This covers Effect.gen, Rpc.toLayer(Effect.gen(...)).pipe(Layer.provide(...)), and similar config-object DSLs across the ecosystem. Also fixes multi-line import source extraction: the previous regex-on-truncated-text approach lost the module name for imports that exceeded 80 chars. Now extracts the source directly from the AST. Tests: 155 existing pass; 8 new in tests/ir/factory-declarations.test.ts. --- src/ir/ast-walker.ts | 188 +++++++++++++++++++++++++- tests/ir/factory-declarations.test.ts | 109 +++++++++++++++ 2 files changed, 291 insertions(+), 6 deletions(-) create mode 100644 tests/ir/factory-declarations.test.ts diff --git a/src/ir/ast-walker.ts b/src/ir/ast-walker.ts index 0d55351..c95a68f 100644 --- a/src/ir/ast-walker.ts +++ b/src/ir/ast-walker.ts @@ -281,14 +281,172 @@ function emitTier2(node: SyntaxNode): string | null { } } +// Extract the module path from an import_statement via AST, not text regex. +// Handles multi-line named imports where text truncation would drop the source. +function findImportSource(node: SyntaxNode): string | null { + for (let i = 0; i < node.childCount; i++) { + const c = node.child(i)!; + if (c.type === "string") { + for (let j = 0; j < c.childCount; j++) { + const cc = c.child(j)!; + if (cc.type === "string_fragment") return cc.text; + } + return c.text.replace(/^['"]|['"]$/g, ""); + } + } + return null; +} + +function argumentsNode(call: SyntaxNode): SyntaxNode | null { + const byField = call.childForFieldName("arguments"); + if (byField) return byField; + for (let i = 0; i < call.childCount; i++) { + const c = call.child(i)!; + if (c.type === "arguments") return c; + } + return null; +} + +const ANON_FN_TYPES = new Set([ + "arrow_function", + "function_expression", + "function", + "generator_function", + "generator_function_declaration", +]); + +// Collect anonymous-function arguments from a call expression, walking through +// nested calls, member-expression receivers (e.g. `.pipe(...)` chains), and +// parenthesized expressions. Covers Effect.gen, Rpc.toLayer(...).pipe(...), +// createSlice, defineConfig, chained builders, etc. +function collectAnonFnArgs(call: SyntaxNode): SyntaxNode[] { + const out: SyntaxNode[] = []; + const seen = new Set(); + const visit = (n: SyntaxNode | null, budget: number) => { + if (!n || budget < 0 || seen.has(n.id)) return; + seen.add(n.id); + if (ANON_FN_TYPES.has(n.type)) { + out.push(n); + return; + } + if (n.type === "parenthesized_expression") { + for (let i = 0; i < n.childCount; i++) visit(n.child(i)!, budget); + return; + } + if (n.type === "call_expression") { + // Descend into arguments. + const args = argumentsNode(n); + if (args) { + for (let i = 0; i < args.childCount; i++) visit(args.child(i)!, budget - 1); + } + // Also descend into the callee receiver to handle `X.pipe(...)` wrappers + // where the factory call `X` is what holds the anonymous function. + const fn = n.childForFieldName("function"); + if (fn && fn.type === "member_expression") { + const obj = fn.childForFieldName("object"); + if (obj) visit(obj, budget - 1); + } + } + }; + visit(call, 4); + return out; +} + +function fnBody(fn: SyntaxNode): SyntaxNode | null { + const byField = fn.childForFieldName("body"); + if (byField) return byField; + for (let i = 0; i < fn.childCount; i++) { + const c = fn.child(i)!; + if (c.type === "statement_block" || c.type === "block") return c; + } + return null; +} + +// Walk the body of an anonymous function looking for `return { key: (...) => ... }` +// patterns and emit each arrow-function pair as a METHOD. This surfaces the +// handler-object idiom used by Effect Rpc (`Rpc.toLayer(Effect.gen(function* () { +// return { handler1, handler2 } }))`), Redux Toolkit `createSlice({ reducers })`, +// and similar config-object DSLs. +function walkFactoryBody(fn: SyntaxNode, depth: number, lines: string[]): void { + const body = fnBody(fn); + if (!body) return; + const visit = (n: SyntaxNode, d: number) => { + if (d > 6) return; + if (n.type === "return_statement") { + for (let i = 0; i < n.childCount; i++) { + const c = n.child(i)!; + if (c.type === "object") { + emitObjectMethods(c, depth, lines); + } + } + return; + } + // Recurse into nested blocks so the return can be anywhere in the function. + for (let i = 0; i < n.childCount; i++) visit(n.child(i)!, d + 1); + }; + for (let i = 0; i < body.childCount; i++) visit(body.child(i)!, 0); + + // Also walk nested T1 declarations inside the factory body so helpers + // defined alongside the returned object still appear in IR. Skip + // return_statements: their handler-object contents are already surfaced + // as METHOD lines above, and re-walking would emit a noisy RET duplicate. + for (let i = 0; i < body.childCount; i++) { + const child = body.child(i)!; + if (child.type === "return_statement") continue; + walkNode(child, depth + 1, lines); + } +} + +function emitObjectMethods(obj: SyntaxNode, depth: number, lines: string[]): void { + const indent = " ".repeat(depth); + for (let i = 0; i < obj.childCount; i++) { + const child = obj.child(i)!; + if (child.type !== "pair") continue; + const key = child.childForFieldName("key"); + const value = child.childForFieldName("value"); + if (!key || !value) continue; + const keyName = key.text; + if (keyName.startsWith("_") || keyName.startsWith("#")) continue; + if (ANON_FN_TYPES.has(value.type)) { + const params = value.childForFieldName("parameters")?.text ?? "()"; + const asyncPrefix = isAsync(value) ? "ASYNC " : ""; + lines.push(`${indent}${asyncPrefix}METHOD:${keyName}${collapseText(params, 40)}`); + } + } +} + +// True when a lexical_declaration's RHS is a call expression whose arguments +// include an anonymous function (directly or one nested call level deep). +function isFactoryDeclaration(node: SyntaxNode): { name: string; fns: SyntaxNode[] } | null { + let declarator: SyntaxNode | null = null; + for (let i = 0; i < node.childCount; i++) { + const c = node.child(i)!; + if (c.type === "variable_declarator") { + declarator = c; + break; + } + } + if (!declarator) return null; + const name = declarator.childForFieldName("name")?.text; + const value = declarator.childForFieldName("value"); + if (!name || !value || value.type !== "call_expression") return null; + const fns = collectAnonFnArgs(value); + if (fns.length === 0) return null; + return { name, fns }; +} + function emitTier1(node: SyntaxNode): string | null { const exported = isExported(node); const outPrefix = exported ? "OUT " : ""; switch (node.type) { case "import_statement": { - const text = collapseText(node.text, 80); - return `USE:${text}`; + // Extract module source from AST (string > string_fragment) instead of + // regex on collapsed text. Collapsed text truncates long imports before + // the `from "..."` clause, breaking downstream module-name extraction. + const source = findImportSource(node); + if (source) return `USE:${source}`; + return `USE:${collapseText(node.text, 80)}`; } case "function_declaration": { @@ -577,12 +735,29 @@ function walkNode(node: SyntaxNode, depth: number, lines: string[]): void { } case "T3_COMPRESS": { - // At deep nesting, expression details are noise — skip + // At deep nesting, expression details are noise - skip if (depth > 4) break; + + // Factory-declaration pattern: `const X = call(fn)` where the call + // receives an anonymous function. Covers Effect.gen, Redux createSlice, + // TanStack createRouter, defineConfig, etc. Emit the declaration and + // walk the anonymous function body so its returned methods appear. + if (node.type === "lexical_declaration") { + const factory = isFactoryDeclaration(node); + if (factory) { + const exported = isExported(node) || node.parent?.type === "export_statement"; + const prefix = exported ? "OUT " : ""; + const indent = " ".repeat(depth); + lines.push(`${indent}${prefix}FN:${factory.name}`); + for (const fn of factory.fns) walkFactoryBody(fn, depth + 1, lines); + break; + } + } + const line = emitTier3(node); const indent = " ".repeat(depth); if (line) lines.push(indent + line); - // Don't walk children — one-liner is enough + // Don't walk children - one-liner is enough break; } @@ -627,8 +802,9 @@ export async function astWalkIR(code: string, filePath: string): Promise` strings via findImportSource, + // so no fallback regex is needed. + useBlock.push(line.slice(4)); } else { if (useBlock.length > 0) { if (useBlock.length <= 3) { diff --git a/tests/ir/factory-declarations.test.ts b/tests/ir/factory-declarations.test.ts new file mode 100644 index 0000000..0abb7e6 --- /dev/null +++ b/tests/ir/factory-declarations.test.ts @@ -0,0 +1,109 @@ +import { describe, it, expect } from "vitest"; +import { astWalkIR } from "../../src/ir/ast-walker.js"; + +describe("astWalkIR - factory declarations (call expressions with anonymous functions)", () => { + describe("import source extraction", () => { + it("preserves module source for long multi-line imports", async () => { + const code = `import { + AAAAAAAAAAAAA, BBBBBBBBBBBBB, CCCCCCCCCCCCC, + DDDDDDDDDDDDD, EEEEEEEEEEEEE, FFFFFFFFFFFFF, + type GGGGGGGGGGGGG, +} from "some-really-long-module-path/deep/nested";`; + const ir = await astWalkIR(code, "app.ts"); + expect(ir).toContain("some-really-long-module-path/deep/nested"); + expect(ir).not.toContain("import {"); + }); + + it("extracts source from single-quote imports", async () => { + const code = `import { a } from 'module-x';`; + const ir = await astWalkIR(code, "app.ts"); + expect(ir).toContain("USE:module-x"); + }); + }); + + describe("Effect.gen handler extraction", () => { + it("surfaces handler methods inside Effect.gen factories", async () => { + const code = ` +export const ChatRpcLive = ChatRpc.toLayer( + Effect.gen(function* () { + const repo = yield* ChatRepository; + return { + send: (input: SendInput) => Effect.gen(function* () { return yield* repo.save(input); }), + list: (input: ListInput) => repo.list(input), + }; + }) +);`; + const ir = await astWalkIR(code, "rpc-live.ts"); + expect(ir).toContain("OUT FN:ChatRpcLive"); + expect(ir).toContain("METHOD:send(input: SendInput)"); + expect(ir).toContain("METHOD:list(input: ListInput)"); + }); + + it("extracts handlers through .pipe(...) wrappers", async () => { + const code = ` +export const Live = Service.toLayer( + Effect.gen(function* () { + return { + foo: (x: number) => x * 2, + bar: (y: string) => y.toUpperCase(), + }; + }) +).pipe( + Layer.provide(Dep1.Default), + Layer.provide(Dep2.Default), +);`; + const ir = await astWalkIR(code, "live.ts"); + expect(ir).toContain("OUT FN:Live"); + expect(ir).toContain("METHOD:foo(x: number)"); + expect(ir).toContain("METHOD:bar(y: string)"); + }); + + it("marks async arrow-function methods", async () => { + const code = ` +export const X = make(Effect.gen(function* () { + return { fetch: async (id: string) => db.get(id) }; +}));`; + const ir = await astWalkIR(code, "x.ts"); + expect(ir).toContain("ASYNC METHOD:fetch(id: string)"); + }); + + it("skips private (_ and #) method keys", async () => { + const code = ` +export const X = make(function() { + return { + _internal: (x: number) => x, + public: (y: string) => y, + }; +});`; + const ir = await astWalkIR(code, "x.ts"); + expect(ir).toContain("METHOD:public(y: string)"); + expect(ir).not.toContain("_internal"); + }); + + it("does not emit METHOD lines for non-factory call declarations", async () => { + const code = ` +export const config = loadConfig({ env: "prod", debug: false }); +export function main() { return 1; }`; + const ir = await astWalkIR(code, "x.ts"); + expect(ir).toContain("FN:main"); + expect(ir).not.toContain("METHOD:"); + }); + }); + + describe("multiple handlers + nested declarations", () => { + it("emits helpers defined alongside the returned handler object", async () => { + const code = ` +export const Live = make(Effect.gen(function* () { + const repo = yield* Repo; + function helper(x: number) { return x + 1; } + return { + run: (input: string) => helper(input.length), + }; +}));`; + const ir = await astWalkIR(code, "live.ts"); + expect(ir).toContain("OUT FN:Live"); + expect(ir).toContain("METHOD:run(input: string)"); + expect(ir).toContain("FN:helper(x: number)"); + }); + }); +}); From d625e58d1eecd3531034a9a3862e5c2ccca81ae5 Mon Sep 17 00:00:00 2001 From: Neco Date: Tue, 14 Apr 2026 09:56:36 +0800 Subject: [PATCH 2/2] fix(ir): emit METHOD for shorthand and method_definition in factories Handles { send, list } and { send(x) {} } returned from factory bodies, common in Vite, TanStack, Zustand style APIs. Addresses review feedback on #1. --- src/ir/ast-walker.ts | 38 +++++++++++++++++------ tests/ir/factory-declarations.test.ts | 43 +++++++++++++++++++++++++++ 2 files changed, 72 insertions(+), 9 deletions(-) diff --git a/src/ir/ast-walker.ts b/src/ir/ast-walker.ts index c95a68f..61dcd63 100644 --- a/src/ir/ast-walker.ts +++ b/src/ir/ast-walker.ts @@ -401,16 +401,36 @@ function emitObjectMethods(obj: SyntaxNode, depth: number, lines: string[]): voi const indent = " ".repeat(depth); for (let i = 0; i < obj.childCount; i++) { const child = obj.child(i)!; - if (child.type !== "pair") continue; - const key = child.childForFieldName("key"); - const value = child.childForFieldName("value"); - if (!key || !value) continue; - const keyName = key.text; - if (keyName.startsWith("_") || keyName.startsWith("#")) continue; - if (ANON_FN_TYPES.has(value.type)) { - const params = value.childForFieldName("parameters")?.text ?? "()"; - const asyncPrefix = isAsync(value) ? "ASYNC " : ""; + if (child.type === "pair") { + const key = child.childForFieldName("key"); + const value = child.childForFieldName("value"); + if (!key || !value) continue; + const keyName = key.text; + if (keyName.startsWith("_") || keyName.startsWith("#")) continue; + if (ANON_FN_TYPES.has(value.type)) { + const params = value.childForFieldName("parameters")?.text ?? "()"; + const asyncPrefix = isAsync(value) ? "ASYNC " : ""; + lines.push(`${indent}${asyncPrefix}METHOD:${keyName}${collapseText(params, 40)}`); + } + continue; + } + // shorthand property: { send, list } + if (child.type === "shorthand_property_identifier") { + const keyName = child.text; + if (keyName.startsWith("_") || keyName.startsWith("#")) continue; + lines.push(`${indent}METHOD:${keyName}`); + continue; + } + // method definition: { send(x) {} } + if (child.type === "method_definition") { + const key = child.childForFieldName("name"); + if (!key) continue; + const keyName = key.text; + if (keyName.startsWith("_") || keyName.startsWith("#")) continue; + const params = child.childForFieldName("parameters")?.text ?? "()"; + const asyncPrefix = isAsync(child) ? "ASYNC " : ""; lines.push(`${indent}${asyncPrefix}METHOD:${keyName}${collapseText(params, 40)}`); + continue; } } } diff --git a/tests/ir/factory-declarations.test.ts b/tests/ir/factory-declarations.test.ts index 0abb7e6..3ba31f8 100644 --- a/tests/ir/factory-declarations.test.ts +++ b/tests/ir/factory-declarations.test.ts @@ -90,6 +90,49 @@ export function main() { return 1; }`; }); }); + describe("shorthand properties and method definitions", () => { + it("emits METHOD for shorthand properties (Vite/Zustand style)", async () => { + const code = ` +export const store = create(() => { + const send = (x: number) => x; + const list = () => []; + return { send, list }; +});`; + const ir = await astWalkIR(code, "store.ts"); + expect(ir).toContain("METHOD:send"); + expect(ir).toContain("METHOD:list"); + }); + + it("emits METHOD for method_definition shorthand (TanStack style)", async () => { + const code = ` +export const plugin = defineConfig(() => { + return { + send(input: SendInput) { return input; }, + async list(filter: string) { return []; }, + }; +});`; + const ir = await astWalkIR(code, "plugin.ts"); + expect(ir).toContain("METHOD:send(input: SendInput)"); + expect(ir).toContain("ASYNC METHOD:list(filter: string)"); + }); + + it("skips private keys in shorthand and method_definition", async () => { + const code = ` +export const x = make(() => { + const _internal = 1; + return { + _internal, + public(y: string) { return y; }, + _hidden(z: number) { return z; }, + }; +});`; + const ir = await astWalkIR(code, "x.ts"); + expect(ir).toContain("METHOD:public(y: string)"); + expect(ir).not.toContain("_internal"); + expect(ir).not.toContain("_hidden"); + }); + }); + describe("multiple handlers + nested declarations", () => { it("emits helpers defined alongside the returned handler object", async () => { const code = `