Skip to content

Commit c1288d2

Browse files
committed
feat(core): add Go generics support to scanner
Extend GoScanner to detect and extract generic type information: - Generic structs: type Stack[T any] struct - Generic interfaces: type Comparable[T any] interface - Generic functions: func Map[T, U any](...) - Methods on generic types: func (s *Stack[T]) Push(...) Metadata additions: - isGeneric: true for any generic declaration - typeParameters: ['T any', 'K comparable'] array of type params Updated tree-sitter query to handle generic_type nodes in receivers. 9 new tests, 1552 total passing. Closes #130
1 parent 762caea commit c1288d2

File tree

3 files changed

+252
-8
lines changed

3 files changed

+252
-8
lines changed
Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
package generics
2+
3+
// Stack is a generic stack data structure.
4+
type Stack[T any] struct {
5+
items []T
6+
}
7+
8+
// Push adds an item to the stack.
9+
func (s *Stack[T]) Push(item T) {
10+
s.items = append(s.items, item)
11+
}
12+
13+
// Pop removes and returns the top item.
14+
func (s *Stack[T]) Pop() (T, bool) {
15+
if len(s.items) == 0 {
16+
var zero T
17+
return zero, false
18+
}
19+
item := s.items[len(s.items)-1]
20+
s.items = s.items[:len(s.items)-1]
21+
return item, true
22+
}
23+
24+
// Pair holds two values of potentially different types.
25+
type Pair[K comparable, V any] struct {
26+
Key K
27+
Value V
28+
}
29+
30+
// NewPair creates a new Pair.
31+
func NewPair[K comparable, V any](key K, value V) *Pair[K, V] {
32+
return &Pair[K, V]{Key: key, Value: value}
33+
}
34+
35+
// Map applies a function to each element of a slice.
36+
func Map[T, U any](slice []T, fn func(T) U) []U {
37+
result := make([]U, len(slice))
38+
for i, v := range slice {
39+
result[i] = fn(v)
40+
}
41+
return result
42+
}
43+
44+
// Filter returns elements that satisfy the predicate.
45+
func Filter[T any](slice []T, predicate func(T) bool) []T {
46+
var result []T
47+
for _, v := range slice {
48+
if predicate(v) {
49+
result = append(result, v)
50+
}
51+
}
52+
return result
53+
}
54+
55+
// Reduce reduces a slice to a single value.
56+
func Reduce[T, U any](slice []T, initial U, fn func(U, T) U) U {
57+
result := initial
58+
for _, v := range slice {
59+
result = fn(result, v)
60+
}
61+
return result
62+
}
63+
64+
// Comparable is an interface for comparable types.
65+
type Comparable[T any] interface {
66+
Compare(other T) int
67+
}
68+
69+
// Ordered is an interface for ordered types.
70+
type Ordered interface {
71+
~int | ~int8 | ~int16 | ~int32 | ~int64 |
72+
~uint | ~uint8 | ~uint16 | ~uint32 | ~uint64 |
73+
~float32 | ~float64 | ~string
74+
}
75+
76+
// Min returns the minimum of two ordered values.
77+
func Min[T Ordered](a, b T) T {
78+
if a < b {
79+
return a
80+
}
81+
return b
82+
}

