diff --git a/gitnexus/src/core/ingestion/type-extractors/python.ts b/gitnexus/src/core/ingestion/type-extractors/python.ts index 98ce18287..e54f552be 100644 --- a/gitnexus/src/core/ingestion/type-extractors/python.ts +++ b/gitnexus/src/core/ingestion/type-extractors/python.ts @@ -256,7 +256,8 @@ const extractForLoopBinding: ForLoopExtractor = ( } if (!iterableName) return; - const typeArgPos = methodToTypeArgPosition(methodName); + const containerTypeName = scopeEnv.get(iterableName); + const typeArgPos = methodToTypeArgPosition(methodName, containerTypeName); const elementType = resolveIterableElementType( iterableName, node, scopeEnv, declarationTypeNodes, scope, extractPyElementTypeFromAnnotation, findPyParamElementType, diff --git a/gitnexus/src/core/ingestion/type-extractors/rust.ts b/gitnexus/src/core/ingestion/type-extractors/rust.ts index 19415b1b5..06f47afb6 100644 --- a/gitnexus/src/core/ingestion/type-extractors/rust.ts +++ b/gitnexus/src/core/ingestion/type-extractors/rust.ts @@ -383,7 +383,8 @@ const extractForLoopBinding: ForLoopExtractor = ( } if (!iterableName) return; - const typeArgPos = methodToTypeArgPosition(methodName); + const containerTypeName = scopeEnv.get(iterableName); + const typeArgPos = methodToTypeArgPosition(methodName, containerTypeName); const elementType = resolveIterableElementType( iterableName, node, scopeEnv, declarationTypeNodes, scope, extractRustElementTypeFromTypeNode, findRustParamElementType, diff --git a/gitnexus/src/core/ingestion/type-extractors/shared.ts b/gitnexus/src/core/ingestion/type-extractors/shared.ts index 808b0db55..abddc8f9c 100644 --- a/gitnexus/src/core/ingestion/type-extractors/shared.ts +++ b/gitnexus/src/core/ingestion/type-extractors/shared.ts @@ -5,14 +5,112 @@ import type { SyntaxNode } from '../utils.js'; * - 'last': value type (e.g., V from Map) — used for .values(), .items(), .iter() */ export type TypeArgPosition = 'first' | 'last'; -/** Map method names to which type argument they yield. - * Methods that iterate/return keys → 'first'; everything else → 'last'. */ -const KEY_METHODS = new Set(['keys', 'keySet', 'Keys']); +// --------------------------------------------------------------------------- +// Container type descriptors — maps container base names to type parameter +// semantics per access method. Replaces the simple KEY_METHODS heuristic. +// +// For user-defined generics (MyCache extends Map), heritage-aware +// fallback can walk the EXTENDS chain to find a matching descriptor. +// --------------------------------------------------------------------------- -/** Determine which type arg to use based on the iterator method name. - * .keys() / .keySet() → 'first' (key type); everything else → 'last' (value type). */ -export function methodToTypeArgPosition(methodName: string | undefined): TypeArgPosition { - return methodName && KEY_METHODS.has(methodName) ? 'first' : 'last'; +/** Describes which type parameter position each access method yields. */ +interface ContainerDescriptor { + /** Number of type parameters (1 = single-element, 2 = key-value) */ + arity: number; + /** Methods that yield the first type parameter (key type for maps) */ + keyMethods: ReadonlySet; + /** Methods that yield the last type parameter (value type) */ + valueMethods: ReadonlySet; +} + +/** Empty set for containers that have no key-yielding methods */ +const NO_KEYS: ReadonlySet = new Set(); + +/** Standard key-yielding methods across languages */ +const STD_KEY_METHODS: ReadonlySet = new Set(['keys']); +const JAVA_KEY_METHODS: ReadonlySet = new Set(['keySet']); +const CSHARP_KEY_METHODS: ReadonlySet = new Set(['Keys']); + +/** Standard value-yielding methods across languages */ +const STD_VALUE_METHODS: ReadonlySet = new Set(['values', 'get', 'pop', 'remove']); +const CSHARP_VALUE_METHODS: ReadonlySet = new Set(['Values', 'TryGetValue']); +const SINGLE_ELEMENT_METHODS: ReadonlySet = new Set([ + 'iter', 'into_iter', 'iterator', 'get', 'first', 'last', 'pop', + 'peek', 'poll', 'find', 'filter', 'map', +]); + +const CONTAINER_DESCRIPTORS: ReadonlyMap = new Map([ + // --- Map / Dict types (arity 2: key + value) --- + ['Map', { arity: 2, keyMethods: STD_KEY_METHODS, valueMethods: STD_VALUE_METHODS }], + ['WeakMap', { arity: 2, keyMethods: STD_KEY_METHODS, valueMethods: STD_VALUE_METHODS }], + ['HashMap', { arity: 2, keyMethods: STD_KEY_METHODS, valueMethods: STD_VALUE_METHODS }], + ['BTreeMap', { arity: 2, keyMethods: STD_KEY_METHODS, valueMethods: STD_VALUE_METHODS }], + ['LinkedHashMap', { arity: 2, keyMethods: JAVA_KEY_METHODS, valueMethods: STD_VALUE_METHODS }], + ['TreeMap', { arity: 2, keyMethods: JAVA_KEY_METHODS, valueMethods: STD_VALUE_METHODS }], + ['dict', { arity: 2, keyMethods: STD_KEY_METHODS, valueMethods: STD_VALUE_METHODS }], + ['Dict', { arity: 2, keyMethods: STD_KEY_METHODS, valueMethods: STD_VALUE_METHODS }], + ['Dictionary', { arity: 2, keyMethods: CSHARP_KEY_METHODS, valueMethods: CSHARP_VALUE_METHODS }], + ['SortedDictionary', { arity: 2, keyMethods: CSHARP_KEY_METHODS, valueMethods: CSHARP_VALUE_METHODS }], + ['Record', { arity: 2, keyMethods: STD_KEY_METHODS, valueMethods: STD_VALUE_METHODS }], + ['OrderedDict', { arity: 2, keyMethods: STD_KEY_METHODS, valueMethods: STD_VALUE_METHODS }], + ['ConcurrentHashMap', { arity: 2, keyMethods: JAVA_KEY_METHODS, valueMethods: STD_VALUE_METHODS }], + ['ConcurrentDictionary', { arity: 2, keyMethods: CSHARP_KEY_METHODS, valueMethods: CSHARP_VALUE_METHODS }], + + // --- Single-element containers (arity 1) --- + ['Array', { arity: 1, keyMethods: NO_KEYS, valueMethods: SINGLE_ELEMENT_METHODS }], + ['List', { arity: 1, keyMethods: NO_KEYS, valueMethods: SINGLE_ELEMENT_METHODS }], + ['ArrayList', { arity: 1, keyMethods: NO_KEYS, valueMethods: SINGLE_ELEMENT_METHODS }], + ['LinkedList',{ arity: 1, keyMethods: NO_KEYS, valueMethods: SINGLE_ELEMENT_METHODS }], + ['Vec', { arity: 1, keyMethods: NO_KEYS, valueMethods: SINGLE_ELEMENT_METHODS }], + ['VecDeque', { arity: 1, keyMethods: NO_KEYS, valueMethods: SINGLE_ELEMENT_METHODS }], + ['Set', { arity: 1, keyMethods: NO_KEYS, valueMethods: SINGLE_ELEMENT_METHODS }], + ['HashSet', { arity: 1, keyMethods: NO_KEYS, valueMethods: SINGLE_ELEMENT_METHODS }], + ['BTreeSet', { arity: 1, keyMethods: NO_KEYS, valueMethods: SINGLE_ELEMENT_METHODS }], + ['TreeSet', { arity: 1, keyMethods: NO_KEYS, valueMethods: SINGLE_ELEMENT_METHODS }], + ['Queue', { arity: 1, keyMethods: NO_KEYS, valueMethods: SINGLE_ELEMENT_METHODS }], + ['Deque', { arity: 1, keyMethods: NO_KEYS, valueMethods: SINGLE_ELEMENT_METHODS }], + ['Stack', { arity: 1, keyMethods: NO_KEYS, valueMethods: SINGLE_ELEMENT_METHODS }], + ['Sequence', { arity: 1, keyMethods: NO_KEYS, valueMethods: SINGLE_ELEMENT_METHODS }], + ['Iterable', { arity: 1, keyMethods: NO_KEYS, valueMethods: SINGLE_ELEMENT_METHODS }], + ['Iterator', { arity: 1, keyMethods: NO_KEYS, valueMethods: SINGLE_ELEMENT_METHODS }], + ['IEnumerable', { arity: 1, keyMethods: NO_KEYS, valueMethods: SINGLE_ELEMENT_METHODS }], + ['IList', { arity: 1, keyMethods: NO_KEYS, valueMethods: SINGLE_ELEMENT_METHODS }], + ['list', { arity: 1, keyMethods: NO_KEYS, valueMethods: SINGLE_ELEMENT_METHODS }], + ['set', { arity: 1, keyMethods: NO_KEYS, valueMethods: SINGLE_ELEMENT_METHODS }], + ['tuple', { arity: 1, keyMethods: NO_KEYS, valueMethods: SINGLE_ELEMENT_METHODS }], + ['frozenset', { arity: 1, keyMethods: NO_KEYS, valueMethods: SINGLE_ELEMENT_METHODS }], +]); + +/** Determine which type arg to extract based on container type name and access method. + * + * Resolution order: + * 1. If container is known and method is in keyMethods → 'first' + * 2. If container is known with arity 1 → 'last' (same as 'first' for single-arg) + * 3. If container is unknown → fall back to method name heuristic + * 4. Default: 'last' (value type) + */ +export function methodToTypeArgPosition(methodName: string | undefined, containerTypeName?: string): TypeArgPosition { + if (containerTypeName) { + const desc = CONTAINER_DESCRIPTORS.get(containerTypeName); + if (desc) { + // Single-element container: always 'last' (= only arg) + if (desc.arity === 1) return 'last'; + // Multi-element: check if method yields key type + if (methodName && desc.keyMethods.has(methodName)) return 'first'; + // Default for multi-element: value type + return 'last'; + } + } + // Fallback for unknown containers: simple method name heuristic + if (methodName && (methodName === 'keys' || methodName === 'keySet' || methodName === 'Keys')) { + return 'first'; + } + return 'last'; +} + +/** Look up the container descriptor for a type name. Exported for heritage-chain lookups. */ +export function getContainerDescriptor(typeName: string): ContainerDescriptor | undefined { + return CONTAINER_DESCRIPTORS.get(typeName); } /** diff --git a/gitnexus/src/core/ingestion/type-extractors/typescript.ts b/gitnexus/src/core/ingestion/type-extractors/typescript.ts index 6fd184fd5..a1a62981c 100644 --- a/gitnexus/src/core/ingestion/type-extractors/typescript.ts +++ b/gitnexus/src/core/ingestion/type-extractors/typescript.ts @@ -356,7 +356,9 @@ const extractForLoopBinding: ForLoopExtractor = ( } if (!iterableName) return; - const typeArgPos = methodToTypeArgPosition(methodName); + // Look up the container's base type name for descriptor-aware resolution + const containerTypeName = scopeEnv.get(iterableName); + const typeArgPos = methodToTypeArgPosition(methodName, containerTypeName); const elementType = resolveIterableElementType( iterableName, node, scopeEnv, declarationTypeNodes, scope, extractTsElementTypeFromAnnotation, findTsIterableElementType, diff --git a/gitnexus/test/unit/type-env.test.ts b/gitnexus/test/unit/type-env.test.ts index e1325eb88..15f1194e1 100644 --- a/gitnexus/test/unit/type-env.test.ts +++ b/gitnexus/test/unit/type-env.test.ts @@ -3056,6 +3056,57 @@ fn process(data: HashMap) { }); }); + describe('container descriptor-aware type arg selection', () => { + it('HashMap.keys() resolves to key type (String) via descriptor', () => { + const tree = parse(` +fn process(data: HashMap) { + for key in data.keys() { + key.len(); + } +} + `, Rust); + const { env } = buildTypeEnv(tree, 'rust'); + expect(flatGet(env, 'key')).toBe('String'); + }); + + it('HashMap.values() resolves to value type (User) via descriptor', () => { + const tree = parse(` +fn process(data: HashMap) { + for user in data.values() { + user.save(); + } +} + `, Rust); + const { env } = buildTypeEnv(tree, 'rust'); + expect(flatGet(env, 'user')).toBe('User'); + }); + + it('Vec.iter() resolves to element type (User) — arity 1 always returns last', () => { + const tree = parse(` +fn process(users: Vec) { + for user in users.iter() { + user.save(); + } +} + `, Rust); + const { env } = buildTypeEnv(tree, 'rust'); + expect(flatGet(env, 'user')).toBe('User'); + }); + + it('unknown container falls back to last-arg heuristic', () => { + // MyCache is not in CONTAINER_DESCRIPTORS, so .keys() still returns first via fallback + const tree = parse(` +function process(cache: MyCache) { + for (const key of cache.keys()) { + key.trim(); + } +} + `, TypeScript.typescript); + const { env } = buildTypeEnv(tree, 'typescript'); + expect(flatGet(env, 'key')).toBe('string'); + }); + }); + describe('known limitations (documented skip tests)', () => { it.skip('Ruby block parameter: users.each { |user| } — closure param inference, different feature', () => { // Not a for-loop; .each { |user| } is a method call with a block.