feat: Phase 3 — return type inference, generic args extraction, Ruby YARD type extractor

Three architectural improvements to the type resolution system:

1. Return type inference — wire extractMethodSignature returnType through
   SymbolDefinition into call-processor. When var = callee() and callee
   has a known return type, bind var to that type. Handles Promise<T>
   unwrapping, nullable stripping, pointer/reference removal.

2. Generic type argument extraction — new extractGenericTypeArgs() utility
   that extracts type parameters from List<User> → ['User']. Handles
   TS/Java/Kotlin/C#/Rust generic syntax. Building block for for-loop
   variable typing.

3. Ruby dedicated type extractor — replaces the stub with YARD annotation
   parsing (@param name [Type]), handling qualified types, nullable types,
   and singleton methods. Ruby now has real type resolution.

Unit tests: 127 → 192+ (type-env) + 65 (symbol-table, call-processor) + 18 (generics)
Integration tests: 8+ new test cases with fixtures across TS/Python/Go/Java/Ruby
This commit is contained in:
Gergo Magyar 2026-03-14 19:51:58 +00:00
parent 62242d5f44
commit f8685781d0
26 changed files with 1103 additions and 9 deletions

View file

@ -343,6 +343,85 @@ const resolveCallTarget = (
return toResolveResult(filteredCandidates[0], tiered.tier);
};
// ── Return type text helpers ─────────────────────────────────────────────
// extractSimpleTypeName works on AST nodes; this operates on raw return-type
// text already stored in SymbolDefinition (e.g. "User", "Promise<User>",
// "User | null", "*User"). Extracts the base user-defined type name.
/** Primitive / built-in types that should NOT produce a receiver binding. */
const PRIMITIVE_TYPES = new Set([
'string', 'number', 'boolean', 'void', 'int', 'float', 'double', 'long',
'short', 'byte', 'char', 'bool', 'str', 'i8', 'i16', 'i32', 'i64',
'u8', 'u16', 'u32', 'u64', 'f32', 'f64', 'usize', 'isize',
'undefined', 'null', 'None', 'nil',
]);
/**
* Extract a simple type name from raw return-type text.
* Handles common patterns:
* "User" → "User"
* "Promise<User>" → "User" (unwrap wrapper generics)
* "Option<User>" → "User"
* "Result<User, Error>" → "User" (first type arg)
* "User | null" → "User" (strip nullable union)
* "User?" → "User" (strip nullable suffix)
* "*User" → "User" (Go pointer)
* "&User" → "User" (Rust reference)
* Returns undefined for complex types or primitives.
*/
const WRAPPER_GENERICS = new Set([
'Promise', 'Observable', 'Option', 'Some', 'Result',
'Optional', 'Future', 'Task', 'ValueTask',
'List', 'Array', 'Vec', 'Set', 'Iterable',
'Sequence', 'MutableList', 'ArrayList',
]);
export const extractReturnTypeName = (raw: string): string | undefined => {
let text = raw.trim();
if (!text) return undefined;
// Strip pointer/reference prefixes: *User, &User, &mut User
text = text.replace(/^[&*]+\s*(mut\s+)?/, '');
// Strip nullable suffix: User?
text = text.replace(/\?$/, '');
// Handle union types: "User | null" → "User"
if (text.includes('|')) {
const parts = text.split('|').map(p => p.trim()).filter(p =>
p !== 'null' && p !== 'undefined' && p !== 'void' && p !== 'None' && p !== 'nil'
);
if (parts.length === 1) text = parts[0];
else return undefined; // genuine union — too complex
}
// Handle generics: Promise<User> → unwrap if wrapper, else take base
const genericMatch = text.match(/^(\w+)\s*<(.+)>$/);
if (genericMatch) {
const [, base, args] = genericMatch;
if (WRAPPER_GENERICS.has(base)) {
// Take the first type argument (e.g., Result<User, Error> → User)
const firstArg = args.split(',')[0].trim();
return extractReturnTypeName(firstArg);
}
// Non-wrapper generic: return the base type (e.g., Map<K,V> → Map)
return PRIMITIVE_TYPES.has(base.toLowerCase()) ? undefined : base;
}
// Handle qualified names: models.User → User, com.example.User → User
if (text.includes('.')) {
text = text.split('.').pop()!;
}
// Final check: skip primitives
if (PRIMITIVE_TYPES.has(text) || PRIMITIVE_TYPES.has(text.toLowerCase())) return undefined;
// Must start with uppercase (class/type convention) or be a valid identifier
if (!/^[A-Z_]\w*$/.test(text)) return undefined;
return text;
};
// ── Scope key helpers ────────────────────────────────────────────────────
// Scope keys use the format "funcName@startIndex" (produced by type-env.ts).
// Source IDs use "Label:filepath:funcName" (produced by parse-worker.ts).
@ -386,6 +465,19 @@ export const processCallsFromExtracted = async (
if (isClass) {
if (!fileReceiverTypes.has(filePath)) fileReceiverTypes.set(filePath, new Map());
fileReceiverTypes.get(filePath)!.set(receiverKey(extractFuncNameFromScope(scope), varName), calleeName);
} else {
// Return type inference: if the callee is a function/method with a known
// return type, bind the variable to that return type.
const callableDefs = tiered?.candidates.filter(d =>
d.type === 'Function' || d.type === 'Method'
);
if (callableDefs && callableDefs.length === 1 && callableDefs[0].returnType) {
const typeName = extractReturnTypeName(callableDefs[0].returnType);
if (typeName) {
if (!fileReceiverTypes.has(filePath)) fileReceiverTypes.set(filePath, new Map());
fileReceiverTypes.get(filePath)!.set(receiverKey(extractFuncNameFromScope(scope), varName), typeName);
}
}
}
}
}