packages/core/src/scanner/__tests__/go.test.ts

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -255,4 +255,103 @@ describe('GoScanner', () => {
255255
});
256256
});
257257
});
258+
259+
describe('generics support', () => {
260+
let genericsDocuments: Document[];
261+
262+
beforeAll(async () => {
263+
genericsDocuments = await scanner.scan(['generics.go'], fixturesDir);
264+
});
265+
266+
describe('generic structs', () => {
267+
it('should extract generic struct Stack[T any]', () => {
268+
const stack = genericsDocuments.find(
269+
(d) => d.metadata.name === 'Stack' && d.type === 'class'
270+
);
271+
expect(stack).toBeDefined();
272+
expect(stack?.metadata.custom?.isGeneric).toBe(true);
273+
expect(stack?.metadata.custom?.typeParameters).toEqual(['T any']);
274+
expect(stack?.metadata.signature).toContain('[T any]');
275+
});
276+
277+
it('should extract generic struct with multiple type parameters', () => {
278+
const pair = genericsDocuments.find(
279+
(d) => d.metadata.name === 'Pair' && d.type === 'class'
280+
);
281+
expect(pair).toBeDefined();
282+
expect(pair?.metadata.custom?.isGeneric).toBe(true);
283+
expect(pair?.metadata.custom?.typeParameters).toEqual(['K comparable', 'V any']);
284+
});
285+
});
286+
287+
describe('generic functions', () => {
288+
it('should extract generic function Map[T, U any]', () => {
289+
const mapFn = genericsDocuments.find(
290+
(d) => d.metadata.name === 'Map' && d.type === 'function'
291+
);
292+
expect(mapFn).toBeDefined();
293+
expect(mapFn?.metadata.custom?.isGeneric).toBe(true);
294+
expect(mapFn?.metadata.custom?.typeParameters).toEqual(['T', 'U any']);
295+
});
296+
297+
it('should extract generic function with constraints', () => {
298+
const minFn = genericsDocuments.find(
299+
(d) => d.metadata.name === 'Min' && d.type === 'function'
300+
);
301+
expect(minFn).toBeDefined();
302+
expect(minFn?.metadata.custom?.isGeneric).toBe(true);
303+
expect(minFn?.metadata.custom?.typeParameters).toEqual(['T Ordered']);
304+
});
305+
306+
it('should extract NewPair generic constructor', () => {
307+
const newPair = genericsDocuments.find(
308+
(d) => d.metadata.name === 'NewPair' && d.type === 'function'
309+
);
310+
expect(newPair).toBeDefined();
311+
expect(newPair?.metadata.custom?.isGeneric).toBe(true);
312+
expect(newPair?.metadata.custom?.typeParameters).toEqual(['K comparable', 'V any']);
313+
});
314+
});
315+
316+
describe('generic methods', () => {
317+
it('should extract method on generic receiver Stack[T]', () => {
318+
const push = genericsDocuments.find(
319+
(d) => d.metadata.name === 'Stack.Push' && d.type === 'method'
320+
);
321+
expect(push).toBeDefined();
322+
expect(push?.metadata.custom?.receiver).toBe('Stack');
323+
expect(push?.metadata.custom?.isGeneric).toBe(true);
324+
});
325+
326+
it('should extract Pop method on generic Stack', () => {
327+
const pop = genericsDocuments.find(
328+
(d) => d.metadata.name === 'Stack.Pop' && d.type === 'method'
329+
);
330+
expect(pop).toBeDefined();
331+
expect(pop?.metadata.custom?.receiver).toBe('Stack');
332+
expect(pop?.metadata.custom?.receiverPointer).toBe(true);
333+
});
334+
});
335+
336+
describe('generic interfaces', () => {
337+
it('should extract generic interface Comparable[T any]', () => {
338+
const comparable = genericsDocuments.find(
339+
(d) => d.metadata.name === 'Comparable' && d.type === 'interface'
340+
);
341+
expect(comparable).toBeDefined();
342+
expect(comparable?.metadata.custom?.isGeneric).toBe(true);
343+
expect(comparable?.metadata.custom?.typeParameters).toEqual(['T any']);
344+
});
345+
});
346+
347+
describe('non-generic items in generic file', () => {
348+
it('should not mark non-generic interface as generic', () => {
349+
const ordered = genericsDocuments.find(
350+
(d) => d.metadata.name === 'Ordered' && d.type === 'interface'
351+
);
352+
expect(ordered).toBeDefined();
353+
expect(ordered?.metadata.custom?.isGeneric).toBeUndefined();
354+
});
355+
});
356+
});
258357
});

packages/core/src/scanner/go.ts

