Merge PR #18: Fix arrow function/export indexing + type alias extraction

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

__tests__/extraction.test.ts

@@ -198,6 +198,271 @@ function main() {
   });
 });
 
+describe('Arrow Function Export Extraction', () => {
+  it('should extract exported arrow functions assigned to const', () => {
+    const code = `
+export const useAuth = (): AuthContextValue => {
+  return useContext(AuthContext);
+};
+`;
+    const result = extractFromSource('hooks.ts', code);
+
+    expect(result.nodes).toHaveLength(1);
+    expect(result.nodes[0]).toMatchObject({
+      kind: 'function',
+      name: 'useAuth',
+      isExported: true,
+    });
+  });
+
+  it('should extract exported function expressions assigned to const', () => {
+    const code = `
+export const processData = function(input: string): string {
+  return input.trim();
+};
+`;
+    const result = extractFromSource('utils.ts', code);
+
+    expect(result.nodes).toHaveLength(1);
+    expect(result.nodes[0]).toMatchObject({
+      kind: 'function',
+      name: 'processData',
+      isExported: true,
+    });
+  });
+
+  it('should not extract non-exported arrow functions as exported', () => {
+    const code = `
+const internalHelper = () => {
+  return 42;
+};
+`;
+    const result = extractFromSource('internal.ts', code);
+
+    const helperNode = result.nodes.find((n) => n.name === 'internalHelper');
+    expect(helperNode).toBeDefined();
+    expect(helperNode?.isExported).toBeFalsy();
+  });
+
+  it('should still skip truly anonymous arrow functions', () => {
+    const code = `
+const items = [1, 2, 3].map((x) => x * 2);
+`;
+    const result = extractFromSource('anon.ts', code);
+
+    // The inline arrow function passed to .map() has no variable_declarator parent
+    // and should remain anonymous (skipped)
+    const anonFunctions = result.nodes.filter(
+      (n) => n.kind === 'function' && n.name === '<anonymous>'
+    );
+    expect(anonFunctions).toHaveLength(0);
+  });
+
+  it('should extract multiple exported arrow functions from the same file', () => {
+    const code = `
+export const add = (a: number, b: number): number => a + b;
+
+export const subtract = (a: number, b: number): number => a - b;
+
+const internal = () => 'not exported';
+`;
+    const result = extractFromSource('math.ts', code);
+
+    const exported = result.nodes.filter((n) => n.kind === 'function' && n.isExported);
+    expect(exported).toHaveLength(2);
+    expect(exported.map((n) => n.name).sort()).toEqual(['add', 'subtract']);
+
+    const internalNode = result.nodes.find((n) => n.name === 'internal');
+    expect(internalNode).toBeDefined();
+    expect(internalNode?.isExported).toBeFalsy();
+  });
+
+  it('should extract arrow functions in JavaScript files', () => {
+    const code = `
+export const fetchData = async () => {
+  const response = await fetch('/api/data');
+  return response.json();
+};
+`;
+    const result = extractFromSource('api.js', code);
+
+    expect(result.nodes).toHaveLength(1);
+    expect(result.nodes[0]).toMatchObject({
+      kind: 'function',
+      name: 'fetchData',
+      isExported: true,
+    });
+  });
+});
+
+describe('Type Alias Extraction', () => {
+  it('should extract exported type aliases in TypeScript', () => {
+    const code = `
+export type AuthContextValue = {
+  user: User | null;
+  login: () => void;
+  logout: () => void;
+};
+`;
+    const result = extractFromSource('types.ts', code);
+
+    expect(result.nodes).toHaveLength(1);
+    expect(result.nodes[0]).toMatchObject({
+      kind: 'type_alias',
+      name: 'AuthContextValue',
+      isExported: true,
+    });
+  });
+
+  it('should extract non-exported type aliases', () => {
+    const code = `
+type InternalState = {
+  loading: boolean;
+  error: string | null;
+};
+`;
+    const result = extractFromSource('internal.ts', code);
+
+    expect(result.nodes).toHaveLength(1);
+    expect(result.nodes[0]).toMatchObject({
+      kind: 'type_alias',
+      name: 'InternalState',
+      isExported: false,
+    });
+  });
+
+  it('should extract multiple type aliases from the same file', () => {
+    const code = `
+export type UnitSystem = 'metric' | 'imperial';
+export type DateFormat = 'ISO' | 'US' | 'EU';
+type Internal = string;
+`;
+    const result = extractFromSource('config.ts', code);
+
+    const typeAliases = result.nodes.filter((n) => n.kind === 'type_alias');
+    expect(typeAliases).toHaveLength(3);
+
+    const exported = typeAliases.filter((n) => n.isExported);
+    expect(exported).toHaveLength(2);
+    expect(exported.map((n) => n.name).sort()).toEqual(['DateFormat', 'UnitSystem']);
+  });
+});
+
+describe('Exported Variable Extraction', () => {
+  it('should extract exported const with call expression (Zustand store)', () => {
+    const code = `
+export const useUIStore = create<UIState>((set) => ({
+  isOpen: false,
+  toggle: () => set((s) => ({ isOpen: !s.isOpen })),
+}));
+`;
+    const result = extractFromSource('store.ts', code);
+
+    const varNode = result.nodes.find((n) => n.kind === 'variable' && n.name === 'useUIStore');
+    expect(varNode).toBeDefined();
+    expect(varNode?.isExported).toBe(true);
+  });
+
+  it('should extract exported const with object literal', () => {
+    const code = `
+export const config = {
+  apiUrl: 'https://api.example.com',
+  timeout: 5000,
+};
+`;
+    const result = extractFromSource('config.ts', code);
+
+    const varNode = result.nodes.find((n) => n.kind === 'variable' && n.name === 'config');
+    expect(varNode).toBeDefined();
+    expect(varNode?.isExported).toBe(true);
+  });
+
+  it('should extract exported const with array literal', () => {
+    const code = `
+export const SCREEN_NAMES = ['home', 'settings', 'profile'] as const;
+`;
+    const result = extractFromSource('constants.ts', code);
+
+    const varNode = result.nodes.find((n) => n.kind === 'variable' && n.name === 'SCREEN_NAMES');
+    expect(varNode).toBeDefined();
+    expect(varNode?.isExported).toBe(true);
+  });
+
+  it('should extract exported const with primitive value', () => {
+    const code = `
+export const MAX_RETRIES = 3;
+export const API_VERSION = "v2";
+`;
+    const result = extractFromSource('constants.ts', code);
+
+    const variables = result.nodes.filter((n) => n.kind === 'variable');
+    expect(variables).toHaveLength(2);
+    expect(variables.map((n) => n.name).sort()).toEqual(['API_VERSION', 'MAX_RETRIES']);
+  });
+
+  it('should NOT duplicate arrow functions as both function and variable', () => {
+    const code = `
+export const useAuth = () => {
+  return useContext(AuthContext);
+};
+`;
+    const result = extractFromSource('hooks.ts', code);
+
+    // Should be extracted as function (from arrow function handler), NOT as variable
+    const funcNodes = result.nodes.filter((n) => n.kind === 'function' && n.name === 'useAuth');
+    const varNodes = result.nodes.filter((n) => n.kind === 'variable' && n.name === 'useAuth');
+    expect(funcNodes).toHaveLength(1);
+    expect(varNodes).toHaveLength(0);
+  });
+
+  it('should not extract non-exported const as exported variable', () => {
+    const code = `
+const internalConfig = {
+  debug: true,
+};
+`;
+    const result = extractFromSource('internal.ts', code);
+
+    // Non-exported const should NOT create a variable node
+    // (only export_statement triggers extractExportedVariables)
+    const varNodes = result.nodes.filter((n) => n.kind === 'variable' && n.name === 'internalConfig');
+    expect(varNodes).toHaveLength(0);
+  });
+
+  it('should extract Zod schema exports', () => {
+    const code = `
+export const userSchema = z.object({
+  id: z.string(),
+  name: z.string(),
+  email: z.string().email(),
+});
+`;
+    const result = extractFromSource('schemas.ts', code);
+
+    const varNode = result.nodes.find((n) => n.kind === 'variable' && n.name === 'userSchema');
+    expect(varNode).toBeDefined();
+    expect(varNode?.isExported).toBe(true);
+  });
+
+  it('should extract XState machine exports', () => {
+    const code = `
+export const authMachine = createMachine({
+  id: "auth",
+  initial: "idle",
+  states: {
+    idle: {},
+    authenticated: {},
+  },
+});
+`;
+    const result = extractFromSource('machine.ts', code);
+
+    const varNode = result.nodes.find((n) => n.kind === 'variable' && n.name === 'authMachine');
+    expect(varNode).toBeDefined();
+    expect(varNode?.isExported).toBe(true);
+  });
+});
+
 describe('Python Extraction', () => {
   it('should extract function definitions', () => {
     const code = `

src/extraction/tree-sitter.ts

@@ -104,6 +104,8 @@ interface LanguageExtractor {
   structTypes: string[];
   /** Node types that represent enums */
   enumTypes: string[];
+  /** Node types that represent type aliases (e.g. `type X = ...`) */
+  typeAliasTypes: string[];
   /** Node types that represent imports */
   importTypes: string[];
   /** Node types that represent function calls */
@@ -143,6 +145,7 @@ const EXTRACTORS: Partial<Record<Language, LanguageExtractor>> = {
     interfaceTypes: ['interface_declaration'],
     structTypes: [],
     enumTypes: ['enum_declaration'],
+    typeAliasTypes: ['type_alias_declaration'],
     importTypes: ['import_statement'],
     callTypes: ['call_expression'],
     variableTypes: ['lexical_declaration', 'variable_declaration'],
@@ -172,12 +175,17 @@ const EXTRACTORS: Partial<Record<Language, LanguageExtractor>> = {
       }
       return undefined;
     },
-    isExported: (node, source) => {
-      const parent = node.parent;
-      if (parent?.type === 'export_statement') return true;
-      // Check for 'export' keyword before declaration
-      const text = source.substring(Math.max(0, node.startIndex - 10), node.startIndex);
-      return text.includes('export');
+    isExported: (node, _source) => {
+      // Walk the parent chain to find an export_statement ancestor.
+      // This correctly handles deeply nested nodes like arrow functions
+      // inside variable declarations: `export const X = () => { ... }`
+      // where the arrow_function is 3 levels deep under export_statement.
+      let current = node.parent;
+      while (current) {
+        if (current.type === 'export_statement') return true;
+        current = current.parent;
+      }
+      return false;
     },
     isAsync: (node) => {
       for (let i = 0; i < node.childCount; i++) {
@@ -212,6 +220,7 @@ const EXTRACTORS: Partial<Record<Language, LanguageExtractor>> = {
     interfaceTypes: [],
     structTypes: [],
     enumTypes: [],
+    typeAliasTypes: [],
     importTypes: ['import_statement'],
     callTypes: ['call_expression'],
     variableTypes: ['lexical_declaration', 'variable_declaration'],
@@ -222,11 +231,13 @@ const EXTRACTORS: Partial<Record<Language, LanguageExtractor>> = {
       const params = getChildByField(node, 'parameters');
       return params ? getNodeText(params, source) : undefined;
     },
-    isExported: (node, source) => {
-      const parent = node.parent;
-      if (parent?.type === 'export_statement') return true;
-      const text = source.substring(Math.max(0, node.startIndex - 10), node.startIndex);
-      return text.includes('export');
+    isExported: (node, _source) => {
+      let current = node.parent;
+      while (current) {
+        if (current.type === 'export_statement') return true;
+        current = current.parent;
+      }
+      return false;
     },
     isAsync: (node) => {
       for (let i = 0; i < node.childCount; i++) {
@@ -252,6 +263,7 @@ const EXTRACTORS: Partial<Record<Language, LanguageExtractor>> = {
     interfaceTypes: [],
     structTypes: [],
     enumTypes: [],
+    typeAliasTypes: [],
     importTypes: ['import_statement', 'import_from_statement'],
     callTypes: ['call'],
     variableTypes: ['assignment'], // Python uses assignment for variable declarations
@@ -290,6 +302,7 @@ const EXTRACTORS: Partial<Record<Language, LanguageExtractor>> = {
     interfaceTypes: ['interface_type'],
     structTypes: ['struct_type'],
     enumTypes: [],
+    typeAliasTypes: ['type_spec'], // Go type declarations
     importTypes: ['import_declaration'],
     callTypes: ['call_expression'],
     variableTypes: ['var_declaration', 'short_var_declaration', 'const_declaration'],
@@ -315,6 +328,7 @@ const EXTRACTORS: Partial<Record<Language, LanguageExtractor>> = {
     interfaceTypes: ['trait_item'],
     structTypes: ['struct_item'],
     enumTypes: ['enum_item'],
+    typeAliasTypes: ['type_item'], // Rust type aliases
     importTypes: ['use_declaration'],
     callTypes: ['call_expression'],
     variableTypes: ['let_declaration', 'const_item', 'static_item'],
@@ -356,6 +370,7 @@ const EXTRACTORS: Partial<Record<Language, LanguageExtractor>> = {
     interfaceTypes: ['interface_declaration'],
     structTypes: [],
     enumTypes: ['enum_declaration'],
+    typeAliasTypes: [],
     importTypes: ['import_declaration'],
     callTypes: ['method_invocation'],
     variableTypes: ['local_variable_declaration', 'field_declaration'],
@@ -399,6 +414,7 @@ const EXTRACTORS: Partial<Record<Language, LanguageExtractor>> = {
     interfaceTypes: [],
     structTypes: ['struct_specifier'],
     enumTypes: ['enum_specifier'],
+    typeAliasTypes: ['type_definition'], // typedef
     importTypes: ['preproc_include'],
     callTypes: ['call_expression'],
     variableTypes: ['declaration'],
@@ -413,6 +429,7 @@ const EXTRACTORS: Partial<Record<Language, LanguageExtractor>> = {
     interfaceTypes: [],
     structTypes: ['struct_specifier'],
     enumTypes: ['enum_specifier'],
+    typeAliasTypes: ['type_definition', 'alias_declaration'], // typedef and using
     importTypes: ['preproc_include'],
     callTypes: ['call_expression'],
     variableTypes: ['declaration'],
@@ -443,6 +460,7 @@ const EXTRACTORS: Partial<Record<Language, LanguageExtractor>> = {
     interfaceTypes: ['interface_declaration'],
     structTypes: ['struct_declaration'],
     enumTypes: ['enum_declaration'],
+    typeAliasTypes: [],
     importTypes: ['using_directive'],
     callTypes: ['invocation_expression'],
     variableTypes: ['local_declaration_statement', 'field_declaration'],
@@ -488,6 +506,7 @@ const EXTRACTORS: Partial<Record<Language, LanguageExtractor>> = {
     interfaceTypes: ['interface_declaration'],
     structTypes: [],
     enumTypes: ['enum_declaration'],
+    typeAliasTypes: [],
     importTypes: ['namespace_use_declaration'],
     callTypes: ['function_call_expression', 'member_call_expression', 'scoped_call_expression'],
     variableTypes: ['property_declaration', 'const_declaration'],
@@ -522,6 +541,7 @@ const EXTRACTORS: Partial<Record<Language, LanguageExtractor>> = {
     interfaceTypes: [], // Ruby uses modules
     structTypes: [],
     enumTypes: [],
+    typeAliasTypes: [],
     importTypes: ['call'], // require/require_relative
     callTypes: ['call', 'method_call'],
     variableTypes: ['assignment'], // Ruby uses assignment like Python
@@ -553,6 +573,7 @@ const EXTRACTORS: Partial<Record<Language, LanguageExtractor>> = {
     interfaceTypes: ['protocol_declaration'],
     structTypes: ['struct_declaration'],
     enumTypes: ['enum_declaration'],
+    typeAliasTypes: ['typealias_declaration'],
     importTypes: ['import_declaration'],
     callTypes: ['call_expression'],
     variableTypes: ['property_declaration', 'constant_declaration'],
@@ -613,6 +634,7 @@ const EXTRACTORS: Partial<Record<Language, LanguageExtractor>> = {
     interfaceTypes: ['class_declaration'], // Interfaces use class_declaration with 'interface' modifier
     structTypes: [], // Kotlin uses data classes
     enumTypes: ['class_declaration'], // Enums use class_declaration with 'enum' modifier
+    typeAliasTypes: ['type_alias'],
     importTypes: ['import_header'],
     callTypes: ['call_expression'],
     variableTypes: ['property_declaration'],
@@ -668,6 +690,7 @@ const EXTRACTORS: Partial<Record<Language, LanguageExtractor>> = {
     interfaceTypes: [],
     structTypes: [],
     enumTypes: ['enum_declaration'],
+    typeAliasTypes: ['type_alias'],
     importTypes: ['import_or_export'],
     callTypes: [],  // Dart calls use identifier+selector, handled via function body traversal
     variableTypes: [],
@@ -929,12 +952,22 @@ export class TreeSitterExtractor {
       this.extractEnum(node);
       skipChildren = true; // extractEnum visits body children
     }
+    // Check for type alias declarations (e.g. `type X = ...` in TypeScript)
+    else if (this.extractor.typeAliasTypes.includes(nodeType)) {
+      this.extractTypeAlias(node);
+    }
     // Check for variable declarations (const, let, var, etc.)
     // Only extract top-level variables (not inside functions/methods)
     else if (this.extractor.variableTypes.includes(nodeType) && this.nodeStack.length === 0) {
       this.extractVariable(node);
       skipChildren = true; // extractVariable handles children
     }
+    // Check for export statements containing non-function variable declarations
+    // e.g. `export const X = create(...)`, `export const X = { ... }`
+    else if (nodeType === 'export_statement') {
+      this.extractExportedVariables(node);
+      // Don't skip children — still need to visit inner nodes (functions, calls, etc.)
+    }
     // Check for imports
     else if (this.extractor.importTypes.includes(nodeType)) {
       this.extractImport(node);
@@ -1033,7 +1066,23 @@ export class TreeSitterExtractor {
   private extractFunction(node: SyntaxNode): void {
     if (!this.extractor) return;
 
-    const name = extractName(node, this.source, this.extractor);
+    let name = extractName(node, this.source, this.extractor);
+    // For arrow functions and function expressions assigned to variables,
+    // resolve the name from the parent variable_declarator.
+    // e.g. `export const useAuth = () => { ... }` — the arrow_function node
+    // has no `name` field; the name lives on the variable_declarator.
+    if (
+      name === '<anonymous>' &&
+      (node.type === 'arrow_function' || node.type === 'function_expression')
+    ) {
+      const parent = node.parent;
+      if (parent?.type === 'variable_declarator') {
+        const varName = getChildByField(parent, 'name');
+        if (varName) {
+          name = getNodeText(varName, this.source);
+        }
+      }
+    }
     if (name === '<anonymous>') return; // Skip anonymous functions
 
     const docstring = getPrecedingDocstring(node, this.source);
@@ -1345,6 +1394,77 @@ export class TreeSitterExtractor {
     }
   }
 
+  /**
+   * Extract a type alias (e.g. `export type X = ...` in TypeScript)
+   */
+  private extractTypeAlias(node: SyntaxNode): void {
+    if (!this.extractor) return;
+
+    const name = extractName(node, this.source, this.extractor);
+    if (name === '<anonymous>') return;
+    const docstring = getPrecedingDocstring(node, this.source);
+    const isExported = this.extractor.isExported?.(node, this.source);
+
+    this.createNode('type_alias', name, node, {
+      docstring,
+      isExported,
+    });
+  }
+
+  /**
+   * Extract an exported variable declaration that isn't a function.
+   * Handles patterns like:
+   *   export const X = create(...)
+   *   export const X = { ... }
+   *   export const X = [...]
+   *   export const X = "value"
+   *
+   * This is called for `export_statement` nodes that contain a
+   * `lexical_declaration` with `variable_declarator` children whose
+   * values are NOT already handled by functionTypes (arrow_function,
+   * function_expression).
+   */
+  private extractExportedVariables(exportNode: SyntaxNode): void {
+    if (!this.extractor) return;
+
+    // Find the lexical_declaration or variable_declaration child
+    for (let i = 0; i < exportNode.namedChildCount; i++) {
+      const decl = exportNode.namedChild(i);
+      if (!decl || (decl.type !== 'lexical_declaration' && decl.type !== 'variable_declaration')) {
+        continue;
+      }
+
+      // Iterate over each variable_declarator in the declaration
+      for (let j = 0; j < decl.namedChildCount; j++) {
+        const declarator = decl.namedChild(j);
+        if (!declarator || declarator.type !== 'variable_declarator') continue;
+
+        const nameNode = getChildByField(declarator, 'name');
+        if (!nameNode) continue;
+        const name = getNodeText(nameNode, this.source);
+
+        // Skip if the value is a function type — those are already handled
+        // by extractFunction via the functionTypes dispatch
+        const value = getChildByField(declarator, 'value');
+        if (value) {
+          const valueType = value.type;
+          if (
+            this.extractor.functionTypes.includes(valueType)
+          ) {
+            continue; // Already handled by extractFunction
+          }
+        }
+
+        const docstring = getPrecedingDocstring(exportNode, this.source);
+
+        this.createNode('variable', name, declarator, {
+          docstring,
+          isExported: true,
+        });
+      }
+    }
+  }
+
   /**
    * Extract an import
    *