Skip to content

Commit 654532c

Browse files
committed
fix(native): add prototype method extraction to Rust engine (#1327)
Implement parity with the WASM JS extractor for pre-ES6 prototype OOP patterns. Extractor (crates/codegraph-core/src/extractors/javascript.rs): - `Foo.prototype.bar = function(){}` → emits `Foo.bar` definition (kind: method) - `Foo.prototype.bar = identifier` → seeds typeMap['Foo.bar'] = identifier (confidence 0.9) - `Foo.prototype = { bar: fn, ... }` → same rules per property (pair, method_definition, shorthand_property_identifier) Built-in globals (Array, Object, …) are excluded via `is_js_builtin_global` guard. Adds 6 unit tests covering all three patterns plus edge cases. Edge builder (crates/codegraph-core/src/edge_builder.rs): - After a typeMap-resolved type lookup, check typeMap['TypeName.method'] for prototype aliases (`Foo.prototype.bar = identifierAlias`), mirroring the protoAlias fallback added to call-resolver.ts in the WASM path. - Inline new-expression receiver: extract class name from `(new Foo).bar()` receivers using string parsing (mirrors the `^\(?\s*new\s+[A-Z...]` regex in call-resolver.ts), enabling resolution without a named variable binding. Verified against the integration test in tests/integration/prototype-method-resolution.test.ts (all 3 tests pass with native engine). docs check acknowledged Closes #1327
1 parent 832f7fc commit 654532c

2 files changed

Lines changed: 261 additions & 1 deletion

File tree

crates/codegraph-core/src/edge_builder.rs

Lines changed: 49 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -407,13 +407,36 @@ fn resolve_call_targets<'a>(
407407
};
408408
let type_lookup = type_map.get(effective_receiver)
409409
.or_else(|| type_map.get(receiver.as_str()));
410-
if let Some(&(type_name, _conf)) = type_lookup {
410+
// Inline new-expression receiver: `(new Foo).bar()` — extract the constructor name
411+
// when no typeMap entry exists for the complex receiver expression.
412+
// Mirrors the regex `/^\(?\s*new\s+([A-Z_$][A-Za-z0-9_$]*)/` in call-resolver.ts.
413+
let inline_new_type = if type_lookup.is_none() {
414+
extract_inline_new_type(receiver)
415+
} else {
416+
None
417+
};
418+
// Use typeMap-resolved type or inline-new-extracted type, whichever is available.
419+
let resolved_type = type_lookup.map(|&(t, _)| t).or(inline_new_type.as_deref());
420+
if let Some(type_name) = resolved_type {
411421
let qualified = format!("{}.{}", type_name, call.name);
412422
let typed: Vec<&NodeInfo> = ctx.nodes_by_name
413423
.get(qualified.as_str())
414424
.map(|v| v.iter().filter(|n| n.kind == "method").copied().collect())
415425
.unwrap_or_default();
416426
if !typed.is_empty() { return typed; }
427+
// Prototype alias: `Foo.prototype.bar = identifier` seeds typeMap['Foo.bar'] = identifier.
428+
// After the direct method lookup misses (no definition emitted for this method),
429+
// check if the typeMap holds an alias to a standalone function.
430+
// Mirrors the protoAlias fallback in resolveByMethodOrGlobal in call-resolver.ts.
431+
if let Some(&(proto_target, _)) = type_map.get(qualified.as_str()) {
432+
let resolved: Vec<&NodeInfo> = ctx.nodes_by_name
433+
.get(proto_target)
434+
.map(|v| v.iter()
435+
.filter(|n| import_resolution::compute_confidence(rel_path, &n.file, None) >= 0.5)
436+
.copied().collect())
437+
.unwrap_or_default();
438+
if !resolved.is_empty() { return resolved; }
439+
}
417440
}
418441
// 4.5. Phase 8.3d: composite pts key — `obj.prop = fn` seeds typeMap['obj.prop']
419442
let composite_key = format!("{}.{}", receiver, call.name);
@@ -477,6 +500,31 @@ fn resolve_call_targets<'a>(
477500
Vec::new()
478501
}
479502