View file

@ -79,6 +79,7 @@ const processParsingWithWorkers = async (
for (const sym of result.symbols) {
symbolTable.add(sym.filePath, sym.name, sym.nodeId, sym.type, {
parameterCount: sym.parameterCount,
returnType: sym.returnType,
ownerId: sym.ownerId,
});
}
@ -244,6 +245,7 @@ const processParsingSequential = async (
symbolTable.add(file.path, nodeName, nodeId, nodeLabel, {
parameterCount: methodSig?.parameterCount,
returnType: methodSig?.returnType,
ownerId: enclosingClassId ?? undefined,
});

View file

@ -3,6 +3,8 @@ export interface SymbolDefinition {
filePath: string;
type: string; // 'Function', 'Class', etc.
parameterCount?: number;
/** Raw return type text extracted from AST (e.g. 'User', 'Promise<User>') */
returnType?: string;
/** Links Method/Constructor to owning Class/Struct/Trait nodeId */
ownerId?: string;
}
@ -16,7 +18,7 @@ export interface SymbolTable {
name: string,
nodeId: string,
type: string,
metadata?: { parameterCount?: number; ownerId?: string }
metadata?: { parameterCount?: number; returnType?: string; ownerId?: string }
) => void;
/**
@ -62,13 +64,14 @@ export const createSymbolTable = (): SymbolTable => {
name: string,
nodeId: string,
type: string,
metadata?: { parameterCount?: number; ownerId?: string }
metadata?: { parameterCount?: number; returnType?: string; ownerId?: string }
) => {
const def: SymbolDefinition = {
nodeId,
filePath,
type,
...(metadata?.parameterCount !== undefined ? { parameterCount: metadata.parameterCount } : {}),
...(metadata?.returnType !== undefined ? { returnType: metadata.returnType } : {}),
...(metadata?.ownerId !== undefined ? { ownerId: metadata.ownerId } : {}),
};

View file

@ -15,6 +15,7 @@ import { typeConfig as pythonConfig } from './python.js';
import { typeConfig as swiftConfig } from './swift.js';
import { typeConfig as cCppConfig } from './c-cpp.js';
import { typeConfig as phpConfig } from './php.js';
import { typeConfig as rubyConfig } from './ruby.js';
export const typeConfigs = {
[SupportedLanguages.JavaScript]: typescriptConfig,
@ -29,12 +30,8 @@ export const typeConfigs = {
[SupportedLanguages.C]: cCppConfig,
[SupportedLanguages.CPlusPlus]: cCppConfig,
[SupportedLanguages.PHP]: phpConfig,
[SupportedLanguages.Ruby]: {
declarationNodeTypes: new Set<string>(),
extractDeclaration: () => {},
extractParameter: () => {},
} as LanguageTypeConfig,
[SupportedLanguages.Ruby]: rubyConfig,
} satisfies Record<SupportedLanguages, LanguageTypeConfig>;
export type { LanguageTypeConfig, TypeBindingExtractor, ParameterExtractor } from './types.js';
export { TYPED_PARAMETER_TYPES, extractSimpleTypeName, extractVarName, findChildByType } from './shared.js';
export { TYPED_PARAMETER_TYPES, extractSimpleTypeName, extractGenericTypeArgs, extractVarName, findChildByType } from './shared.js';

View file

@ -0,0 +1,205 @@
import type { SyntaxNode } from '../utils.js';
import type { LanguageTypeConfig, ParameterExtractor, TypeBindingExtractor, InitializerExtractor, ClassNameLookup } from './types.js';
import { extractVarName } from './shared.js';
/**
* Ruby type extractor — YARD annotation parsing.
*
* Ruby has no static type system, but the YARD documentation convention
* provides de facto type annotations via comments:
*
* # @param name [String] the user's name
* # @param repo [UserRepo] the repository
* # @return [User]
* def create(name, repo)
* repo.save
* end
*
* This extractor parses `@param name [Type]` patterns from comment nodes
* preceding method definitions and binds parameter names to their types.
*
* Resolution tiers:
* - Tier 0: YARD @param annotations (extractDeclaration pre-populates env)
* - Tier 1: Constructor inference via `user = User.new` (handled by CONSTRUCTOR_BINDING_SCANNERS in type-env.ts)
*/
/** Regex to extract @param annotations: `@param name [Type]` */
const YARD_PARAM_RE = /@param\s+(\w+)\s+\[([^\]]+)\]/g;
/**
* Extract the simple type name from a YARD type string.
* Handles:
* - Simple types: "String" → "String"
* - Qualified types: "Models::User" → "User"
* - Generic types: "Array<User>" → "Array"
* - Nullable types: "String, nil" → "String"
* - Union types: "String, Integer" → undefined (ambiguous)
*/
const extractYardTypeName = (yardType: string): string | undefined => {
const trimmed = yardType.trim();
// Handle nullable: "Type, nil" or "nil, Type"
const parts = trimmed.split(',').map(p => p.trim()).filter(p => p !== 'nil');
if (parts.length !== 1) return undefined; // ambiguous union
const typePart = parts[0];
// Handle qualified: "Models::User" → "User"
const segments = typePart.split('::');
const last = segments[segments.length - 1];
// Handle generic: "Array<User>" → "Array"
const genericMatch = last.match(/^(\w+)\s*[<{(]/);
if (genericMatch) return genericMatch[1];
// Simple identifier check
if (/^\w+$/.test(last)) return last;
return undefined;
};
/**
* Collect YARD @param annotations from comment nodes preceding a method definition.
* Returns a map of paramName → typeName.
*
* In tree-sitter-ruby, comments are sibling nodes that appear before the method node.
* We walk backwards through preceding siblings collecting consecutive comment nodes.
*/
const collectYardParams = (methodNode: SyntaxNode): Map<string, string> => {
const params = new Map<string, string>();
// In tree-sitter-ruby, YARD comments preceding a method inside a class body
// are placed as children of the `class` node, NOT as siblings of the `method`
// inside `body_statement`. The AST structure is:
//
// class
// constant = "ClassName"
// comment = "# @param ..." ← sibling of body_statement
// comment = "# @param ..." ← sibling of body_statement
// body_statement
// method ← method is here, no preceding siblings
//
// For top-level methods (outside classes), comments ARE direct siblings.
// We handle both by checking: if method has no preceding comment siblings,
// look at parent (body_statement) siblings instead.
const commentTexts: string[] = [];
const collectComments = (startNode: SyntaxNode): void => {
let sibling = startNode.previousSibling;
while (sibling) {
if (sibling.type === 'comment') {
commentTexts.unshift(sibling.text);
} else if (sibling.isNamed) {
break;
}
sibling = sibling.previousSibling;
}
};
// Try method's own siblings first (top-level methods)
collectComments(methodNode);
// If no comments found and parent is body_statement, check parent's siblings
if (commentTexts.length === 0 && methodNode.parent?.type === 'body_statement') {
collectComments(methodNode.parent);
}
// Parse all comment lines for @param annotations
const commentBlock = commentTexts.join('\n');
let match: RegExpExecArray | null;
// Reset regex state
YARD_PARAM_RE.lastIndex = 0;
while ((match = YARD_PARAM_RE.exec(commentBlock)) !== null) {
const paramName = match[1];
const rawType = match[2];
const typeName = extractYardTypeName(rawType);
if (typeName) {
params.set(paramName, typeName);
}
}
return params;
};
/**
* Ruby declaration node types that may carry YARD annotations.
* `method` is the tree-sitter-ruby node for `def name ... end`.
* `singleton_method` is `def self.name ... end`.
*
* We intercept these in extractDeclaration to pre-populate the env
* with YARD parameter types before the standard parameter walk.
*/
const DECLARATION_NODE_TYPES: ReadonlySet<string> = new Set([
'method',
'singleton_method',
'assignment', // For constructor inference: user = User.new
]);
/**
* Extract YARD annotations from method definitions.
* Pre-populates the scope env with parameter types before the
* standard parameter walk (which won't find types since Ruby has none).
*/
const extractDeclaration: TypeBindingExtractor = (node: SyntaxNode, env: Map<string, string>): void => {
if (node.type !== 'method' && node.type !== 'singleton_method') return;
const yardParams = collectYardParams(node);
if (yardParams.size === 0) return;
// Pre-populate env with YARD type bindings for each parameter
for (const [paramName, typeName] of yardParams) {
env.set(paramName, typeName);
}
};
/**
* Ruby parameter extraction.
* Ruby parameters (identifiers inside method_parameters) have no inline
* type annotations. YARD types are already populated by extractDeclaration,
* so this is a no-op — the bindings are already in the env.
*
* We still register this to maintain the LanguageTypeConfig contract.
*/
const extractParameter: ParameterExtractor = (_node: SyntaxNode, _env: Map<string, string>): void => {
// Ruby parameters have no type annotations.
// YARD types are pre-populated by extractDeclaration.
};
/**
* Ruby constructor inference: user = User.new
* Extracts type from assignments where the RHS is a `.new` call on a constant.
* This complements the CONSTRUCTOR_BINDING_SCANNERS in type-env.ts by also
* handling local class names (already defined in the same file).
*/
const extractInitializer: InitializerExtractor = (node: SyntaxNode, env: Map<string, string>, classNames: ClassNameLookup): void => {
if (node.type !== 'assignment') return;
const left = node.childForFieldName('left');
const right = node.childForFieldName('right');
if (!left || !right) return;
// Support both local variables (identifier) and constants (SERVICE = UserService.new)
if (left.type !== 'identifier' && left.type !== 'constant') return;
const varName = extractVarName(left);
if (!varName || env.has(varName)) return;
// Ruby constructor pattern: ClassName.new(args)
if (right.type !== 'call') return;
const method = right.childForFieldName('method');
if (!method || method.text !== 'new') return;
const receiver = right.childForFieldName('receiver');
if (!receiver || receiver.type !== 'constant') return;
const calleeName = receiver.text;
if (classNames.has(calleeName)) {
env.set(varName, calleeName);
}
};
export const typeConfig: LanguageTypeConfig = {
declarationNodeTypes: DECLARATION_NODE_TYPES,
extractDeclaration,
extractParameter,
extractInitializer,
};

View file

@ -122,6 +122,66 @@ export const TYPED_PARAMETER_TYPES = new Set([
'property_promotion_parameter', // PHP 8.0+ constructor promotion: __construct(private Foo $x)
]);
/**
* Extract type arguments from a generic type node.
* e.g., List<User, String> → ['User', 'String'], Vec<User> → ['User']
*
* Handles language-specific AST structures:
* - TS/Java/Rust/Go: generic_type > type_arguments > type nodes
* - C#: generic_type > type_argument_list > type nodes
* - Kotlin: generic_type > type_arguments > type_projection > type nodes
*
* Note: Go slices/maps use slice_type/map_type, not generic_type — those are
* NOT handled here. Use language-specific extractors for Go container types.
*
* @param typeNode A generic_type or parameterized_type AST node (or any node —
* returns [] for non-generic types).
* @returns Array of resolved type argument names. Unresolvable arguments are omitted.
*/
export const extractGenericTypeArgs = (typeNode: SyntaxNode): string[] => {
// Unwrap wrapper nodes that may sit above the generic_type
if (typeNode.type === 'type_annotation' || typeNode.type === 'type'
|| typeNode.type === 'user_type' || typeNode.type === 'nullable_type'
|| typeNode.type === 'optional_type') {
const inner = typeNode.firstNamedChild;
if (inner) return extractGenericTypeArgs(inner);
return [];
}
// Only process generic/parameterized type nodes
if (typeNode.type !== 'generic_type' && typeNode.type !== 'parameterized_type') {
return [];
}
// Find the type_arguments / type_argument_list child
let argsNode: SyntaxNode | null = null;
for (let i = 0; i < typeNode.namedChildCount; i++) {
const child = typeNode.namedChild(i);
if (child && (child.type === 'type_arguments' || child.type === 'type_argument_list')) {
argsNode = child;
break;
}
}
if (!argsNode) return [];
const result: string[] = [];
for (let i = 0; i < argsNode.namedChildCount; i++) {
let argNode = argsNode.namedChild(i);
if (!argNode) continue;
// Kotlin: type_arguments > type_projection > user_type > type_identifier
if (argNode.type === 'type_projection') {
argNode = argNode.firstNamedChild;
if (!argNode) continue;
}
const name = extractSimpleTypeName(argNode);
if (name) result.push(name);
}
return result;
};
/** Find the first named child with the given node type */
export const findChildByType = (node: SyntaxNode, type: string): SyntaxNode | null => {
for (let i = 0; i < node.namedChildCount; i++) {

View file

@ -79,6 +79,7 @@ interface ParsedSymbol {
nodeId: string;
type: string;
parameterCount?: number;
returnType?: string;
ownerId?: string;
}
@ -1078,6 +1079,7 @@ const processFileGroup = (
nodeId,
type: nodeLabel,
...(parameterCount !== undefined ? { parameterCount } : {}),
...(returnType !== undefined ? { returnType } : {}),
...(enclosingClassId ? { ownerId: enclosingClassId } : {}),
});

View file

@ -0,0 +1,12 @@
package main
import "example.com/returntype/models"
func GetUser(name string) *models.User {
return &models.User{Name: name}
}
func processUser() {
user := GetUser("alice")
user.Save()
}

View file

@ -0,0 +1,3 @@
module example.com/returntype
go 1.21

View file

@ -0,0 +1,9 @@
package models
type User struct {
Name string
}
func (u *User) Save() bool {
return true
}

View file

@ -0,0 +1,9 @@
import services.UserService;
public class App {
public static void processUser() {
UserService svc = new UserService();
var user = svc.getUser("alice");
user.save();
}
}

View file

@ -0,0 +1,13 @@
package models;
public class User {
private String name;
public User(String name) {
this.name = name;
}
public boolean save() {
return true;
}
}

View file

@ -0,0 +1,9 @@
package services;
import models.User;
public class UserService {
public User getUser(String name) {
return new User(name);
}
}

View file

@ -0,0 +1,5 @@
from service import get_user
def process_user():
user = get_user('alice')
user.save()

View file

@ -0,0 +1,6 @@
class User:
def __init__(self, name: str):
self.name = name
def save(self) -> bool:
return True

View file

@ -0,0 +1,4 @@
from models import User
def get_user(name: str) -> User:
return User(name)

View file

@ -0,0 +1,15 @@
class UserRepo
def save
true
end
def find_by_name(name)
true
end
end
class User
def greet
"hello"
end
end

View file

@ -0,0 +1,11 @@
require_relative './models'
class UserService
# @param repo [UserRepo] the repository
# @param user [User] the user to create
# @return [Boolean]
def create(repo, user)
repo.save
user.greet
end
end

View file

@ -0,0 +1,11 @@
import { getUser, fetchUserAsync } from './service';
function processUser() {
const user = getUser('alice');
user.save();
}
async function processUserAsync() {
const user = await fetchUserAsync('bob');
user.save();
}

View file

@ -0,0 +1,15 @@
export class User {
name: string;
constructor(name: string) {
this.name = name;
}
save(): boolean {
return true;
}
getName(): string {
return this.name;
}
}

View file

@ -0,0 +1,9 @@
import { User } from './models';
export function getUser(name: string): User {
return new User(name);
}
export function fetchUserAsync(name: string): Promise<User> {
return Promise.resolve(new User(name));
}

View file

@ -505,3 +505,46 @@ describe('Ruby constant constructor binding resolution', () => {
expect(validateCall).toBeDefined();
});
});
// ---------------------------------------------------------------------------
// YARD annotation type resolution: @param repo [UserRepo] → repo.save resolves
// ---------------------------------------------------------------------------
describe('Ruby YARD annotation type resolution', () => {
let result: PipelineResult;
beforeAll(async () => {
result = await runPipelineFromRepo(
path.join(FIXTURES, 'ruby-yard-annotations'),
() => {},
);
}, 60000);
it('detects UserRepo, User, and UserService classes', () => {
expect(getNodesByLabel(result, 'Class')).toContain('UserRepo');
expect(getNodesByLabel(result, 'Class')).toContain('User');
expect(getNodesByLabel(result, 'Class')).toContain('UserService');
});
it('detects save, find_by_name, greet, and create methods', () => {
const methods = getNodesByLabel(result, 'Method');
expect(methods).toContain('save');
expect(methods).toContain('find_by_name');
expect(methods).toContain('greet');
expect(methods).toContain('create');
});
it('resolves repo.save to UserRepo#save via YARD @param annotation', () => {
const calls = getRelationships(result, 'CALLS');
const saveCall = calls.find(c => c.target === 'save' && c.source === 'create');
expect(saveCall).toBeDefined();
expect(saveCall!.targetFilePath).toContain('models.rb');
});
it('resolves user.greet to User#greet via YARD @param annotation', () => {
const calls = getRelationships(result, 'CALLS');
const greetCall = calls.find(c => c.target === 'greet' && c.source === 'create');
expect(greetCall).toBeDefined();
expect(greetCall!.targetFilePath).toContain('models.rb');
});
});

View file

@ -1,5 +1,5 @@
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { processCallsFromExtracted } from '../../src/core/ingestion/call-processor.js';
import { processCallsFromExtracted, extractReturnTypeName } from '../../src/core/ingestion/call-processor.js';
import { createResolutionContext, type ResolutionContext } from '../../src/core/ingestion/resolution-context.js';
import { createKnowledgeGraph } from '../../src/core/graph/graph.js';
import type { ExtractedCall, FileConstructorBindings } from '../../src/core/ingestion/workers/parse-worker.js';
@ -333,6 +333,147 @@ describe('processCallsFromExtracted', () => {
expect(rels).toHaveLength(0);
});
// ---- Return type inference (Phase 4) ----
it('return type inference: binds variable to return type of callee', async () => {
// getUser() returns User, and User has a save() method
ctx.symbols.add('src/utils.ts', 'getUser', 'Function:src/utils.ts:getUser', 'Function', { returnType: 'User' });
ctx.symbols.add('src/models.ts', 'User', 'Class:src/models.ts:User', 'Class');
ctx.symbols.add('src/models.ts', 'save', 'Method:src/models.ts:save', 'Method', { ownerId: 'Class:src/models.ts:User' });
ctx.importMap.set('src/index.ts', new Set(['src/utils.ts', 'src/models.ts']));
// Binding: user = getUser() — getUser is not a class, so constructor path fails,
// but return type inference should kick in
const constructorBindings: FileConstructorBindings[] = [{
filePath: 'src/index.ts',
bindings: [
{ scope: 'main@0', varName: 'user', calleeName: 'getUser' },
],
}];
const calls: ExtractedCall[] = [{
filePath: 'src/index.ts',
calledName: 'save',
sourceId: 'Function:src/index.ts:main',
receiverName: 'user',
callForm: 'member',
}];
await processCallsFromExtracted(graph, calls, ctx, undefined, constructorBindings);
const rels = graph.relationships.filter(r => r.type === 'CALLS');
expect(rels).toHaveLength(1);
expect(rels[0].targetId).toBe('Method:src/models.ts:save');
});
it('return type inference: unwraps Promise<User> to User', async () => {
ctx.symbols.add('src/api.ts', 'fetchUser', 'Function:src/api.ts:fetchUser', 'Function', { returnType: 'Promise<User>' });
ctx.symbols.add('src/models.ts', 'User', 'Class:src/models.ts:User', 'Class');
ctx.symbols.add('src/models.ts', 'save', 'Method:src/models.ts:save', 'Method', { ownerId: 'Class:src/models.ts:User' });
ctx.importMap.set('src/index.ts', new Set(['src/api.ts', 'src/models.ts']));
const constructorBindings: FileConstructorBindings[] = [{
filePath: 'src/index.ts',
bindings: [
{ scope: 'main@0', varName: 'user', calleeName: 'fetchUser' },
],
}];
const calls: ExtractedCall[] = [{
filePath: 'src/index.ts',
calledName: 'save',
sourceId: 'Function:src/index.ts:main',
receiverName: 'user',
callForm: 'member',
}];
await processCallsFromExtracted(graph, calls, ctx, undefined, constructorBindings);
const rels = graph.relationships.filter(r => r.type === 'CALLS');
expect(rels).toHaveLength(1);
expect(rels[0].targetId).toBe('Method:src/models.ts:save');
});
it('return type inference: skips when return type is primitive', async () => {
ctx.symbols.add('src/utils.ts', 'getCount', 'Function:src/utils.ts:getCount', 'Function', { returnType: 'number' });
ctx.importMap.set('src/index.ts', new Set(['src/utils.ts']));
const constructorBindings: FileConstructorBindings[] = [{
filePath: 'src/index.ts',
bindings: [
{ scope: 'main@0', varName: 'count', calleeName: 'getCount' },
],
}];
const calls: ExtractedCall[] = [{
filePath: 'src/index.ts',
calledName: 'toString',
sourceId: 'Function:src/index.ts:main',
receiverName: 'count',
callForm: 'member',
}];
await processCallsFromExtracted(graph, calls, ctx, undefined, constructorBindings);
// No binding should be created for primitive return types
const rels = graph.relationships.filter(r => r.type === 'CALLS');
expect(rels).toHaveLength(0);
});
it('return type inference: skips ambiguous callees (multiple definitions)', async () => {
ctx.symbols.add('src/a.ts', 'getData', 'Function:src/a.ts:getData', 'Function', { returnType: 'User' });
ctx.symbols.add('src/b.ts', 'getData', 'Function:src/b.ts:getData', 'Function', { returnType: 'Repo' });
const constructorBindings: FileConstructorBindings[] = [{
filePath: 'src/index.ts',
bindings: [
{ scope: 'main@0', varName: 'data', calleeName: 'getData' },
],
}];
const calls: ExtractedCall[] = [{
filePath: 'src/index.ts',
calledName: 'save',
sourceId: 'Function:src/index.ts:main',
receiverName: 'data',
callForm: 'member',
}];
await processCallsFromExtracted(graph, calls, ctx, undefined, constructorBindings);
// Ambiguous callee — don't guess
const rels = graph.relationships.filter(r => r.type === 'CALLS');
expect(rels).toHaveLength(0);
});
it('return type inference: prefers constructor binding over return type', async () => {
// If the callee IS a class, constructor binding wins (existing behavior)
ctx.symbols.add('src/models.ts', 'User', 'Class:src/models.ts:User', 'Class');
ctx.symbols.add('src/models.ts', 'save', 'Method:src/models.ts:save', 'Method', { ownerId: 'Class:src/models.ts:User' });
ctx.importMap.set('src/index.ts', new Set(['src/models.ts']));
const constructorBindings: FileConstructorBindings[] = [{
filePath: 'src/index.ts',
bindings: [
{ scope: 'main@0', varName: 'user', calleeName: 'User' },
],
}];
const calls: ExtractedCall[] = [{
filePath: 'src/index.ts',
calledName: 'save',
sourceId: 'Function:src/index.ts:main',
receiverName: 'user',
callForm: 'member',
}];
await processCallsFromExtracted(graph, calls, ctx, undefined, constructorBindings);
const rels = graph.relationships.filter(r => r.type === 'CALLS');
expect(rels).toHaveLength(1);
expect(rels[0].targetId).toBe('Method:src/models.ts:save');
});
// ---- Scope-aware constructor bindings (Phase 3) ----
it('scope-aware bindings: same varName in different functions resolves to correct type', async () => {
@ -377,3 +518,86 @@ describe('processCallsFromExtracted', () => {
expect(rels[1].sourceId).toBe('Function:src/index.ts:processRepo');
});
});
describe('extractReturnTypeName', () => {
it('extracts simple type name', () => {
expect(extractReturnTypeName('User')).toBe('User');
});
it('unwraps Promise<User>', () => {
expect(extractReturnTypeName('Promise<User>')).toBe('User');
});
it('unwraps Option<User>', () => {
expect(extractReturnTypeName('Option<User>')).toBe('User');
});
it('unwraps Result<User, Error> to first type arg', () => {
expect(extractReturnTypeName('Result<User, Error>')).toBe('User');
});
it('strips nullable union: User | null', () => {
expect(extractReturnTypeName('User | null')).toBe('User');
});
it('strips nullable union: User | undefined', () => {
expect(extractReturnTypeName('User | undefined')).toBe('User');
});
it('strips nullable suffix: User?', () => {
expect(extractReturnTypeName('User?')).toBe('User');
});
it('strips Go pointer: *User', () => {
expect(extractReturnTypeName('*User')).toBe('User');
});
it('strips Rust reference: &User', () => {
expect(extractReturnTypeName('&User')).toBe('User');
});
it('strips Rust mutable reference: &mut User', () => {
expect(extractReturnTypeName('&mut User')).toBe('User');
});
it('returns undefined for primitives', () => {
expect(extractReturnTypeName('string')).toBeUndefined();
expect(extractReturnTypeName('number')).toBeUndefined();
expect(extractReturnTypeName('boolean')).toBeUndefined();
expect(extractReturnTypeName('void')).toBeUndefined();
expect(extractReturnTypeName('int')).toBeUndefined();
});
it('returns undefined for genuine union types', () => {
expect(extractReturnTypeName('User | Repo')).toBeUndefined();
});
it('returns undefined for empty string', () => {
expect(extractReturnTypeName('')).toBeUndefined();
});
it('extracts qualified type: models.User → User', () => {
expect(extractReturnTypeName('models.User')).toBe('User');
});
it('handles non-wrapper generics: Map<K, V> → Map', () => {
expect(extractReturnTypeName('Map<string, User>')).toBe('Map');
});
it('handles nested wrapper: Promise<Option<User>>', () => {
// Promise<Option<User>> → unwrap Promise → Option<User> → unwrap Option → User
expect(extractReturnTypeName('Promise<Option<User>>')).toBe('User');
});
it('unwraps Vec<User>', () => {
expect(extractReturnTypeName('Vec<User>')).toBe('User');
});
it('unwraps Optional<User>', () => {
expect(extractReturnTypeName('Optional<User>')).toBe('User');
});
it('returns undefined for lowercase non-class types', () => {
expect(extractReturnTypeName('error')).toBeUndefined();
});
});

View file

@ -0,0 +1,219 @@
import { describe, it, expect } from 'vitest';
import { extractGenericTypeArgs } from '../../src/core/ingestion/type-extractors/shared.js';
import type { SyntaxNode } from '../../src/core/ingestion/utils.js';
/**
* Create a minimal mock SyntaxNode for testing type extraction.
* Only the properties used by extractSimpleTypeName / extractGenericTypeArgs
* are populated — everything else is left as stubs.
*/
function mockNode(
type: string,
opts: {
text?: string;
namedChildren?: SyntaxNode[];
fields?: Record<string, SyntaxNode>;
} = {},
): SyntaxNode {
const children = opts.namedChildren ?? [];
const fields = opts.fields ?? {};
const text = opts.text ?? children.map((c) => c.text).join(', ');
return {
type,
text,
namedChildCount: children.length,
namedChild: (i: number) => children[i] ?? null,
firstNamedChild: children[0] ?? null,
lastNamedChild: children[children.length - 1] ?? null,
childForFieldName: (name: string) => fields[name] ?? null,
} as unknown as SyntaxNode;
}
// Helper: build a generic_type node with type_arguments
function genericType(
baseName: string,
typeArgNames: string[],
opts?: { argsNodeType?: string; wrapInProjection?: boolean },
): SyntaxNode {
const argsNodeType = opts?.argsNodeType ?? 'type_arguments';
const baseNode = mockNode('type_identifier', { text: baseName });
let argChildren = typeArgNames.map((name) =>
mockNode('type_identifier', { text: name }),
);
// Kotlin wraps each arg in type_projection > user_type > type_identifier
if (opts?.wrapInProjection) {
argChildren = typeArgNames.map((name) => {
const typeId = mockNode('type_identifier', { text: name });
const userType = mockNode('user_type', { namedChildren: [typeId] });
return mockNode('type_projection', { namedChildren: [userType] });
}) as unknown as SyntaxNode[];
}
const typeArgsNode = mockNode(argsNodeType, {
namedChildren: argChildren,
});
return mockNode('generic_type', {
namedChildren: [baseNode, typeArgsNode],
fields: { name: baseNode },
});
}
describe('extractGenericTypeArgs', () => {
describe('single type argument', () => {
it('extracts from TypeScript Array<User>', () => {
const node = genericType('Array', ['User']);
expect(extractGenericTypeArgs(node)).toEqual(['User']);
});
it('extracts from Java List<User>', () => {
const node = genericType('List', ['User']);
expect(extractGenericTypeArgs(node)).toEqual(['User']);
});
it('extracts from Rust Vec<User>', () => {
const node = genericType('Vec', ['User']);
expect(extractGenericTypeArgs(node)).toEqual(['User']);
});
it('extracts from C# List<User> (type_argument_list)', () => {
const node = genericType('List', ['User'], {
argsNodeType: 'type_argument_list',
});
expect(extractGenericTypeArgs(node)).toEqual(['User']);
});
});
describe('multiple type arguments', () => {
it('extracts from Java Map<String, User>', () => {
const node = genericType('Map', ['String', 'User']);
expect(extractGenericTypeArgs(node)).toEqual(['String', 'User']);
});
it('extracts from TS Map<string, number>', () => {
const node = genericType('Map', ['string', 'number']);
expect(extractGenericTypeArgs(node)).toEqual(['string', 'number']);
});
});
describe('Kotlin type_projection wrapping', () => {
it('extracts from Kotlin List<User> through type_projection', () => {
const node = genericType('List', ['User'], { wrapInProjection: true });
expect(extractGenericTypeArgs(node)).toEqual(['User']);
});
it('extracts from Kotlin Map<String, User> through type_projection', () => {
const node = genericType('Map', ['String', 'User'], {
wrapInProjection: true,
});
expect(extractGenericTypeArgs(node)).toEqual(['String', 'User']);
});
});
describe('parameterized_type (Java/Kotlin alternate node type)', () => {
it('extracts type arguments from parameterized_type', () => {
const baseNode = mockNode('type_identifier', { text: 'List' });
const argNode = mockNode('type_identifier', { text: 'User' });
const typeArgsNode = mockNode('type_arguments', {
namedChildren: [argNode],
});
const node = mockNode('parameterized_type', {
namedChildren: [baseNode, typeArgsNode],
fields: { name: baseNode },
});
expect(extractGenericTypeArgs(node)).toEqual(['User']);
});
});
describe('wrapper node unwrapping', () => {
it('unwraps type_annotation before extracting', () => {
const inner = genericType('Array', ['User']);
const wrapper = mockNode('type_annotation', { namedChildren: [inner] });
expect(extractGenericTypeArgs(wrapper)).toEqual(['User']);
});
it('unwraps nullable_type before extracting', () => {
const inner = genericType('List', ['User']);
const wrapper = mockNode('nullable_type', { namedChildren: [inner] });
expect(extractGenericTypeArgs(wrapper)).toEqual(['User']);
});
it('unwraps user_type before extracting (Kotlin)', () => {
const inner = genericType('MutableList', ['String']);
const wrapper = mockNode('user_type', { namedChildren: [inner] });
expect(extractGenericTypeArgs(wrapper)).toEqual(['String']);
});
});
describe('non-generic types return empty array', () => {
it('returns [] for plain type_identifier', () => {
const node = mockNode('type_identifier', { text: 'User' });
expect(extractGenericTypeArgs(node)).toEqual([]);
});
it('returns [] for identifier', () => {
const node = mockNode('identifier', { text: 'foo' });
expect(extractGenericTypeArgs(node)).toEqual([]);
});
it('returns [] for union_type', () => {
const node = mockNode('union_type', {
namedChildren: [
mockNode('type_identifier', { text: 'string' }),
mockNode('type_identifier', { text: 'number' }),
],
});
expect(extractGenericTypeArgs(node)).toEqual([]);
});
});
describe('nested generic types as arguments', () => {
it('extracts outer type arg names for nested generics', () => {
// Map<String, List<User>> — the second arg is itself a generic_type
// extractGenericTypeArgs should extract 'List' (via extractSimpleTypeName)
const innerGeneric = genericType('List', ['User']);
const stringNode = mockNode('type_identifier', { text: 'String' });
const typeArgsNode = mockNode('type_arguments', {
namedChildren: [stringNode, innerGeneric],
});
const baseNode = mockNode('type_identifier', { text: 'Map' });
const node = mockNode('generic_type', {
namedChildren: [baseNode, typeArgsNode],
fields: { name: baseNode },
});
// extractSimpleTypeName on a generic_type returns the base name
expect(extractGenericTypeArgs(node)).toEqual(['String', 'List']);
});
});
describe('edge cases', () => {
it('returns [] for generic_type with no type_arguments child', () => {
const baseNode = mockNode('type_identifier', { text: 'List' });
const node = mockNode('generic_type', {
namedChildren: [baseNode],
fields: { name: baseNode },
});
expect(extractGenericTypeArgs(node)).toEqual([]);
});
it('skips unresolvable type arguments', () => {
// If a child can't be resolved by extractSimpleTypeName, it is omitted
const baseNode = mockNode('type_identifier', { text: 'Fn' });
const unresolvedArg = mockNode('function_type', { text: '() => void' });
const resolvedArg = mockNode('type_identifier', { text: 'User' });
const typeArgsNode = mockNode('type_arguments', {
namedChildren: [unresolvedArg, resolvedArg],
});
const node = mockNode('generic_type', {
namedChildren: [baseNode, typeArgsNode],
fields: { name: baseNode },
});
expect(extractGenericTypeArgs(node)).toEqual(['User']);
});
});
});

View file

@ -101,6 +101,42 @@ describe('SymbolTable', () => {
});
});
describe('returnType metadata', () => {
it('stores returnType in SymbolDefinition', () => {
table.add('src/utils.ts', 'getUser', 'func:getUser', 'Function', { returnType: 'User' });
const def = table.lookupExactFull('src/utils.ts', 'getUser');
expect(def).toBeDefined();
expect(def!.returnType).toBe('User');
});
it('returnType is available via lookupFuzzy', () => {
table.add('src/utils.ts', 'getUser', 'func:getUser', 'Function', { returnType: 'Promise<User>' });
const results = table.lookupFuzzy('getUser');
expect(results).toHaveLength(1);
expect(results[0].returnType).toBe('Promise<User>');
});
it('omits returnType when not provided', () => {
table.add('src/utils.ts', 'helper', 'func:helper', 'Function');
const def = table.lookupExactFull('src/utils.ts', 'helper');
expect(def).toBeDefined();
expect(def!.returnType).toBeUndefined();
});
it('stores returnType alongside parameterCount and ownerId', () => {
table.add('src/models.ts', 'save', 'method:save', 'Method', {
parameterCount: 1,
returnType: 'boolean',
ownerId: 'class:User',
});
const def = table.lookupExactFull('src/models.ts', 'save');
expect(def).toBeDefined();
expect(def!.parameterCount).toBe(1);
expect(def!.returnType).toBe('boolean');
expect(def!.ownerId).toBe('class:User');
});
});
describe('clear', () => {
it('resets all state', () => {
table.add('src/a.ts', 'foo', 'func:foo', 'Function');

View file

@ -520,6 +520,86 @@ class UserService {
});
});
describe('Ruby YARD annotations', () => {
it('extracts @param type bindings from YARD comments', () => {
const tree = parse(`
class UserService
# @param repo [UserRepo] the repository
# @param name [String] the user's name
def create(repo, name)
repo.save
end
end
`, Ruby);
const { env } = buildTypeEnv(tree, 'ruby');
expect(flatGet(env, 'repo')).toBe('UserRepo');
expect(flatGet(env, 'name')).toBe('String');
});
it('handles qualified YARD types (Models::User → User)', () => {
const tree = parse(`
# @param user [Models::User] the user
def process(user)
end
`, Ruby);
const { env } = buildTypeEnv(tree, 'ruby');
expect(flatGet(env, 'user')).toBe('User');
});
it('handles nullable YARD types (String, nil → String)', () => {
const tree = parse(`
# @param name [String, nil] optional name
def greet(name)
end
`, Ruby);
const { env } = buildTypeEnv(tree, 'ruby');
expect(flatGet(env, 'name')).toBe('String');
});
it('skips ambiguous union YARD types (String, Integer → undefined)', () => {
const tree = parse(`
# @param value [String, Integer] mixed type
def process(value)
end
`, Ruby);
const { env } = buildTypeEnv(tree, 'ruby');
expect(flatGet(env, 'value')).toBeUndefined();
});
it('extracts no types when no YARD comments present', () => {
const tree = parse(`
def create(repo, name)
repo.save
end
`, Ruby);
const { env } = buildTypeEnv(tree, 'ruby');
expect(flatSize(env)).toBe(0);
});
it('extracts types from singleton method YARD comments', () => {
const tree = parse(`
class UserService
# @param name [String] the user's name
def self.find(name)
name
end
end
`, Ruby);
const { env } = buildTypeEnv(tree, 'ruby');
expect(flatGet(env, 'name')).toBe('String');
});
it('handles generic YARD types (Array<User> → Array)', () => {
const tree = parse(`
# @param users [Array<User>] list of users
def process(users)
end
`, Ruby);
const { env } = buildTypeEnv(tree, 'ruby');
expect(flatGet(env, 'users')).toBe('Array');
});
});
describe('super/base/parent resolution', () => {
it('resolves super to parent class name (TypeScript)', () => {
const code = `