Skip to content

Commit c62e3b8

Browse files
authored
feat(resolver): extract generator functions as definitions (JS/TS) (#1333)
* feat(resolver): extract generator functions as definitions (JS/TS) Register function* declarations and const gen = function*() {} as first-class function definitions in both the WASM and native engines. Before this change, generator_function_declaration nodes were absent from the definition registry, so: - Calls inside generator bodies were attributed to the file level - yield* gen8() and other inter-generator calls produced no edges - The generators jelly-micro fixture scored 0% recall (0/9 named edges) Changes: - parser.ts + wasm-worker-entry.ts: add generator_function_declaration and generator_function (var-declared) to COMMON_QUERY_PATTERNS - extractors/javascript.ts: walk-path case, handleVariableDecl, extractConstDeclarators skip guard, export kindMaps, return type map - extractors/javascript.rs: equivalent changes for native engine parity - Add generators jelly-micro fixture (9 named edges, 100% recall) - Add 4 unit tests for generator extraction and yield* call capture Closes #1319 docs check acknowledged * test(parsers): strengthen generator call attribution assertion with line-range check
1 parent 2c63d85 commit c62e3b8

7 files changed

Lines changed: 198 additions & 10 deletions

File tree

crates/codegraph-core/src/extractors/javascript.rs

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -353,7 +353,7 @@ fn find_descriptor_value<'a>(node: &Node<'a>, source: &'a [u8]) -> Option<&'a st
353353
/// Mirrors `extractReturnTypeMapWalk` in src/extractors/javascript.ts.
354354
fn match_js_return_type_map(node: &Node, source: &[u8], symbols: &mut FileSymbols, _depth: usize) {
355355
match node.kind() {
356-
"function_declaration" => {
356+
"function_declaration" | "generator_function_declaration" => {
357357
let Some(name_n) = node.child_by_field_name("name") else { return };
358358
let fn_name = node_text(&name_n, source);
359359
if fn_name == "constructor" { return; }
@@ -381,9 +381,9 @@ fn match_js_return_type_map(node: &Node, source: &[u8], symbols: &mut FileSymbol
381381
let Some(name_n) = node.child_by_field_name("name") else { return };
382382
if name_n.kind() != "identifier" { return; }
383383
let Some(value_n) = node.child_by_field_name("value") else { return };
384-
// Only arrow_function and function_expression match the TS reference;
384+
// Only arrow_function, function_expression and generator_function match the TS reference;
385385
// "function" is not a valid tree-sitter value-expression kind here.
386-
if !matches!(value_n.kind(), "arrow_function" | "function_expression") {
386+
if !matches!(value_n.kind(), "arrow_function" | "function_expression" | "generator_function") {
387387
return;
388388
}
389389
let var_name = node_text(&name_n, source);
@@ -487,7 +487,7 @@ fn match_js_call_assignments(node: &Node, source: &[u8], symbols: &mut FileSymbo
487487

488488
fn match_js_node(node: &Node, source: &[u8], symbols: &mut FileSymbols, _depth: usize) {
489489
match node.kind() {
490-
"function_declaration" => handle_function_decl(node, source, symbols),
490+
"function_declaration" | "generator_function_declaration" => handle_function_decl(node, source, symbols),
491491
"class_declaration" | "abstract_class_declaration" => {
492492
handle_class_decl(node, source, symbols)
493493
}
@@ -650,7 +650,7 @@ fn handle_var_decl(node: &Node, source: &[u8], symbols: &mut FileSymbols) {
650650
let value_n = declarator.child_by_field_name("value");
651651
let (Some(name_n), Some(value_n)) = (name_n, value_n) else { continue };
652652
let vt = value_n.kind();
653-
if vt == "arrow_function" || vt == "function_expression" || vt == "function" {
653+
if vt == "arrow_function" || vt == "function_expression" || vt == "function" || vt == "generator_function" {
654654
let children = extract_js_parameters(&value_n, source);
655655
symbols.definitions.push(Definition {
656656
name: node_text(&name_n, source).to_string(),
@@ -804,7 +804,7 @@ fn handle_export_stmt(node: &Node, source: &[u8], symbols: &mut FileSymbols) {
804804

805805
fn handle_export_declaration(node: &Node, decl: &Node, source: &[u8], symbols: &mut FileSymbols) {
806806
let (kind_str, field) = match decl.kind() {
807-
"function_declaration" => ("function", "name"),
807+
"function_declaration" | "generator_function_declaration" => ("function", "name"),
808808
"class_declaration" | "abstract_class_declaration" => ("class", "name"),
809809
"interface_declaration" => ("interface", "name"),
810810
"type_alias_declaration" => ("type", "name"),

src/domain/parser.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -152,8 +152,10 @@ interface WasmExtractResult {
152152
// Shared patterns for all JS/TS/TSX (class_declaration excluded — name type differs)
153153
const COMMON_QUERY_PATTERNS: string[] = [
154154
'(function_declaration name: (identifier) @fn_name) @fn_node',
155+
'(generator_function_declaration name: (identifier) @fn_name) @fn_node',
155156
'(variable_declarator name: (identifier) @varfn_name value: (arrow_function) @varfn_value)',
156157
'(variable_declarator name: (identifier) @varfn_name value: (function_expression) @varfn_value)',
158+
'(variable_declarator name: (identifier) @varfn_name value: (generator_function) @varfn_value)',
157159
'(method_definition name: (property_identifier) @meth_name) @meth_node',
158160
'(method_definition name: (private_property_identifier) @meth_name) @meth_node',
159161
'(import_statement source: (string) @imp_source) @imp_node',

src/domain/wasm-worker-entry.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,8 +109,10 @@ function grammarPath(name: string): string {
109109

110110
const COMMON_QUERY_PATTERNS: string[] = [
111111
'(function_declaration name: (identifier) @fn_name) @fn_node',
112+
'(generator_function_declaration name: (identifier) @fn_name) @fn_node',
112113
'(variable_declarator name: (identifier) @varfn_name value: (arrow_function) @varfn_value)',
113114
'(variable_declarator name: (identifier) @varfn_name value: (function_expression) @varfn_value)',
115+
'(variable_declarator name: (identifier) @varfn_name value: (generator_function) @varfn_value)',
114116
'(method_definition name: (property_identifier) @meth_name) @meth_node',
115117
'(method_definition name: (private_property_identifier) @meth_name) @meth_node',
116118
'(import_statement source: (string) @imp_source) @imp_node',

src/extractors/javascript.ts

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -189,6 +189,7 @@ function handleExportCapture(
189189
const declType = decl.type;
190190
const kindMap: Record<string, string> = {
191191
function_declaration: 'function',
192+
generator_function_declaration: 'function',
192193
class_declaration: 'class',
193194
abstract_class_declaration: 'class',
194195
interface_declaration: 'interface',
@@ -482,7 +483,12 @@ function extractConstDeclarators(declNode: TreeSitterNode, definitions: Definiti
482483
if (nameN?.type !== 'identifier' || !valueN) continue;
483484
// Skip functions — already captured by query patterns
484485
const valType = valueN.type;
485-
if (valType === 'arrow_function' || valType === 'function_expression' || valType === 'function')
486+
if (
487+
valType === 'arrow_function' ||
488+
valType === 'function_expression' ||
489+
valType === 'function' ||
490+
valType === 'generator_function'
491+
)
486492
continue;
487493
if (isConstantValue(valueN)) {
488494
definitions.push({
@@ -629,6 +635,7 @@ function extractSymbolsWalk(tree: TreeSitterTree): ExtractorOutput {
629635
function walkJavaScriptNode(node: TreeSitterNode, ctx: ExtractorOutput): void {
630636
switch (node.type) {
631637
case 'function_declaration':
638+
case 'generator_function_declaration':
632639
handleFunctionDecl(node, ctx);
633640
break;
634641
case 'class_declaration':
@@ -809,7 +816,8 @@ function handleVariableDecl(node: TreeSitterNode, ctx: ExtractorOutput): void {
809816
if (
810817
valType === 'arrow_function' ||
811818
valType === 'function_expression' ||
812-
valType === 'function'
819+
valType === 'function' ||
820+
valType === 'generator_function'
813821
) {
814822
const varFnChildren = extractParameters(valueN);
815823
ctx.definitions.push({
@@ -941,6 +949,7 @@ function handleExportStmt(node: TreeSitterNode, ctx: ExtractorOutput): void {
941949
const declType = decl.type;
942950
const kindMap: Record<string, string> = {
943951
function_declaration: 'function',
952+
generator_function_declaration: 'function',
944953
class_declaration: 'class',
945954
abstract_class_declaration: 'class',
946955
interface_declaration: 'interface',
@@ -1205,7 +1214,7 @@ function extractReturnTypeMapWalk(
12051214
return;
12061215
}
12071216

1208-
if (t === 'function_declaration') {
1217+
if (t === 'function_declaration' || t === 'generator_function_declaration') {
12091218
const nameNode = node.childForFieldName('name');
12101219
if (nameNode?.type === 'identifier' && nameNode.text !== 'constructor') {
12111220
const fnName = currentClass ? `${currentClass}.${nameNode.text}` : nameNode.text;
@@ -1234,7 +1243,11 @@ function extractReturnTypeMapWalk(
12341243
const valueN = node.childForFieldName('value');
12351244
if (nameN?.type === 'identifier' && valueN) {
12361245
const vt = valueN.type;
1237-
if (vt === 'arrow_function' || vt === 'function_expression') {
1246+
if (
1247+
vt === 'arrow_function' ||
1248+
vt === 'function_expression' ||
1249+
vt === 'generator_function'
1250+
) {
12381251
const fnName = currentClass ? `${currentClass}.${nameN.text}` : nameN.text;
12391252
storeReturnType(valueN, fnName, returnTypeMap);
12401253
}
Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
{
2+
"$schema": "../../../expected-edges.schema.json",
3+
"language": "javascript",
4+
"description": "Hand-annotated call edges for generator function resolution benchmark",
5+
"edges": [
6+
{
7+
"source": { "name": "gen2", "file": "generators.js" },
8+
"target": { "name": "gen1", "file": "generators.js" },
9+
"kind": "calls",
10+
"mode": "static",
11+
"notes": "yield* delegation inside generator — inner call_expression"
12+
},
13+
{
14+
"source": { "name": "gen3", "file": "generators.js" },
15+
"target": { "name": "gen9", "file": "generators.js" },
16+
"kind": "calls",
17+
"mode": "static",
18+
"notes": "direct call to another generator from within a generator"
19+
},
20+
{
21+
"source": { "name": "gen4", "file": "generators.js" },
22+
"target": { "name": "gen4helper", "file": "generators.js" },
23+
"kind": "calls",
24+
"mode": "static",
25+
"notes": "call to regular function from inside a generator"
26+
},
27+
{
28+
"source": { "name": "gen5", "file": "generators.js" },
29+
"target": { "name": "gen2", "file": "generators.js" },
30+
"kind": "calls",
31+
"mode": "static",
32+
"notes": "yield* delegation to another named generator"
33+
},
34+
{
35+
"source": { "name": "gen5", "file": "generators.js" },
36+
"target": { "name": "gen4", "file": "generators.js" },
37+
"kind": "calls",
38+
"mode": "static",
39+
"notes": "second yield* delegation from gen5"
40+
},
41+
{
42+
"source": { "name": "gen6", "file": "generators.js" },
43+
"target": { "name": "gen7", "file": "generators.js" },
44+
"kind": "calls",
45+
"mode": "static",
46+
"notes": "direct call to sibling generator"
47+
},
48+
{
49+
"source": { "name": "gen7", "file": "generators.js" },
50+
"target": { "name": "gen6", "file": "generators.js" },
51+
"kind": "calls",
52+
"mode": "static",
53+
"notes": "direct call to sibling generator (mutual recursion)"
54+
},
55+
{
56+
"source": { "name": "gen9", "file": "generators.js" },
57+
"target": { "name": "gen8", "file": "generators.js" },
58+
"kind": "calls",
59+
"mode": "static",
60+
"notes": "yield* delegation — inner call_expression resolves to gen8"
61+
},
62+
{
63+
"source": { "name": "gen10", "file": "generators.js" },
64+
"target": { "name": "gen8", "file": "generators.js" },
65+
"kind": "calls",
66+
"mode": "static",
67+
"notes": "call from variable-declared generator (const gen10 = function*(){})"
68+
}
69+
]
70+
}
Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
// Jelly micro-test: generators
2+
// Tests call resolution in/between generator functions.
3+
4+
function* gen1() {
5+
yield 42;
6+
}
7+
8+
function* gen2() {
9+
yield* gen1(); // yield* delegation → edge gen2 → gen1
10+
}
11+
12+
function* gen3() {
13+
const it = gen9(); // direct call → edge gen3 → gen9
14+
it.next();
15+
yield it;
16+
}
17+
18+
function gen4helper() {
19+
return 1;
20+
}
21+
22+
function* gen4() {
23+
yield gen4helper(); // call to regular function → edge gen4 → gen4helper
24+
}
25+
26+
function* gen5() {
27+
yield* gen2(); // yield* delegation → edge gen5 → gen2
28+
yield* gen4(); // yield* delegation → edge gen5 → gen4
29+
}
30+
31+
function* gen6() {
32+
yield gen7(); // call to sibling generator → edge gen6 → gen7
33+
}
34+
35+
function* gen7() {
36+
yield gen6(); // call to sibling generator → edge gen7 → gen6
37+
}
38+
39+
function* gen8() {
40+
yield 1;
41+
yield 2;
42+
}
43+
44+
function* gen9() {
45+
yield* gen8(); // yield* delegation → edge gen9 → gen8
46+
}
47+
48+
// Variable-declared generator
49+
const gen10 = function* () {
50+
yield gen8(); // call from var-declared generator → edge gen10 → gen8
51+
};
52+
53+
// Entry: call some generators
54+
gen3();
55+
gen5();
56+
gen10();

tests/parsers/javascript.test.ts

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,51 @@ describe('JavaScript parser', () => {
3636
);
3737
});
3838

39+
it('extracts generator function declarations', () => {
40+
const symbols = parseJS(`function* gen() { yield 1; }`);
41+
expect(symbols.definitions).toContainEqual(
42+
expect.objectContaining({ name: 'gen', kind: 'function' }),
43+
);
44+
});
45+
46+
it('extracts variable-declared generator functions', () => {
47+
const symbols = parseJS(`const gen = function*() { yield 1; };`);
48+
expect(symbols.definitions).toContainEqual(
49+
expect.objectContaining({ name: 'gen', kind: 'function' }),
50+
);
51+
});
52+
53+
it('attributes calls inside generator body to the generator', () => {
54+
// Use multi-line generators so line ranges are non-overlapping and the
55+
// attribution can be verified by line number containment.
56+
const symbols = parseJS(
57+
'function* gen9() {\n yield* gen8();\n}\nfunction* gen8() { yield 1; }',
58+
);
59+
const gen9Def = symbols.definitions.find((d) => d.name === 'gen9');
60+
const gen8Def = symbols.definitions.find((d) => d.name === 'gen8');
61+
expect(gen9Def).toBeDefined();
62+
expect(gen8Def).toBeDefined();
63+
64+
// The call to gen8 must exist.
65+
const gen8Call = symbols.calls.find((c) => c.name === 'gen8');
66+
expect(gen8Call).toBeDefined();
67+
68+
// The call's line must fall within gen9's range — proving it is attributed
69+
// to gen9's body, not to file level or to gen8 itself.
70+
expect(gen8Call!.line).toBeGreaterThanOrEqual(gen9Def!.line);
71+
expect(gen8Call!.line).toBeLessThanOrEqual(gen9Def!.endLine!);
72+
73+
// Negative: the call must NOT fall within gen8's own range (not self-attributed).
74+
const callIsInsideGen8 =
75+
gen8Call!.line >= gen8Def!.line && gen8Call!.line <= (gen8Def!.endLine ?? gen8Def!.line);
76+
expect(callIsInsideGen8).toBe(false);
77+
});
78+
79+
it('captures calls inside yield* expressions', () => {
80+
const symbols = parseJS(`function* delegator() { yield* inner(); }`);
81+
expect(symbols.calls).toContainEqual(expect.objectContaining({ name: 'inner' }));
82+
});
83+
3984
it('extracts class declarations', () => {
4085
const symbols = parseJS(`class Foo { bar() {} }`);
4186
expect(symbols.definitions).toContainEqual(

0 commit comments

Comments
 (0)