503+
/// Extract the constructor name from an inline `new` receiver expression.
504+
///
505+
/// Mirrors the regex `/^\(?\s*new\s+([A-Z_$][A-Za-z0-9_$]*)/` used in call-resolver.ts.
506+
/// Handles `(new Foo)` and `(new Foo('arg'))` receivers that arise when the call site
507+
/// is `(new Foo).method()` without a named variable binding.
508+
///
509+
/// Only extracts PascalCase (uppercase-initial) names to avoid false positives on
510+
/// lowercase constructor calls (rare but present in legacy code).
511+
fn extract_inline_new_type(receiver: &str) -> Option<String> {
512+
let s = receiver.trim_start_matches('(').trim_start();
513+
let s = s.strip_prefix("new")?;
514+
if !s.starts_with(|c: char| c.is_whitespace()) { return None; }
515+
let s = s.trim_start();
516+
let end = s.find(|c: char| !c.is_alphanumeric() && c != '_' && c != '$')
517+
.unwrap_or(s.len());
518+
let name = &s[..end];
519+
if name.is_empty() { return None; }
520+
let first = name.chars().next()?;
521+
if first.is_uppercase() || first == '_' || first == '$' {
522+
Some(name.to_string())
523+
} else {
524+
None
525+
}
526+
}
527+
480528
/// Sort targets by confidence descending.
481529
fn sort_targets_by_confidence(targets: &mut Vec<&NodeInfo>, rel_path: &str, imported_from: Option<&str>) {
482530
if targets.len() > 1 {

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

Lines changed: 212 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@ impl SymbolExtractor for JsExtractor {
2929
walk_ast_nodes(&tree.root_node(), source, &mut symbols.ast_nodes);
3030
walk_tree(&tree.root_node(), source, &mut symbols, match_js_type_map);
3131
walk_tree(&tree.root_node(), source, &mut symbols, match_js_return_type_map);
32+
// Pre-ES6 prototype methods: `Foo.prototype.bar = fn` and `Foo.prototype = { bar: fn }`
33+
walk_tree(&tree.root_node(), source, &mut symbols, match_js_prototype_methods);
3234
// call_assignments runs after type_map is populated (needs receiver types)
3335
walk_tree(&tree.root_node(), source, &mut symbols, match_js_call_assignments);
3436
symbols
@@ -292,6 +294,143 @@ fn push_return_type_entry(symbols: &mut FileSymbols, fn_name: &str, type_name: &
292294
});
293295
}
294296

297+
// ── Prototype-method extraction ─────────────────────────────────────────────
298+
299+
/// Walk the AST collecting pre-ES6 prototype assignments.
300+
///
301+
/// Mirrors `extractPrototypeMethodsWalk` in `src/extractors/javascript.ts`.
302+
///
303+
/// Three patterns are handled:
304+
/// 1. `Foo.prototype.bar = function(){}` → emits `Foo.bar` as a method definition
305+
/// 2. `Foo.prototype.bar = identifier` → seeds `typeMap['Foo.bar'] = identifier`
306+
/// 3. `Foo.prototype = { bar: fn, ... }` → same rules applied per property
307+
fn match_js_prototype_methods(node: &Node, source: &[u8], symbols: &mut FileSymbols, _depth: usize) {
308+
if node.kind() != "expression_statement" { return; }
309+
let Some(expr) = node.child(0) else { return };
310+
if expr.kind() != "assignment_expression" { return; }
311+
let lhs = expr.child_by_field_name("left");
312+
let rhs = expr.child_by_field_name("right");
313+
if let (Some(lhs), Some(rhs)) = (lhs, rhs) {
314+
handle_js_prototype_assignment(&lhs, &rhs, source, symbols);
315+
}
316+
}
317+
318+
fn handle_js_prototype_assignment(lhs: &Node, rhs: &Node, source: &[u8], symbols: &mut FileSymbols) {
319+
if lhs.kind() != "member_expression" { return; }
320+
let Some(lhs_obj) = lhs.child_by_field_name("object") else { return };
321+
let Some(lhs_prop) = lhs.child_by_field_name("property") else { return };
322+
323+
// Pattern 1: `Foo.prototype.bar = rhs`
324+
// lhs.object is `Foo.prototype` (member_expression), lhs.property is `bar`
325+
if lhs_obj.kind() == "member_expression"
326+
&& matches!(lhs_prop.kind(), "property_identifier" | "identifier")
327+
{
328+
let proto_obj = lhs_obj.child_by_field_name("object");
329+
let proto_prop = lhs_obj.child_by_field_name("property");
330+
if let (Some(proto_obj), Some(proto_prop)) = (proto_obj, proto_prop) {
331+
if proto_obj.kind() == "identifier"
332+
&& node_text(&proto_prop, source) == "prototype"
333+
&& !is_js_builtin_global(node_text(&proto_obj, source))
334+
{
335+
emit_js_prototype_method(
336+
node_text(&proto_obj, source),
337+
node_text(&lhs_prop, source),
338+
rhs,
339+
source,
340+
symbols,
341+
);
342+
}
343+
}
344+
return;
345+
}
346+
347+
// Pattern 2: `Foo.prototype = { bar: fn, ... }`
348+
// lhs.object is `Foo` (identifier), lhs.property is `prototype`, rhs is object literal
349+
if lhs_obj.kind() == "identifier"
350+
&& node_text(&lhs_prop, source) == "prototype"
351+
&& !is_js_builtin_global(node_text(&lhs_obj, source))
352+
&& rhs.kind() == "object"
353+
{
354+
extract_js_prototype_object_literal(node_text(&lhs_obj, source), rhs, source, symbols);
355+
}
356+
}
357+
358+
/// Emit one prototype method definition or typeMap alias for `ClassName.methodName = rhs`.
359+
///
360+
/// Mirrors `emitPrototypeMethod` in `src/extractors/javascript.ts`.
361+
fn emit_js_prototype_method(class_name: &str, method_name: &str, rhs: &Node, source: &[u8], symbols: &mut FileSymbols) {
362+
let full_name = format!("{}.{}", class_name, method_name);
363+
match rhs.kind() {
364+
"function_expression" | "arrow_function" => {
365+
symbols.definitions.push(Definition {
366+
name: full_name,
367+
kind: "method".to_string(),
368+
line: start_line(rhs),
369+
end_line: Some(end_line(rhs)),
370+
decorators: None,
371+
complexity: None,
372+
cfg: None,
373+
children: None,
374+
});
375+
}
376+
"identifier" => {
377+
let rhs_name = node_text(rhs, source);
378+
if !is_js_builtin_global(rhs_name) {
379+
push_type_map_entry(symbols, full_name, rhs_name.to_string());
380+
}
381+
}
382+
_ => {}
383+
}
384+
}
385+
386+
/// Iterate over an object literal assigned to `Foo.prototype` and emit definitions/aliases.
387+
///
388+
/// Mirrors `extractPrototypeObjectLiteral` in `src/extractors/javascript.ts`.
389+
fn extract_js_prototype_object_literal(class_name: &str, obj_node: &Node, source: &[u8], symbols: &mut FileSymbols) {
390+
for i in 0..obj_node.child_count() {
391+
let Some(child) = obj_node.child(i) else { continue };
392+
match child.kind() {
393+
"method_definition" => {
394+
let Some(name_node) = child.child_by_field_name("name") else { continue };
395+
symbols.definitions.push(Definition {
396+
name: format!("{}.{}", class_name, node_text(&name_node, source)),
397+
kind: "method".to_string(),
398+
line: start_line(&child),
399+
end_line: Some(end_line(&child)),
400+
decorators: None,
401+
complexity: None,
402+
cfg: None,
403+
children: None,
404+
});
405+
}
406+
"shorthand_property_identifier" => {
407+
let prop_name = node_text(&child, source);
408+
if !is_js_builtin_global(prop_name) {
409+
push_type_map_entry(
410+
symbols,
411+
format!("{}.{}", class_name, prop_name),
412+
prop_name.to_string(),
413+
);
414+
}
415+
}
416+
"pair" => {
417+
let key_node = child.child_by_field_name("key");
418+
let value_node = child.child_by_field_name("value");
419+
if let (Some(key_node), Some(value_node)) = (key_node, value_node) {
420+
let method_name = if key_node.kind() == "string" {
421+
node_text(&key_node, source).replace(['\'', '"'], "")
422+
} else {
423+
node_text(&key_node, source).to_string()
424+
};
425+
if method_name.is_empty() { continue; }
426+
emit_js_prototype_method(class_name, &method_name, &value_node, source, symbols);
427+
}
428+
}
429+
_ => {}
430+
}
431+
}
432+
}
433+
295434
// ── Call-assignment extraction (Phase 8.2 parity) ───────────────────────────
296435

297436
/// Walk the AST recording variable assignments from call expressions into
@@ -2292,4 +2431,77 @@ mod tests {
22922431
"compute call should have receiver='calc'"
22932432
);
22942433
}
2434+
2435+
// ── Prototype-method extraction ─────────────────────────────────────────
2436+
2437+
#[test]
2438+
fn prototype_direct_method_emits_definition() {
2439+
let s = parse_js(
2440+
"function C() {}\n\
2441+
C.prototype.foo = function() { return 1; };",
2442+
);
2443+
let def = s.definitions.iter().find(|d| d.name == "C.foo");
2444+
assert!(def.is_some(), "C.foo definition missing; got: {:?}", s.definitions.iter().map(|d| &d.name).collect::<Vec<_>>());
2445+
assert_eq!(def.unwrap().kind, "method");
2446+
}
2447+
2448+
#[test]
2449+
fn prototype_identifier_alias_seeds_type_map() {
2450+
let s = parse_js(
2451+
"let f = () => {};\n\
2452+
class A {}\n\
2453+
A.prototype.t = f;",
2454+
);
2455+
let entry = s.type_map.iter().find(|e| e.name == "A.t");
2456+
assert!(entry.is_some(), "type_map entry A.t missing; got: {:?}", s.type_map.iter().map(|e| &e.name).collect::<Vec<_>>());
2457+
assert_eq!(entry.unwrap().type_name, "f");
2458+
}
2459+
2460+
#[test]
2461+
fn prototype_object_literal_emits_definitions() {
2462+
let s = parse_js(
2463+
"function C() {}\n\
2464+
C.prototype = {\n\
2465+
foo: function() {},\n\
2466+
bar: function() {},\n\
2467+
};",
2468+
);
2469+
let foo = s.definitions.iter().find(|d| d.name == "C.foo");
2470+
let bar = s.definitions.iter().find(|d| d.name == "C.bar");
2471+
assert!(foo.is_some(), "C.foo missing");
2472+
assert_eq!(foo.unwrap().kind, "method");
2473+
assert!(bar.is_some(), "C.bar missing");
2474+
}
2475+
2476+
#[test]
2477+
fn prototype_object_literal_shorthand_method() {
2478+
let s = parse_js(
2479+
"function C() {}\n\
2480+
C.prototype = {\n\
2481+
greet() { return 'hi'; },\n\
2482+
};",
2483+
);
2484+
let def = s.definitions.iter().find(|d| d.name == "C.greet");
2485+
assert!(def.is_some(), "C.greet definition missing; got: {:?}", s.definitions.iter().map(|d| &d.name).collect::<Vec<_>>());
2486+
assert_eq!(def.unwrap().kind, "method");
2487+
}
2488+
2489+
#[test]
2490+
fn prototype_object_literal_shorthand_property_seeds_type_map() {
2491+
let s = parse_js(
2492+
"function helper() {}\n\
2493+
function C() {}\n\
2494+
C.prototype = { helper };",
2495+
);
2496+
let entry = s.type_map.iter().find(|e| e.name == "C.helper");
2497+
assert!(entry.is_some(), "type_map entry C.helper missing; got: {:?}", s.type_map.iter().map(|e| &e.name).collect::<Vec<_>>());
2498+
assert_eq!(entry.unwrap().type_name, "helper");
2499+
}
2500+
2501+
#[test]
2502+
fn prototype_builtin_globals_are_excluded() {
2503+
let s = parse_js("Array.prototype.custom = function() {};");
2504+
let def = s.definitions.iter().find(|d| d.name.contains("Array"));
2505+
assert!(def.is_none(), "built-in prototype assignment should be ignored; got: {:?}", def);
2506+
}
22952507
}

0 commit comments

Comments
 (0)