Lines changed: 71 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -21,15 +21,17 @@ const GO_QUERIES = {
2121
name: (identifier) @name) @definition
2222
`,
2323

24-
// Method declarations with receivers
24+
// Method declarations with receivers (handles both regular and generic types)
2525
methods: `
2626
(method_declaration
2727
receiver: (parameter_list
2828
(parameter_declaration
2929
name: (identifier)? @receiver_name
3030
type: [
3131
(pointer_type (type_identifier) @receiver_type)
32+
(pointer_type (generic_type (type_identifier) @receiver_type))
3233
(type_identifier) @receiver_type
34+
(generic_type (type_identifier) @receiver_type)
3335
])) @receiver
3436
name: (field_identifier) @name) @definition
3537
`,
@@ -196,6 +198,9 @@ export class GoScanner implements Scanner {
196198
const exported = this.isExported(name);
197199
const snippet = this.truncateSnippet(fullText);
198200

201+
// Check for generics
202+
const { isGeneric, typeParameters } = this.extractTypeParameters(signature);
203+
199204
documents.push({
200205
id: `${file}:${name}:${startLine}`,
201206
text: this.buildEmbeddingText('function', name, signature, docstring),
@@ -210,7 +215,10 @@ export class GoScanner implements Scanner {
210215
exported,
211216
docstring,
212217
snippet,
213-
custom: isTestFile ? { isTest: true } : undefined,
218+
custom: {
219+
...(isTestFile ? { isTest: true } : {}),
220+
...(isGeneric ? { isGeneric, typeParameters } : {}),
221+
},
214222
},
215223
});
216224
}
@@ -240,7 +248,9 @@ export class GoScanner implements Scanner {
240248

241249
const methodName = nameCapture.node.text;
242250
const receiverType = receiverTypeCapture?.node.text || 'Unknown';
243-
const name = `${receiverType}.${methodName}`;
251+
// Strip type parameters from receiver for cleaner name (Stack[T] -> Stack)
252+
const baseReceiverType = receiverType.replace(/\[.*\]/, '');
253+
const name = `${baseReceiverType}.${methodName}`;
244254
const startLine = defCapture.node.startPosition.row + 1;
245255
const endLine = defCapture.node.endPosition.row + 1;
246256
const fullText = defCapture.node.text;
@@ -253,6 +263,12 @@ export class GoScanner implements Scanner {
253263
const receiverText = receiverCapture?.node.text || '';
254264
const receiverPointer = receiverText.includes('*');
255265

266+
// Check for generics (receiver has type params like Stack[T])
267+
const receiverHasGenerics = receiverType.includes('[');
268+
const { isGeneric: signatureHasGenerics, typeParameters } =
269+
this.extractTypeParameters(signature);
270+
const isGeneric = receiverHasGenerics || signatureHasGenerics;
271+
256272
documents.push({
257273
id: `${file}:${name}:${startLine}`,
258274
text: this.buildEmbeddingText('method', name, signature, docstring),
@@ -268,9 +284,10 @@ export class GoScanner implements Scanner {
268284
docstring,
269285
snippet,
270286
custom: {
271-
receiver: receiverType,
287+
receiver: baseReceiverType,
272288
receiverPointer,
273289
...(isTestFile ? { isTest: true } : {}),
290+
...(isGeneric ? { isGeneric, typeParameters } : {}),
274291
},
275292
},
276293
});
@@ -301,7 +318,13 @@ export class GoScanner implements Scanner {
301318
const startLine = defCapture.node.startPosition.row + 1;
302319
const endLine = defCapture.node.endPosition.row + 1;
303320
const fullText = defCapture.node.text;
304-
const signature = `type ${name} struct`;
321+
322+
// Check for generics in the full declaration text
323+
const { isGeneric, typeParameters } = this.extractTypeParameters(fullText);
324+
const signature = isGeneric
325+
? `type ${name}[${typeParameters?.join(', ')}] struct`
326+
: `type ${name} struct`;
327+
305328
const docstring = extractGoDocComment(sourceText, startLine);
306329
const exported = this.isExported(name);
307330
const snippet = this.truncateSnippet(fullText);
@@ -320,7 +343,10 @@ export class GoScanner implements Scanner {
320343
exported,
321344
docstring,
322345
snippet,
323-
custom: isTestFile ? { isTest: true } : undefined,
346+
custom: {
347+
...(isTestFile ? { isTest: true } : {}),
348+
...(isGeneric ? { isGeneric, typeParameters } : {}),
349+
},
324350
},
325351
});
326352
}
@@ -350,7 +376,13 @@ export class GoScanner implements Scanner {
350376
const startLine = defCapture.node.startPosition.row + 1;
351377
const endLine = defCapture.node.endPosition.row + 1;
352378
const fullText = defCapture.node.text;
353-
const signature = `type ${name} interface`;
379+
380+
// Check for generics in the full declaration text
381+
const { isGeneric, typeParameters } = this.extractTypeParameters(fullText);
382+
const signature = isGeneric
383+
? `type ${name}[${typeParameters?.join(', ')}] interface`
384+
: `type ${name} interface`;
385+
354386
const docstring = extractGoDocComment(sourceText, startLine);
355387
const exported = this.isExported(name);
356388
const snippet = this.truncateSnippet(fullText);
@@ -369,7 +401,10 @@ export class GoScanner implements Scanner {
369401
exported,
370402
docstring,
371403
snippet,
372-
custom: isTestFile ? { isTest: true } : undefined,
404+
custom: {
405+
...(isTestFile ? { isTest: true } : {}),
406+
...(isGeneric ? { isGeneric, typeParameters } : {}),
407+
},
373408
},
374409
});
375410
}
@@ -498,6 +533,34 @@ export class GoScanner implements Scanner {
498533
return fullText.slice(0, braceIndex).trim();
499534
}
500535

536+
/**
537+
* Extract type parameters from a generic declaration.
538+
* Returns { isGeneric, typeParameters } where typeParameters is an array like ["T any", "K comparable"]
539+
*/
540+
private extractTypeParameters(signature: string): {
541+
isGeneric: boolean;
542+
typeParameters?: string[];
543+
} {
544+
// Match type parameters in brackets: func Name[T any, K comparable](...) or type Name[T any] struct
545+
// Look for [ before ( for functions, or [ after type name for types
546+
const typeParamMatch = signature.match(/\[([^\]]+)\]/);
547+
if (!typeParamMatch) {
548+
return { isGeneric: false };
549+
}
550+
551+
const params = typeParamMatch[1];
552+
// Split by comma, but handle constraints like "~int | ~string"
553+
const typeParameters = params
554+
.split(/,\s*/)
555+
.map((p) => p.trim())
556+
.filter((p) => p.length > 0);
557+
558+
return {
559+
isGeneric: true,
560+
typeParameters,
561+
};
562+
}
563+
501564
/**
502565
* Build embedding text for vector search
503566
*/

0 commit comments

Comments
 (0)