From 6978acc92e9901bb6402c88a2bf0530615497e1c Mon Sep 17 00:00:00 2001 From: ctype_lab Date: Thu, 6 Aug 2026 16:49:14 +0900 Subject: [PATCH] feat(extraction): model union declarations distinctly --- __tests__/extraction.test.ts | 28 +++++++++++---- src/extraction/languages/c-cpp.ts | 23 +++++++------ src/extraction/languages/objc.ts | 13 +++---- src/extraction/languages/rust.ts | 8 ++--- src/extraction/tree-sitter-types.ts | 2 ++ src/extraction/tree-sitter.ts | 53 +++++++++++++++++++++-------- src/types.ts | 1 + 7 files changed, 84 insertions(+), 44 deletions(-) diff --git a/__tests__/extraction.test.ts b/__tests__/extraction.test.ts index d8a4315..2926588 100644 --- a/__tests__/extraction.test.ts +++ b/__tests__/extraction.test.ts @@ -1194,11 +1194,11 @@ impl Describe for Reg { `; const result = extractFromSource('reg.rs', code); - // A union is a type definition, not an alias — it must be a node, or the - // impl below has no source endpoint to hang off. + // A union is a first-class type definition, not an alias — it must be a + // node, or the impl below has no source endpoint to hang off. const reg = result.nodes.find((n) => n.name === 'Reg'); expect(reg).toBeDefined(); - expect(reg?.kind).toBe('struct'); + expect(reg?.kind).toBe('union'); const implRef = result.unresolvedReferences.find( (r) => r.referenceKind === 'implements' && r.referenceName === 'Describe' @@ -5704,7 +5704,7 @@ static unsigned int hdr_raw(union packet_hdr *h) { return h->raw; } const hdr = result.nodes.find((n) => n.name === 'packet_hdr'); expect(hdr).toBeDefined(); - expect(hdr?.kind).toBe('struct'); + expect(hdr?.kind).toBe('union'); // Same rule as `struct Foo;`: bodiless is a forward declaration, so it must // not mint a phantom node beside the real definition. @@ -5725,7 +5725,7 @@ typedef union { const result = extractFromSource('word.c', code); const word = result.nodes.find((n) => n.name === 'word_t'); - expect(word?.kind).toBe('struct'); + expect(word?.kind).toBe('union'); // Resolved through the typedef the same way `typedef struct { … } X;` is, // so the anonymous union body does not become its own node. expect(result.nodes.some((n) => n.name === '')).toBe(false); @@ -5742,7 +5742,7 @@ union Value { const result = extractFromSource('value.cpp', code); const value = result.nodes.find((n) => n.name === 'Value'); - expect(value?.kind).toBe('struct'); + expect(value?.kind).toBe('union'); const asInt = result.nodes.find((n) => n.name === 'as_int'); expect(asInt).toBeDefined(); @@ -8446,6 +8446,22 @@ void helperFunction(int count) { expect(imports).toContain('MyClass.h'); }); + it('extracts union declarations as first-class union nodes', () => { + const code = ` +typedef union { + unsigned int raw; + float value; +} NumberBits; + +union opaque_bits; +`; + const result = extractFromSource('NumberBits.m', code); + + const numberBits = result.nodes.find((n) => n.name === 'NumberBits'); + expect(numberBits?.kind).toBe('union'); + expect(result.nodes.some((n) => n.name === 'opaque_bits')).toBe(false); + }); + it('should record inheritance and protocol conformance', () => { const result = extractFromSource('App.m', sample); const extendsRefs = result.unresolvedReferences.filter((r) => r.referenceKind === 'extends'); diff --git a/src/extraction/languages/c-cpp.ts b/src/extraction/languages/c-cpp.ts index 708985e..cdd573b 100644 --- a/src/extraction/languages/c-cpp.ts +++ b/src/extraction/languages/c-cpp.ts @@ -186,11 +186,10 @@ export const cExtractor: LanguageExtractor = { classTypes: [], methodTypes: [], interfaceTypes: [], - // `union U { … };` is a type DEFINITION, same as `struct U { … };` — it - // declares a named type whose members other code refers to. Extracted with - // kind `struct` because NodeKind has no `union`; a bodiless `union U;` is a - // forward declaration and still falls out via extractStruct's body guard. - structTypes: ['struct_specifier', 'union_specifier'], + structTypes: ['struct_specifier'], + // A bodiless `union U;` is a forward declaration; the aggregate extractor + // applies the same body requirement it uses for C structs. + unionTypes: ['union_specifier'], enumTypes: ['enum_specifier'], enumMemberTypes: ['enumerator'], typeAliasTypes: ['type_definition'], // typedef @@ -219,10 +218,11 @@ export const cExtractor: LanguageExtractor = { if (!child) continue; if (child.type === 'enum_specifier' && getChildByField(child, 'body')) return 'enum'; if ( - (child.type === 'struct_specifier' || child.type === 'union_specifier') && + child.type === 'struct_specifier' && getChildByField(child, 'body') ) return 'struct'; + if (child.type === 'union_specifier' && getChildByField(child, 'body')) return 'union'; } return undefined; }, @@ -1561,10 +1561,10 @@ export const cppExtractor: LanguageExtractor = { skipBodilessClass: true, methodTypes: ['function_definition'], interfaceTypes: [], - // See the C extractor: a named `union U { … };` is a definition, not an - // alias. C++ unions additionally carry member functions, which extract - // through the same body walk as a struct's. - structTypes: ['struct_specifier', 'union_specifier'], + structTypes: ['struct_specifier'], + // C++ unions additionally carry member functions, which extract through the + // same aggregate-body walk as structs while preserving their distinct kind. + unionTypes: ['union_specifier'], enumTypes: ['enum_specifier'], enumMemberTypes: ['enumerator'], typeAliasTypes: ['type_definition', 'alias_declaration'], // typedef and using @@ -1601,10 +1601,11 @@ export const cppExtractor: LanguageExtractor = { if (!child) continue; if (child.type === 'enum_specifier' && getChildByField(child, 'body')) return 'enum'; if ( - (child.type === 'struct_specifier' || child.type === 'union_specifier') && + child.type === 'struct_specifier' && getChildByField(child, 'body') ) return 'struct'; + if (child.type === 'union_specifier' && getChildByField(child, 'body')) return 'union'; } return undefined; }, diff --git a/src/extraction/languages/objc.ts b/src/extraction/languages/objc.ts index 668e226..9bada34 100644 --- a/src/extraction/languages/objc.ts +++ b/src/extraction/languages/objc.ts @@ -102,9 +102,9 @@ export const objcExtractor: LanguageExtractor = { methodTypes: ['method_definition'], interfaceTypes: ['protocol_declaration'], interfaceKind: 'protocol', - // Objective-C is a C superset: `union U { … };` is a definition, same as in - // the C extractor. - structTypes: ['struct_specifier', 'union_specifier'], + structTypes: ['struct_specifier'], + // Objective-C is a C superset: union declarations preserve their own kind. + unionTypes: ['union_specifier'], enumTypes: ['enum_specifier'], enumMemberTypes: ['enumerator'], typeAliasTypes: ['type_definition'], @@ -130,12 +130,9 @@ export const objcExtractor: LanguageExtractor = { const child = node.namedChild(i); if (!child) continue; if (child.type === 'enum_specifier' && getChildByField(child, 'body')) return 'enum'; - // `typedef union { … } name;` resolves like `typedef struct` — see the C extractor. - if ( - (child.type === 'struct_specifier' || child.type === 'union_specifier') && - getChildByField(child, 'body') - ) + if (child.type === 'struct_specifier' && getChildByField(child, 'body')) return 'struct'; + if (child.type === 'union_specifier' && getChildByField(child, 'body')) return 'union'; } return undefined; }, diff --git a/src/extraction/languages/rust.ts b/src/extraction/languages/rust.ts index 2015e62..6d91bf6 100644 --- a/src/extraction/languages/rust.ts +++ b/src/extraction/languages/rust.ts @@ -41,10 +41,10 @@ export const rustExtractor: LanguageExtractor = { classTypes: [], // Rust has impl blocks methodTypes: ['function_item', 'function_signature_item'], interfaceTypes: ['trait_item'], - // `union U { … }` is a definition like `struct U { … }` — same `body:` - // (`field_declaration_list`) and the same `impl Trait for U` attachment - // point. Extracted with kind `struct` because NodeKind has no `union`. - structTypes: ['struct_item', 'union_item'], + structTypes: ['struct_item'], + // Unions share struct member syntax and impl attachment, but retain their + // distinct semantic kind in the graph. + unionTypes: ['union_item'], enumTypes: ['enum_item'], enumMemberTypes: ['enum_variant'], typeAliasTypes: ['type_item'], // Rust type aliases diff --git a/src/extraction/tree-sitter-types.ts b/src/extraction/tree-sitter-types.ts index 2a02b47..6808895 100644 --- a/src/extraction/tree-sitter-types.ts +++ b/src/extraction/tree-sitter-types.ts @@ -102,6 +102,8 @@ export interface LanguageExtractor { interfaceTypes: string[]; /** Node types that represent structs */ structTypes: string[]; + /** Node types that represent unions */ + unionTypes?: string[]; /** Node types that represent enums */ enumTypes: string[]; /** Node types that represent enum members/cases (e.g. Swift: 'enum_entry', Rust: 'enum_variant') */ diff --git a/src/extraction/tree-sitter.ts b/src/extraction/tree-sitter.ts index 9e53e62..8d71d7f 100644 --- a/src/extraction/tree-sitter.ts +++ b/src/extraction/tree-sitter.ts @@ -1066,6 +1066,11 @@ export class TreeSitterExtractor { this.extractStruct(node); skipChildren = true; // extractStruct visits body children } + // Check for union declarations + else if (this.extractor.unionTypes?.includes(nodeType)) { + this.extractUnion(node); + skipChildren = true; // extractUnion visits body children + } // Check for enum declarations else if (this.extractor.enumTypes.includes(nodeType)) { this.extractEnum(node); @@ -1487,7 +1492,7 @@ export class TreeSitterExtractor { /** * Check if the current node stack indicates we are inside a class-like node - * (class, struct, interface, trait). File nodes do not count as class-like. + * (class, struct, union, interface, trait). File nodes do not count as class-like. */ private isInsideClassLikeNode(): boolean { if (this.nodeStack.length === 0) return false; @@ -1498,6 +1503,7 @@ export class TreeSitterExtractor { return ( parentNode.kind === 'class' || parentNode.kind === 'struct' || + parentNode.kind === 'union' || parentNode.kind === 'interface' || parentNode.kind === 'trait' || parentNode.kind === 'enum' || @@ -1807,7 +1813,7 @@ export class TreeSitterExtractor { (n) => n.name === receiverType && n.filePath === this.filePath && - (n.kind === 'struct' || n.kind === 'class' || n.kind === 'enum' || n.kind === 'trait') + (n.kind === 'struct' || n.kind === 'union' || n.kind === 'class' || n.kind === 'enum' || n.kind === 'trait') ); if (ownerNode) { this.edges.push({ @@ -1873,6 +1879,16 @@ export class TreeSitterExtractor { * Extract a struct */ private extractStruct(node: SyntaxNode): void { + this.extractAggregate(node, 'struct'); + } + + /** Extract a union while sharing the member-walk behavior of aggregate types. */ + private extractUnion(node: SyntaxNode): void { + this.extractAggregate(node, 'union'); + } + + /** Extract a struct-like declaration without conflating its semantic kind. */ + private extractAggregate(node: SyntaxNode, kind: 'struct' | 'union'): void { if (!this.extractor) return; // Skip forward declarations and type references (no body = not a definition) @@ -1886,24 +1902,24 @@ export class TreeSitterExtractor { const visibility = this.extractor.getVisibility?.(node); const isExported = this.extractor.isExported?.(node, this.source); - const structNode = this.createNode('struct', name, node, { + const aggregateNode = this.createNode(kind, name, node, { docstring, visibility, isExported, }); - if (!structNode) return; + if (!aggregateNode) return; // Extract inheritance (e.g. Swift: struct HTTPMethod: RawRepresentable) - this.extractInheritance(node, structNode.id); + this.extractInheritance(node, aggregateNode.id); // C# primary-constructor parameter dependencies (`struct P(int x)`, and // `record struct M(decimal Amount)` which the grammar nests here). - this.extractCsharpPrimaryCtorParamRefs(node, structNode.id); + this.extractCsharpPrimaryCtorParamRefs(node, aggregateNode.id); // Push to stack for field extraction (bodiless positional records have // no members to visit) if (body) { - this.nodeStack.push(structNode.id); + this.nodeStack.push(aggregateNode.id); for (let i = 0; i < body.namedChildCount; i++) { const child = body.namedChild(i); if (child) { @@ -2905,17 +2921,20 @@ export class TreeSitterExtractor { // (e.g. Go: `type Foo struct { ... }` is a type_spec wrapping struct_type) const resolvedKind = this.extractor.resolveTypeAliasKind?.(node, this.source); - if (resolvedKind === 'struct') { - const structNode = this.createNode('struct', name, node, { docstring, isExported }); - if (!structNode) return true; + if (resolvedKind === 'struct' || resolvedKind === 'union') { + const aggregateNode = this.createNode(resolvedKind, name, node, { docstring, isExported }); + if (!aggregateNode) return true; // Visit body children for field extraction - this.nodeStack.push(structNode.id); - // Try Go-style 'type' field first, then find inner struct child (C typedef struct) + this.nodeStack.push(aggregateNode.id); + // Try Go-style 'type' field first, then find the matching inner aggregate child. const typeChild = getChildByField(node, 'type') - || this.findChildByTypes(node, this.extractor.structTypes); + || this.findChildByTypes( + node, + resolvedKind === 'union' ? (this.extractor.unionTypes ?? []) : this.extractor.structTypes + ); if (typeChild) { // Extract struct embedding (e.g. Go: `type DB struct { *Head; Queryable }`) - this.extractInheritance(typeChild, structNode.id); + this.extractInheritance(typeChild, aggregateNode.id); const body = getChildByField(typeChild, this.extractor.bodyField) || typeChild; for (let i = 0; i < body.namedChildCount; i++) { const child = body.namedChild(i); @@ -5271,6 +5290,10 @@ export class TreeSitterExtractor { this.extractStruct(node); return; } + if (this.extractor!.unionTypes?.includes(nodeType)) { + this.extractUnion(node); + return; + } if (this.extractor!.enumTypes.includes(nodeType)) { this.extractEnum(node); return; @@ -5745,7 +5768,7 @@ export class TreeSitterExtractor { */ private findNodeByName(name: string): string | undefined { for (const node of this.nodes) { - if (node.name === name && (node.kind === 'struct' || node.kind === 'enum' || node.kind === 'class')) { + if (node.name === name && (node.kind === 'struct' || node.kind === 'union' || node.kind === 'enum' || node.kind === 'class')) { return node.id; } } diff --git a/src/types.ts b/src/types.ts index a1861bb..b0ebfe4 100644 --- a/src/types.ts +++ b/src/types.ts @@ -42,6 +42,7 @@ export const NODE_KINDS = [ 'export', 'route', 'component', + 'union', ] as const; export type NodeKind = (typeof NODE_KINDS)[number];