diff --git a/crates/codegraph-core/src/extractors/javascript.rs b/crates/codegraph-core/src/extractors/javascript.rs index a4390f809..acd1006bb 100644 --- a/crates/codegraph-core/src/extractors/javascript.rs +++ b/crates/codegraph-core/src/extractors/javascript.rs @@ -353,7 +353,7 @@ fn find_descriptor_value<'a>(node: &Node<'a>, source: &'a [u8]) -> Option<&'a st /// Mirrors `extractReturnTypeMapWalk` in src/extractors/javascript.ts. fn match_js_return_type_map(node: &Node, source: &[u8], symbols: &mut FileSymbols, _depth: usize) { match node.kind() { - "function_declaration" => { + "function_declaration" | "generator_function_declaration" => { let Some(name_n) = node.child_by_field_name("name") else { return }; let fn_name = node_text(&name_n, source); if fn_name == "constructor" { return; } @@ -381,9 +381,9 @@ fn match_js_return_type_map(node: &Node, source: &[u8], symbols: &mut FileSymbol let Some(name_n) = node.child_by_field_name("name") else { return }; if name_n.kind() != "identifier" { return; } let Some(value_n) = node.child_by_field_name("value") else { return }; - // Only arrow_function and function_expression match the TS reference; + // Only arrow_function, function_expression and generator_function match the TS reference; // "function" is not a valid tree-sitter value-expression kind here. - if !matches!(value_n.kind(), "arrow_function" | "function_expression") { + if !matches!(value_n.kind(), "arrow_function" | "function_expression" | "generator_function") { return; } let var_name = node_text(&name_n, source); @@ -487,7 +487,7 @@ fn match_js_call_assignments(node: &Node, source: &[u8], symbols: &mut FileSymbo fn match_js_node(node: &Node, source: &[u8], symbols: &mut FileSymbols, _depth: usize) { match node.kind() { - "function_declaration" => handle_function_decl(node, source, symbols), + "function_declaration" | "generator_function_declaration" => handle_function_decl(node, source, symbols), "class_declaration" | "abstract_class_declaration" => { handle_class_decl(node, source, symbols) } @@ -650,7 +650,7 @@ fn handle_var_decl(node: &Node, source: &[u8], symbols: &mut FileSymbols) { let value_n = declarator.child_by_field_name("value"); let (Some(name_n), Some(value_n)) = (name_n, value_n) else { continue }; let vt = value_n.kind(); - if vt == "arrow_function" || vt == "function_expression" || vt == "function" { + if vt == "arrow_function" || vt == "function_expression" || vt == "function" || vt == "generator_function" { let children = extract_js_parameters(&value_n, source); symbols.definitions.push(Definition { name: node_text(&name_n, source).to_string(), @@ -804,7 +804,7 @@ fn handle_export_stmt(node: &Node, source: &[u8], symbols: &mut FileSymbols) { fn handle_export_declaration(node: &Node, decl: &Node, source: &[u8], symbols: &mut FileSymbols) { let (kind_str, field) = match decl.kind() { - "function_declaration" => ("function", "name"), + "function_declaration" | "generator_function_declaration" => ("function", "name"), "class_declaration" | "abstract_class_declaration" => ("class", "name"), "interface_declaration" => ("interface", "name"), "type_alias_declaration" => ("type", "name"), diff --git a/src/domain/parser.ts b/src/domain/parser.ts index b4703b1ca..663c26b20 100644 --- a/src/domain/parser.ts +++ b/src/domain/parser.ts @@ -152,8 +152,10 @@ interface WasmExtractResult { // Shared patterns for all JS/TS/TSX (class_declaration excluded — name type differs) const COMMON_QUERY_PATTERNS: string[] = [ '(function_declaration name: (identifier) @fn_name) @fn_node', + '(generator_function_declaration name: (identifier) @fn_name) @fn_node', '(variable_declarator name: (identifier) @varfn_name value: (arrow_function) @varfn_value)', '(variable_declarator name: (identifier) @varfn_name value: (function_expression) @varfn_value)', + '(variable_declarator name: (identifier) @varfn_name value: (generator_function) @varfn_value)', '(method_definition name: (property_identifier) @meth_name) @meth_node', '(method_definition name: (private_property_identifier) @meth_name) @meth_node', '(import_statement source: (string) @imp_source) @imp_node', diff --git a/src/domain/wasm-worker-entry.ts b/src/domain/wasm-worker-entry.ts index bac9421ef..d7e7ed442 100644 --- a/src/domain/wasm-worker-entry.ts +++ b/src/domain/wasm-worker-entry.ts @@ -109,8 +109,10 @@ function grammarPath(name: string): string { const COMMON_QUERY_PATTERNS: string[] = [ '(function_declaration name: (identifier) @fn_name) @fn_node', + '(generator_function_declaration name: (identifier) @fn_name) @fn_node', '(variable_declarator name: (identifier) @varfn_name value: (arrow_function) @varfn_value)', '(variable_declarator name: (identifier) @varfn_name value: (function_expression) @varfn_value)', + '(variable_declarator name: (identifier) @varfn_name value: (generator_function) @varfn_value)', '(method_definition name: (property_identifier) @meth_name) @meth_node', '(method_definition name: (private_property_identifier) @meth_name) @meth_node', '(import_statement source: (string) @imp_source) @imp_node', diff --git a/src/extractors/javascript.ts b/src/extractors/javascript.ts index f76f5f012..db9b91521 100644 --- a/src/extractors/javascript.ts +++ b/src/extractors/javascript.ts @@ -189,6 +189,7 @@ function handleExportCapture( const declType = decl.type; const kindMap: Record = { function_declaration: 'function', + generator_function_declaration: 'function', class_declaration: 'class', abstract_class_declaration: 'class', interface_declaration: 'interface', @@ -482,7 +483,12 @@ function extractConstDeclarators(declNode: TreeSitterNode, definitions: Definiti if (nameN?.type !== 'identifier' || !valueN) continue; // Skip functions — already captured by query patterns const valType = valueN.type; - if (valType === 'arrow_function' || valType === 'function_expression' || valType === 'function') + if ( + valType === 'arrow_function' || + valType === 'function_expression' || + valType === 'function' || + valType === 'generator_function' + ) continue; if (isConstantValue(valueN)) { definitions.push({ @@ -629,6 +635,7 @@ function extractSymbolsWalk(tree: TreeSitterTree): ExtractorOutput { function walkJavaScriptNode(node: TreeSitterNode, ctx: ExtractorOutput): void { switch (node.type) { case 'function_declaration': + case 'generator_function_declaration': handleFunctionDecl(node, ctx); break; case 'class_declaration': @@ -809,7 +816,8 @@ function handleVariableDecl(node: TreeSitterNode, ctx: ExtractorOutput): void { if ( valType === 'arrow_function' || valType === 'function_expression' || - valType === 'function' + valType === 'function' || + valType === 'generator_function' ) { const varFnChildren = extractParameters(valueN); ctx.definitions.push({ @@ -941,6 +949,7 @@ function handleExportStmt(node: TreeSitterNode, ctx: ExtractorOutput): void { const declType = decl.type; const kindMap: Record = { function_declaration: 'function', + generator_function_declaration: 'function', class_declaration: 'class', abstract_class_declaration: 'class', interface_declaration: 'interface', @@ -1205,7 +1214,7 @@ function extractReturnTypeMapWalk( return; } - if (t === 'function_declaration') { + if (t === 'function_declaration' || t === 'generator_function_declaration') { const nameNode = node.childForFieldName('name'); if (nameNode?.type === 'identifier' && nameNode.text !== 'constructor') { const fnName = currentClass ? `${currentClass}.${nameNode.text}` : nameNode.text; @@ -1234,7 +1243,11 @@ function extractReturnTypeMapWalk( const valueN = node.childForFieldName('value'); if (nameN?.type === 'identifier' && valueN) { const vt = valueN.type; - if (vt === 'arrow_function' || vt === 'function_expression') { + if ( + vt === 'arrow_function' || + vt === 'function_expression' || + vt === 'generator_function' + ) { const fnName = currentClass ? `${currentClass}.${nameN.text}` : nameN.text; storeReturnType(valueN, fnName, returnTypeMap); } diff --git a/tests/benchmarks/resolution/fixtures/jelly-micro/generators/expected-edges.json b/tests/benchmarks/resolution/fixtures/jelly-micro/generators/expected-edges.json new file mode 100644 index 000000000..7525ec153 --- /dev/null +++ b/tests/benchmarks/resolution/fixtures/jelly-micro/generators/expected-edges.json @@ -0,0 +1,70 @@ +{ + "$schema": "../../../expected-edges.schema.json", + "language": "javascript", + "description": "Hand-annotated call edges for generator function resolution benchmark", + "edges": [ + { + "source": { "name": "gen2", "file": "generators.js" }, + "target": { "name": "gen1", "file": "generators.js" }, + "kind": "calls", + "mode": "static", + "notes": "yield* delegation inside generator — inner call_expression" + }, + { + "source": { "name": "gen3", "file": "generators.js" }, + "target": { "name": "gen9", "file": "generators.js" }, + "kind": "calls", + "mode": "static", + "notes": "direct call to another generator from within a generator" + }, + { + "source": { "name": "gen4", "file": "generators.js" }, + "target": { "name": "gen4helper", "file": "generators.js" }, + "kind": "calls", + "mode": "static", + "notes": "call to regular function from inside a generator" + }, + { + "source": { "name": "gen5", "file": "generators.js" }, + "target": { "name": "gen2", "file": "generators.js" }, + "kind": "calls", + "mode": "static", + "notes": "yield* delegation to another named generator" + }, + { + "source": { "name": "gen5", "file": "generators.js" }, + "target": { "name": "gen4", "file": "generators.js" }, + "kind": "calls", + "mode": "static", + "notes": "second yield* delegation from gen5" + }, + { + "source": { "name": "gen6", "file": "generators.js" }, + "target": { "name": "gen7", "file": "generators.js" }, + "kind": "calls", + "mode": "static", + "notes": "direct call to sibling generator" + }, + { + "source": { "name": "gen7", "file": "generators.js" }, + "target": { "name": "gen6", "file": "generators.js" }, + "kind": "calls", + "mode": "static", + "notes": "direct call to sibling generator (mutual recursion)" + }, + { + "source": { "name": "gen9", "file": "generators.js" }, + "target": { "name": "gen8", "file": "generators.js" }, + "kind": "calls", + "mode": "static", + "notes": "yield* delegation — inner call_expression resolves to gen8" + }, + { + "source": { "name": "gen10", "file": "generators.js" }, + "target": { "name": "gen8", "file": "generators.js" }, + "kind": "calls", + "mode": "static", + "notes": "call from variable-declared generator (const gen10 = function*(){})" + } + ] +} diff --git a/tests/benchmarks/resolution/fixtures/jelly-micro/generators/generators.js b/tests/benchmarks/resolution/fixtures/jelly-micro/generators/generators.js new file mode 100644 index 000000000..54e07235d --- /dev/null +++ b/tests/benchmarks/resolution/fixtures/jelly-micro/generators/generators.js @@ -0,0 +1,56 @@ +// Jelly micro-test: generators +// Tests call resolution in/between generator functions. + +function* gen1() { + yield 42; +} + +function* gen2() { + yield* gen1(); // yield* delegation → edge gen2 → gen1 +} + +function* gen3() { + const it = gen9(); // direct call → edge gen3 → gen9 + it.next(); + yield it; +} + +function gen4helper() { + return 1; +} + +function* gen4() { + yield gen4helper(); // call to regular function → edge gen4 → gen4helper +} + +function* gen5() { + yield* gen2(); // yield* delegation → edge gen5 → gen2 + yield* gen4(); // yield* delegation → edge gen5 → gen4 +} + +function* gen6() { + yield gen7(); // call to sibling generator → edge gen6 → gen7 +} + +function* gen7() { + yield gen6(); // call to sibling generator → edge gen7 → gen6 +} + +function* gen8() { + yield 1; + yield 2; +} + +function* gen9() { + yield* gen8(); // yield* delegation → edge gen9 → gen8 +} + +// Variable-declared generator +const gen10 = function* () { + yield gen8(); // call from var-declared generator → edge gen10 → gen8 +}; + +// Entry: call some generators +gen3(); +gen5(); +gen10(); diff --git a/tests/parsers/javascript.test.ts b/tests/parsers/javascript.test.ts index e1d1a5e6e..1d3409e3d 100644 --- a/tests/parsers/javascript.test.ts +++ b/tests/parsers/javascript.test.ts @@ -36,6 +36,51 @@ describe('JavaScript parser', () => { ); }); + it('extracts generator function declarations', () => { + const symbols = parseJS(`function* gen() { yield 1; }`); + expect(symbols.definitions).toContainEqual( + expect.objectContaining({ name: 'gen', kind: 'function' }), + ); + }); + + it('extracts variable-declared generator functions', () => { + const symbols = parseJS(`const gen = function*() { yield 1; };`); + expect(symbols.definitions).toContainEqual( + expect.objectContaining({ name: 'gen', kind: 'function' }), + ); + }); + + it('attributes calls inside generator body to the generator', () => { + // Use multi-line generators so line ranges are non-overlapping and the + // attribution can be verified by line number containment. + const symbols = parseJS( + 'function* gen9() {\n yield* gen8();\n}\nfunction* gen8() { yield 1; }', + ); + const gen9Def = symbols.definitions.find((d) => d.name === 'gen9'); + const gen8Def = symbols.definitions.find((d) => d.name === 'gen8'); + expect(gen9Def).toBeDefined(); + expect(gen8Def).toBeDefined(); + + // The call to gen8 must exist. + const gen8Call = symbols.calls.find((c) => c.name === 'gen8'); + expect(gen8Call).toBeDefined(); + + // The call's line must fall within gen9's range — proving it is attributed + // to gen9's body, not to file level or to gen8 itself. + expect(gen8Call!.line).toBeGreaterThanOrEqual(gen9Def!.line); + expect(gen8Call!.line).toBeLessThanOrEqual(gen9Def!.endLine!); + + // Negative: the call must NOT fall within gen8's own range (not self-attributed). + const callIsInsideGen8 = + gen8Call!.line >= gen8Def!.line && gen8Call!.line <= (gen8Def!.endLine ?? gen8Def!.line); + expect(callIsInsideGen8).toBe(false); + }); + + it('captures calls inside yield* expressions', () => { + const symbols = parseJS(`function* delegator() { yield* inner(); }`); + expect(symbols.calls).toContainEqual(expect.objectContaining({ name: 'inner' })); + }); + it('extracts class declarations', () => { const symbols = parseJS(`class Foo { bar() {} }`); expect(symbols.definitions).toContainEqual(