diff --git a/packages/core/src/map/__tests__/map.test.ts b/packages/core/src/map/__tests__/map.test.ts index 875f78d..d8a99c3 100644 --- a/packages/core/src/map/__tests__/map.test.ts +++ b/packages/core/src/map/__tests__/map.test.ts @@ -69,6 +69,7 @@ describe('Codebase Map', () => { path: 'packages/cli/src/cli.ts', type: 'function', name: 'main', + signature: 'function main(args: string[]): Promise', startLine: 5, endLine: 50, language: 'typescript', @@ -167,6 +168,30 @@ describe('Codebase Map', () => { expect(nodeWithExports?.exports?.[0].name).toBeDefined(); }); + it('should include signatures in exports when available', async () => { + const indexer = createMockIndexer(); + const map = await generateCodebaseMap(indexer, { depth: 5, includeExports: true }); + + // Find any node with an export that has a signature + const findExportWithSignature = ( + node: typeof map.root + ): { name: string; signature?: string } | null => { + if (node.exports) { + const withSig = node.exports.find((e) => e.signature); + if (withSig) return withSig; + } + for (const child of node.children) { + const found = findExportWithSignature(child); + if (found) return found; + } + return null; + }; + + const exportWithSig = findExportWithSignature(map.root); + expect(exportWithSig).not.toBeNull(); + expect(exportWithSig?.signature).toBe('function main(args: string[]): Promise'); + }); + it('should not include exports when includeExports is false', async () => { const indexer = createMockIndexer(); const map = await generateCodebaseMap(indexer, { depth: 5, includeExports: false }); @@ -259,6 +284,41 @@ describe('Codebase Map', () => { expect(output).toContain('exports:'); }); + it('should show signatures in exports when available', async () => { + const indexer = createMockIndexer(); + const map = await generateCodebaseMap(indexer, { depth: 5, includeExports: true }); + const output = formatCodebaseMap(map, { includeExports: true }); + + // The main function has a signature, should appear in output + expect(output).toContain('function main(args: string[]): Promise'); + }); + + it('should truncate long signatures', async () => { + const longSigResults: SearchResult[] = [ + { + id: 'src/index.ts:longFunction:1', + score: 0.9, + metadata: { + path: 'src/index.ts', + type: 'function', + name: 'longFunction', + signature: + 'function longFunction(param1: string, param2: number, param3: boolean, param4: object): Promise', + exported: true, + }, + }, + ]; + + const indexer = createMockIndexer(longSigResults); + const map = await generateCodebaseMap(indexer, { depth: 5, includeExports: true }); + const output = formatCodebaseMap(map, { includeExports: true }); + + // Should be truncated with ... + expect(output).toContain('...'); + // Should not contain the full signature + expect(output).not.toContain('ComplexReturnType'); + }); + it('should show component counts', async () => { const indexer = createMockIndexer(); const map = await generateCodebaseMap(indexer); @@ -277,6 +337,277 @@ describe('Codebase Map', () => { }); }); + describe('Hot Paths', () => { + it('should compute hot paths from callers data', async () => { + const resultsWithCallers: SearchResult[] = [ + { + id: 'src/core.ts:coreFunction:1', + score: 0.9, + metadata: { + path: 'src/core.ts', + type: 'function', + name: 'coreFunction', + exported: true, + callers: [ + { name: 'caller1', file: 'src/a.ts', startLine: 10 }, + { name: 'caller2', file: 'src/b.ts', startLine: 20 }, + { name: 'caller3', file: 'src/c.ts', startLine: 30 }, + ], + }, + }, + { + id: 'src/utils.ts:utilFunction:1', + score: 0.8, + metadata: { + path: 'src/utils.ts', + type: 'function', + name: 'utilFunction', + exported: true, + callers: [{ name: 'caller1', file: 'src/a.ts', startLine: 15 }], + }, + }, + ]; + + const indexer = createMockIndexer(resultsWithCallers); + const map = await generateCodebaseMap(indexer, { includeHotPaths: true }); + + expect(map.hotPaths.length).toBeGreaterThan(0); + // coreFunction should be first (more callers) + expect(map.hotPaths[0].file).toBe('src/core.ts'); + expect(map.hotPaths[0].incomingRefs).toBe(3); + }); + + it('should compute hot paths from callees data', async () => { + const resultsWithCallees: SearchResult[] = [ + { + id: 'src/main.ts:main:1', + score: 0.9, + metadata: { + path: 'src/main.ts', + type: 'function', + name: 'main', + exported: true, + callees: [ + { name: 'helper', file: 'src/helpers.ts', line: 10 }, + { name: 'helper', file: 'src/helpers.ts', line: 10 }, + ], + }, + }, + { + id: 'src/other.ts:other:1', + score: 0.8, + metadata: { + path: 'src/other.ts', + type: 'function', + name: 'other', + exported: true, + callees: [{ name: 'helper', file: 'src/helpers.ts', line: 10 }], + }, + }, + ]; + + const indexer = createMockIndexer(resultsWithCallees); + const map = await generateCodebaseMap(indexer, { includeHotPaths: true }); + + expect(map.hotPaths.length).toBeGreaterThan(0); + // helpers.ts should be referenced most + expect(map.hotPaths[0].file).toBe('src/helpers.ts'); + expect(map.hotPaths[0].incomingRefs).toBe(3); + }); + + it('should limit hot paths to maxHotPaths', async () => { + const manyRefs: SearchResult[] = Array.from({ length: 20 }, (_, i) => ({ + id: `src/file${i}.ts:fn:1`, + score: 0.9, + metadata: { + path: `src/file${i}.ts`, + type: 'function', + name: `fn${i}`, + exported: true, + callers: Array.from({ length: 20 - i }, (_, j) => ({ + name: `caller${j}`, + file: `src/other${j}.ts`, + startLine: j * 10, + })), + }, + })); + + const indexer = createMockIndexer(manyRefs); + const map = await generateCodebaseMap(indexer, { includeHotPaths: true, maxHotPaths: 3 }); + + expect(map.hotPaths.length).toBe(3); + // Should be sorted by refs descending + expect(map.hotPaths[0].incomingRefs).toBeGreaterThanOrEqual(map.hotPaths[1].incomingRefs); + }); + + it('should not include hot paths when disabled', async () => { + const resultsWithCallers: SearchResult[] = [ + { + id: 'src/core.ts:coreFunction:1', + score: 0.9, + metadata: { + path: 'src/core.ts', + type: 'function', + name: 'coreFunction', + exported: true, + callers: [{ name: 'caller1', file: 'src/a.ts', startLine: 10 }], + }, + }, + ]; + + const indexer = createMockIndexer(resultsWithCallers); + const map = await generateCodebaseMap(indexer, { includeHotPaths: false }); + + expect(map.hotPaths.length).toBe(0); + }); + + it('should format hot paths in output', async () => { + const resultsWithCallers: SearchResult[] = [ + { + id: 'src/core.ts:coreFunction:1', + score: 0.9, + metadata: { + path: 'src/core.ts', + type: 'function', + name: 'coreFunction', + exported: true, + callers: [ + { name: 'caller1', file: 'src/a.ts', startLine: 10 }, + { name: 'caller2', file: 'src/b.ts', startLine: 20 }, + ], + }, + }, + ]; + + const indexer = createMockIndexer(resultsWithCallers); + const map = await generateCodebaseMap(indexer, { includeHotPaths: true }); + const output = formatCodebaseMap(map, { includeHotPaths: true }); + + expect(output).toContain('## Hot Paths'); + expect(output).toContain('src/core.ts'); + expect(output).toContain('2 refs'); + }); + }); + + describe('Smart Depth', () => { + it('should expand dense directories when smartDepth is enabled', async () => { + // Create a structure with varying density + const mixedDensity: SearchResult[] = [ + // Dense directory - 15 components + ...Array.from({ length: 15 }, (_, i) => ({ + id: `packages/core/src/dense/file${i}.ts:fn:1`, + score: 0.9, + metadata: { + path: `packages/core/src/dense/file${i}.ts`, + type: 'function', + name: `fn${i}`, + exported: true, + }, + })), + // Sparse directory - 2 components + ...Array.from({ length: 2 }, (_, i) => ({ + id: `packages/core/src/sparse/file${i}.ts:fn:1`, + score: 0.9, + metadata: { + path: `packages/core/src/sparse/file${i}.ts`, + type: 'function', + name: `fn${i}`, + exported: true, + }, + })), + ]; + + const indexer = createMockIndexer(mixedDensity); + const map = await generateCodebaseMap(indexer, { + depth: 5, + smartDepth: true, + smartDepthThreshold: 10, + }); + + // Find the core node + const findNode = (node: typeof map.root, name: string): typeof map.root | null => { + if (node.name === name) return node; + for (const child of node.children) { + const found = findNode(child, name); + if (found) return found; + } + return null; + }; + + const srcNode = findNode(map.root, 'src'); + expect(srcNode).not.toBeNull(); + + // Dense should be expanded (has children or is at leaf level) + const denseNode = srcNode?.children.find((c) => c.name === 'dense'); + expect(denseNode).toBeDefined(); + expect(denseNode?.componentCount).toBe(15); + + // Sparse should also exist but may be collapsed + const sparseNode = srcNode?.children.find((c) => c.name === 'sparse'); + expect(sparseNode).toBeDefined(); + expect(sparseNode?.componentCount).toBe(2); + }); + + it('should always expand first 2 levels regardless of density', async () => { + const sparseResults: SearchResult[] = [ + { + id: 'packages/tiny/src/file.ts:fn:1', + score: 0.9, + metadata: { + path: 'packages/tiny/src/file.ts', + type: 'function', + name: 'fn', + exported: true, + }, + }, + ]; + + const indexer = createMockIndexer(sparseResults); + const map = await generateCodebaseMap(indexer, { + depth: 5, + smartDepth: true, + smartDepthThreshold: 100, // Very high threshold + }); + + // Should still show packages and tiny (first 2 levels) + const packagesNode = map.root.children.find((c) => c.name === 'packages'); + expect(packagesNode).toBeDefined(); + expect(packagesNode?.children.length).toBeGreaterThan(0); + }); + + it('should not use smart depth when disabled', async () => { + const results: SearchResult[] = Array.from({ length: 5 }, (_, i) => ({ + id: `a/b/c/d/e/file${i}.ts:fn:1`, + score: 0.9, + metadata: { + path: `a/b/c/d/e/file${i}.ts`, + type: 'function', + name: `fn${i}`, + exported: true, + }, + })); + + const indexer = createMockIndexer(results); + const mapWithSmart = await generateCodebaseMap(indexer, { + depth: 3, + smartDepth: true, + smartDepthThreshold: 1, + }); + const mapWithoutSmart = await generateCodebaseMap(indexer, { + depth: 3, + smartDepth: false, + }); + + // Without smart depth, should strictly follow depth limit + const countDepth = (node: typeof mapWithSmart.root, d = 0): number => { + if (node.children.length === 0) return d; + return Math.max(...node.children.map((c) => countDepth(c, d + 1))); + }; + + expect(countDepth(mapWithoutSmart.root)).toBeLessThanOrEqual(3); + }); + }); + describe('Edge Cases', () => { it('should handle empty results', async () => { const indexer = createMockIndexer([]); diff --git a/packages/core/src/map/index.ts b/packages/core/src/map/index.ts index 03e5dd9..409b76b 100644 --- a/packages/core/src/map/index.ts +++ b/packages/core/src/map/index.ts @@ -6,7 +6,7 @@ import * as path from 'node:path'; import type { RepositoryIndexer } from '../indexer'; import type { SearchResult } from '../vector/types'; -import type { CodebaseMap, ExportInfo, MapNode, MapOptions } from './types'; +import type { CodebaseMap, ExportInfo, HotPath, MapNode, MapOptions } from './types'; export * from './types'; @@ -16,6 +16,10 @@ const DEFAULT_OPTIONS: Required = { focus: '', includeExports: true, maxExportsPerDir: 5, + includeHotPaths: true, + maxHotPaths: 5, + smartDepth: false, + smartDepthThreshold: 10, tokenBudget: 2000, }; @@ -46,10 +50,14 @@ export async function generateCodebaseMap( const totalComponents = countComponents(root); const totalDirectories = countDirectories(root); + // Compute hot paths (most referenced files) + const hotPaths = opts.includeHotPaths ? computeHotPaths(allDocs, opts.maxHotPaths) : []; + return { root, totalComponents, totalDirectories, + hotPaths, generatedAt: new Date().toISOString(), }; } @@ -92,8 +100,12 @@ function buildDirectoryTree(docs: SearchResult[], opts: Required): M insertIntoTree(root, dir, dirDocs, opts); } - // Prune tree to depth - pruneToDepth(root, opts.depth); + // Prune tree to depth (smart or fixed) + if (opts.smartDepth) { + smartPruneTree(root, opts.depth, opts.smartDepthThreshold); + } else { + pruneToDepth(root, opts.depth); + } // Sort children alphabetically sortTree(root); @@ -161,6 +173,7 @@ function extractExports(docs: SearchResult[], maxExports: number): ExportInfo[] name: doc.metadata.name as string, type: (doc.metadata.type as string) || 'unknown', file: (doc.metadata.path as string) || (doc.metadata.file as string) || '', + signature: doc.metadata.signature as string | undefined, }); if (exports.length >= maxExports) break; @@ -199,6 +212,38 @@ function pruneToDepth(node: MapNode, depth: number, currentDepth = 0): void { } } +/** + * Smart prune tree - expand dense directories, collapse sparse ones + * Uses information density heuristic: expand if componentCount >= threshold + */ +function smartPruneTree( + node: MapNode, + maxDepth: number, + threshold: number, + currentDepth = 0 +): void { + // Always stop at max depth + if (currentDepth >= maxDepth) { + node.children = []; + return; + } + + // For each child, decide whether to expand or collapse + for (const child of node.children) { + // Expand if: + // 1. We're within first 2 levels (always show some structure) + // 2. OR the child has enough components to be "interesting" + const shouldExpand = currentDepth < 2 || child.componentCount >= threshold; + + if (shouldExpand) { + smartPruneTree(child, maxDepth, threshold, currentDepth + 1); + } else { + // Collapse this branch - it's too sparse to be interesting + child.children = []; + } + } +} + /** * Sort tree children alphabetically */ @@ -227,6 +272,54 @@ function countDirectories(node: MapNode): number { return count; } +/** + * Compute hot paths - files with the most incoming references + */ +function computeHotPaths(docs: SearchResult[], maxPaths: number): HotPath[] { + // Count incoming references per file + const refCounts = new Map(); + + for (const doc of docs) { + const callers = doc.metadata.callers as Array<{ file: string }> | undefined; + if (callers && Array.isArray(callers)) { + // This document is called by others - count it + const filePath = (doc.metadata.path as string) || (doc.metadata.file as string) || ''; + if (filePath) { + const existing = refCounts.get(filePath) || { count: 0 }; + existing.count += callers.length; + existing.component = existing.component || (doc.metadata.name as string); + refCounts.set(filePath, existing); + } + } + } + + // Also count based on callees pointing to files + for (const doc of docs) { + const callees = doc.metadata.callees as Array<{ file: string; name: string }> | undefined; + if (callees && Array.isArray(callees)) { + for (const callee of callees) { + if (callee.file) { + const existing = refCounts.get(callee.file) || { count: 0 }; + existing.count += 1; + refCounts.set(callee.file, existing); + } + } + } + } + + // Sort by count and take top N + const sorted = Array.from(refCounts.entries()) + .map(([file, data]) => ({ + file, + incomingRefs: data.count, + primaryComponent: data.component, + })) + .sort((a, b) => b.incomingRefs - a.incomingRefs) + .slice(0, maxPaths); + + return sorted; +} + /** * Format codebase map as readable text */ @@ -237,7 +330,20 @@ export function formatCodebaseMap(map: CodebaseMap, options: MapOptions = {}): s lines.push('# Codebase Map'); lines.push(''); + // Format hot paths if present + if (opts.includeHotPaths && map.hotPaths.length > 0) { + lines.push('## Hot Paths (most referenced)'); + for (let i = 0; i < map.hotPaths.length; i++) { + const hp = map.hotPaths[i]; + const component = hp.primaryComponent ? ` (${hp.primaryComponent})` : ''; + lines.push(`${i + 1}. \`${hp.file}\`${component} - ${hp.incomingRefs} refs`); + } + lines.push(''); + } + // Format tree + lines.push('## Directory Structure'); + lines.push(''); formatNode(map.root, lines, '', true, opts); lines.push(''); @@ -266,8 +372,16 @@ function formatNode( // Add exports if present if (opts.includeExports && node.exports && node.exports.length > 0) { const exportPrefix = prefix + (isLast ? ' ' : '│ '); - const exportNames = node.exports.map((e) => e.name).join(', '); - lines.push(`${exportPrefix}└── exports: ${exportNames}`); + const exportItems = node.exports.map((e) => { + // Use signature if available, otherwise just name + if (e.signature) { + // Truncate long signatures + const sig = e.signature.length > 60 ? `${e.signature.slice(0, 57)}...` : e.signature; + return sig; + } + return e.name; + }); + lines.push(`${exportPrefix}└── exports: ${exportItems.join(', ')}`); } // Format children diff --git a/packages/core/src/map/types.ts b/packages/core/src/map/types.ts index ad21d29..7e0d507 100644 --- a/packages/core/src/map/types.ts +++ b/packages/core/src/map/types.ts @@ -31,6 +31,8 @@ export interface ExportInfo { type: string; /** File where it's defined */ file: string; + /** Function/method signature (if available) */ + signature?: string; } /** @@ -45,10 +47,30 @@ export interface MapOptions { includeExports?: boolean; /** Maximum exports to show per directory (default: 5) */ maxExportsPerDir?: number; + /** Include hot paths - most referenced files (default: true) */ + includeHotPaths?: boolean; + /** Maximum hot paths to show (default: 5) */ + maxHotPaths?: number; + /** Use smart depth - expand dense directories, collapse sparse ones (default: false) */ + smartDepth?: boolean; + /** Minimum components to expand a directory when using smart depth (default: 10) */ + smartDepthThreshold?: number; /** Token budget for output (default: 2000) */ tokenBudget?: number; } +/** + * Information about a frequently referenced file + */ +export interface HotPath { + /** File path */ + file: string; + /** Number of incoming references (callers) */ + incomingRefs: number; + /** Primary component name in this file */ + primaryComponent?: string; +} + /** * Result of codebase map generation */ @@ -59,6 +81,8 @@ export interface CodebaseMap { totalComponents: number; /** Total number of directories */ totalDirectories: number; + /** Most referenced files (hot paths) */ + hotPaths: HotPath[]; /** Generation timestamp */ generatedAt: string; }