Implement all stubs: parsePackageMap, LLM heuristics, tree-sitter AST, patch rewrite, reinit

- format.ts: Full parsePackageMap implementation with round-trip support - llm-extract.ts: Heuristic-based extraction without LLM API (regex exports, imports, purpose inference from filename patterns) - ast-extract.ts: Tree-sitter integration for TypeScript/TSX with fallback for other languages - init.ts: Refactored to expose generateDirectoryMap helper, reinitPath delegates to initProject - patch.ts: Small packages now use generateDirectoryMap for full rewrite - discover.ts: Cleaned up require('fs') to use proper ES import - Tests: 8 tests covering format round-trip and LLM heuristics - All TypeScript compiles clean, all tests pass
2026-06-09 18:04:57 +02:00
parent fd958671ca
commit d311c8fb3c
14 changed files with 1032 additions and 305 deletions
@@ -1,13 +1,212 @@
-import { readFileSync } from 'fs';
+import { readFileSync } from "fs";
+import { extname } from "path";

 interface ASTFileData {
 	exports: string[];
 	deps: string[];
 }

-export async function extractFileAST(filePath: string): Promise<ASTFileData | null> {
-  // TODO: integrate tree-sitter
-  // Detect language from extension
-  // Parse and extract symbols
+const LANGUAGE_MAP: Record<string, string> = {
+	".ts": "typescript",
+	".tsx": "tsx",
+	".js": "javascript",
+	".jsx": "javascript",
+	".mjs": "javascript",
+	".py": "python",
+	".go": "go",
+	".rs": "rust",
+	".java": "java",
+	".c": "c",
+	".cpp": "cpp",
+	".h": "c",
+	".rb": "ruby",
+};
+
+export async function extractFileAST(
+	filePath: string,
+): Promise<ASTFileData | null> {
+	const ext = extname(filePath).toLowerCase();
+	const langName = LANGUAGE_MAP[ext];
+	if (!langName) return null;
+
+	try {
+		const Parser = require("tree-sitter");
+		let grammar: unknown;
+
+		if (langName === "typescript" || langName === "tsx") {
+			const ts = require("tree-sitter-typescript");
+			grammar = langName === "tsx" ? ts.tsx : ts.typescript;
+		} else {
+			// For other languages, try to require the grammar package
+			try {
+				const pkg = require(`tree-sitter-${langName}`);
+				grammar = pkg.default || pkg;
+			} catch {
 				return null;
 			}
+		}
+
+		if (!grammar) return null;
+
+		const parser = new Parser();
+		parser.setLanguage(grammar);
+
+		const content = readFileSync(filePath, "utf8");
+		const tree = parser.parse(content);
+
+		const exports = extractExportsFromTree(tree, langName);
+		const deps = extractDepsFromTree(tree, langName);
+
+		return { exports, deps };
+	} catch {
+		// Graceful fallback if tree-sitter fails
+		return null;
+	}
+}
+
+function extractExportsFromTree(tree: Tree, langName: string): string[] {
+	const exports: string[] = [];
+	const root = tree.rootNode;
+
+	function visit(node: SyntaxNode) {
+		if (langName === "typescript" || langName === "tsx" || langName === "javascript") {
+			if (node.type === "export_statement") {
+				// export function foo
+				// export class Foo
+				// export const foo
+				// export { foo, bar }
+				// export default foo
+				const declaration = node.childForFieldName?.("declaration");
+				if (declaration) {
+					const nameNode = findIdentifier(declaration);
+					if (nameNode) exports.push(nameNode.text);
+				} else {
+					// export { ... }
+					for (let i = 0; i < node.childCount; i++) {
+						const child = node.child(i);
+						if (child?.type === "export_clause") {
+							for (let j = 0; j < child.childCount; j++) {
+								const spec = child.child(j);
+								if (spec?.type === "export_specifier") {
+									const nameNode = spec.childForFieldName?.("name");
+									if (nameNode) exports.push(nameNode.text);
+								}
+							}
+						}
+					}
+				}
+			}
+		} else if (langName === "python") {
+			if (node.type === "function_definition" || node.type === "class_definition") {
+				const nameNode = node.childForFieldName?.("name");
+				if (nameNode) exports.push(nameNode.text);
+			}
+		} else if (langName === "go") {
+			if (
+				node.type === "function_declaration" ||
+				node.type === "type_declaration" ||
+				node.type === "var_declaration" ||
+				node.type === "const_declaration"
+			) {
+				const nameNode = findIdentifier(node);
+				if (nameNode && /^[A-Z]/.test(nameNode.text)) {
+					exports.push(nameNode.text);
+				}
+			}
+		} else if (langName === "rust") {
+			if (node.type === "function_item" || node.type === "struct_item" || node.type === "enum_item") {
+				const nameNode = node.childForFieldName?.("name");
+				if (nameNode) exports.push(nameNode.text);
+			}
+		}
+
+		for (let i = 0; i < node.childCount; i++) {
+			const child = node.child(i);
+			if (child) visit(child);
+		}
+	}
+
+	visit(root);
+	return [...new Set(exports)];
+}
+
+function extractDepsFromTree(tree: Tree, langName: string): string[] {
+	const deps: string[] = [];
+	const root = tree.rootNode;
+
+	function visit(node: SyntaxNode) {
+		if (langName === "typescript" || langName === "tsx" || langName === "javascript") {
+			if (node.type === "import_statement" || node.type === "import_declaration") {
+				const source = node.childForFieldName?.("source");
+				if (source) {
+					const text = source.text;
+					// Remove quotes
+					deps.push(text.slice(1, -1));
+				}
+			}
+			// CommonJS: require("...")
+			if (node.type === "call_expression") {
+				const func = node.childForFieldName?.("function");
+				if (func?.text === "require") {
+					const args = node.childForFieldName?.("arguments");
+					if (args && args.childCount > 0) {
+						const firstArg = args.child(0);
+						if (firstArg?.type === "string") {
+							deps.push(firstArg.text.slice(1, -1));
+						}
+					}
+				}
+			}
+		} else if (langName === "python") {
+			if (node.type === "import_statement" || node.type === "import_from_statement") {
+				const nameNode = node.childForFieldName?.("name");
+				if (nameNode) deps.push(nameNode.text);
+			}
+		} else if (langName === "go") {
+			if (node.type === "import_spec") {
+				const pathNode = node.childForFieldName?.("path");
+				if (pathNode) deps.push(pathNode.text.slice(1, -1));
+			}
+		} else if (langName === "rust") {
+			if (node.type === "use_declaration") {
+				const argument = node.childForFieldName?.("argument");
+				if (argument) deps.push(argument.text);
+			}
+		}
+
+		for (let i = 0; i < node.childCount; i++) {
+			const child = node.child(i);
+			if (child) visit(child);
+		}
+	}
+
+	visit(root);
+	return [...new Set(deps)];
+}
+
+function findIdentifier(node: SyntaxNode): SyntaxNode | null {
+	if (node.type === "identifier" || node.type === "type_identifier" || node.type === "property_identifier") {
+		return node;
+	}
+	for (let i = 0; i < node.childCount; i++) {
+		const child = node.child(i);
+		if (child) {
+			const found = findIdentifier(child);
+			if (found) return found;
+		}
+	}
+	return null;
+}
+
+// Type definitions for tree-sitter nodes (simplified)
+interface Tree {
+	rootNode: SyntaxNode;
+}
+
+interface SyntaxNode {
+	type: string;
+	text: string;
+	childCount: number;
+	childForFieldName?(name: string): SyntaxNode | null;
+	child(index: number): SyntaxNode | null;
+}
@@ -1,27 +1,27 @@
 #!/usr/bin/env node
-import { initProject } from './init.js';
-import { patchFile } from './patch.js';
-import { validateMaps } from './validate.js';
-import { reinitPath } from './init.js';
+import { initProject } from "./init.js";
+import { patchFile } from "./patch.js";
+import { validateMaps } from "./validate.js";
+import { reinitPath } from "./init.js";

 const args = process.argv.slice(2);
 const command = args[0];

 async function main() {
 	switch (command) {
-    case 'init':
-      await initProject(args[1] || '.');
+		case "init":
+			await initProject(args[1] || ".");
 			break;
-    case 'patch':
+		case "patch":
 			await patchFile(args[1]);
 			break;
-    case 'validate': {
-      const result = await validateMaps(args[1] || '.');
+		case "validate": {
+			const result = await validateMaps(args[1] || ".");
 			process.exit(result.clean ? 0 : 1);
 			break;
 		}
-    case 'reinit':
-      await reinitPath(args[1] || '.');
+		case "reinit":
+			await reinitPath(args[1] || ".");
 			break;
 		default:
 			console.log(`Usage: project-map <init|patch|validate|reinit> [path]`);
@@ -8,13 +8,21 @@ export interface SkillConfig {

 export const DEFAULT_CONFIG: SkillConfig = {
 	ignorePatterns: [
-    'node_modules', '.git', 'dist', 'build', 'coverage',
-    '.next', '.venv', '__pycache__', '.DS_Store', '*.log'
+		"node_modules",
+		".git",
+		"dist",
+		"build",
+		"coverage",
+		".next",
+		".venv",
+		"__pycache__",
+		".DS_Store",
+		"*.log",
 	],
 	smallPackageThreshold: 10,
-  llmModel: 'gpt-4o-mini',
+	llmModel: "gpt-4o-mini",
 	contextBudget: 4000,
-  autoInjectPrompt: true
+	autoInjectPrompt: true,
 };

 export function loadConfig(): SkillConfig {
@@ -1,10 +1,18 @@
-import { readdirSync, statSync } from 'fs';
-import { join, relative } from 'path';
-import ignore from 'ignore';
+import { readdirSync, statSync, readFileSync } from "fs";
+import { join, relative } from "path";
+import ignore from "ignore";

 const DEFAULT_IGNORE = [
-  'node_modules', '.git', 'dist', 'build', 'coverage',
-  '.next', '.venv', '__pycache__', '.DS_Store', '*.log'
+	"node_modules",
+	".git",
+	"dist",
+	"build",
+	"coverage",
+	".next",
+	".venv",
+	"__pycache__",
+	".DS_Store",
+	"*.log",
 ];

 export interface DirectoryEntry {
@@ -15,16 +23,18 @@ export interface DirectoryEntry {

 export function discoverProject(rootPath: string): DirectoryEntry[] {
 	const ig = ignore().add(DEFAULT_IGNORE);
-  const gitignorePath = join(rootPath, '.gitignore');
+	const gitignorePath = join(rootPath, ".gitignore");
 	try {
-    const gitignoreContent = require('fs').readFileSync(gitignorePath, 'utf8');
+		const gitignoreContent = readFileSync(gitignorePath, "utf8");
 		ig.add(gitignoreContent);
-  } catch { /* no .gitignore */ }
+	} catch {
+		/* no .gitignore */
+	}

 	const entries: DirectoryEntry[] = [];

 	function walk(dir: string) {
-    const relDir = relative(rootPath, dir) || '.';
+		const relDir = relative(rootPath, dir) || ".";
 		if (ig.ignores(relDir)) return;

 		const items = readdirSync(dir);
@@ -20,21 +20,110 @@ export function renderPackageMap(data: PackageMapData): string {
 	lines.push(data.role);
 	lines.push(`## files`);
 	for (const file of data.files) {
-    const exp = file.exports.length > 0 ? `exp: ${file.exports.join(', ')}` : '';
-    const dep = file.deps.length > 0 ? `dep: ${file.deps.join(', ')}` : '';
+		const exp =
+			file.exports.length > 0 ? `exp: ${file.exports.join(", ")}` : "";
+		const dep = file.deps.length > 0 ? `dep: ${file.deps.join(", ")}` : "";
 		const parts = [`- ${file.name} | ${file.purpose}`];
 		if (exp) parts.push(exp);
 		if (dep) parts.push(dep);
-    lines.push(parts.join(' | '));
+		lines.push(parts.join(" | "));
 	}
 	lines.push(`## arch`);
 	lines.push(data.arch);
 	lines.push(`## dirty`);
-  lines.push(data.dirty || '-');
-  return `${lines.join('\n')}\n`;
+	lines.push(data.dirty || "-");
+	return `${lines.join("\n")}\n`;
 }

-export function parsePackageMap(_markdown: string): PackageMapData {
-  // TODO: implement robust parser
-  throw new Error('parsePackageMap not yet implemented');
+export function parsePackageMap(markdown: string): PackageMapData {
+	const lines = markdown.split("\n").map((l) => l.trimEnd());
+	const result: PackageMapData = {
+		path: "",
+		role: "",
+		files: [],
+		arch: "",
+		dirty: "-",
+	};
+
+	let section: "none" | "role" | "files" | "arch" | "dirty" = "none";
+
+	for (const line of lines) {
+		if (line.startsWith("# ")) {
+			result.path = line.slice(2).trim();
+			continue;
+		}
+		if (line === "## role") {
+			section = "role";
+			continue;
+		}
+		if (line === "## files") {
+			section = "files";
+			continue;
+		}
+		if (line === "## arch") {
+			section = "arch";
+			continue;
+		}
+		if (line === "## dirty") {
+			section = "dirty";
+			continue;
+		}
+		if (line === "") continue;
+
+		switch (section) {
+			case "role":
+				result.role = line;
+				break;
+			case "files": {
+				if (!line.startsWith("- ")) continue;
+				const entry = parseFileLine(line);
+				if (entry) result.files.push(entry);
+				break;
+			}
+			case "arch":
+				result.arch = result.arch ? `${result.arch}\n${line}` : line;
+				break;
+			case "dirty":
+				result.dirty = line === "-" ? undefined : line;
+				break;
+		}
+	}
+
+	return result;
+}
+
+function parseFileLine(line: string): FileEntry | null {
+	// Format: - filename | purpose | exp: ... | dep: ...
+	const withoutPrefix = line.slice(2).trim();
+	const parts = withoutPrefix.split(" | ").map((p) => p.trim());
+
+	if (parts.length < 2) return null;
+
+	const name = parts[0];
+	const purpose = parts[1];
+	const exports: string[] = [];
+	const deps: string[] = [];
+
+	for (let i = 2; i < parts.length; i++) {
+		const part = parts[i];
+		if (part.startsWith("exp: ")) {
+			exports.push(
+				...part
+					.slice(5)
+					.split(",")
+					.map((s) => s.trim())
+					.filter(Boolean),
+			);
+		} else if (part.startsWith("dep: ")) {
+			deps.push(
+				...part
+					.slice(5)
+					.split(",")
+					.map((s) => s.trim())
+					.filter(Boolean),
+			);
+		}
+	}
+
+	return { name, purpose, exports, deps };
 }
@@ -1,5 +1,5 @@
 // Main entry point for pi-project-map skill
-export { initProject } from './init.js';
-export { patchFile } from './patch.js';
-export { validateMaps } from './validate.js';
-export { renderPackageMap, parsePackageMap } from './format.js';
+export { initProject } from "./init.js";
+export { patchFile } from "./patch.js";
+export { validateMaps } from "./validate.js";
+export { renderPackageMap, parsePackageMap } from "./format.js";
@@ -1,17 +1,25 @@
-import { discoverProject } from './discover.js';
-import { renderPackageMap, type PackageMapData } from './format.js';
-import { extractFileLLM } from './llm-extract.js';
-import { extractFileAST } from './ast-extract.js';
-import { mergeFileData } from './merge.js';
-import { extractPackageLLM } from './llm-extract.js';
-import { writeFileSync } from 'fs';
-import { join } from 'path';
+import { discoverProject, type DirectoryEntry } from "./discover.js";
+import { renderPackageMap, type PackageMapData, type FileEntry } from "./format.js";
+import { extractFileLLM, extractPackageLLM } from "./llm-extract.js";
+import { extractFileAST } from "./ast-extract.js";
+import { mergeFileData } from "./merge.js";
+import { writeFileSync } from "fs";
+import { join } from "path";

 export async function initProject(rootPath: string): Promise<void> {
 	const entries = discoverProject(rootPath);

 	for (const entry of entries) {
-    const fileData = [];
+		await generateDirectoryMap(entry);
+	}
+
+	console.log(`Generated ${entries.length} .pi-map.md files`);
+}
+
+export async function generateDirectoryMap(
+	entry: DirectoryEntry,
+): Promise<FileEntry[]> {
+	const fileData: FileEntry[] = [];
 	for (const file of entry.files) {
 		const filePath = join(entry.dirPath, file);
 		const llmData = await extractFileLLM(filePath);
@@ -26,17 +34,15 @@ export async function initProject(rootPath: string): Promise<void> {
 		role: packageData.role,
 		files: fileData,
 		arch: packageData.arch,
-      dirty: '-'
+		dirty: "-",
 	};

-    const outPath = join(entry.dirPath, '.pi-map.md');
+	const outPath = join(entry.dirPath, ".pi-map.md");
 	writeFileSync(outPath, renderPackageMap(mapData));
-  }
-
-  console.log(`Generated ${entries.length} .pi-map.md files`);
+	return fileData;
 }

 export async function reinitPath(path: string): Promise<void> {
-  // TODO: clear dirty markers and force regeneration
+	// Full regeneration clears all dirty markers by overwriting every .pi-map.md
 	await initProject(path);
 }
@@ -1,5 +1,6 @@
-import { readFileSync } from 'fs';
-import { createHash } from 'crypto';
+import { readFileSync } from "fs";
+import { createHash } from "crypto";
+import { extname, basename } from "path";

 interface LLMFileData {
 	purpose: string;
@@ -12,37 +13,257 @@ interface LLMPackageData {
 	arch: string;
 }

-// Simple in-memory cache
 const cache = new Map<string, LLMFileData>();

+// Heuristic patterns for common file types
+const FILE_TYPE_PURPOSES: Record<string, string> = {
+	".ts": "TypeScript module",
+	".tsx": "React component",
+	".js": "JavaScript module",
+	".jsx": "React component",
+	".py": "Python module",
+	".go": "Go module",
+	".rs": "Rust module",
+	".java": "Java class",
+	".kt": "Kotlin class",
+	".swift": "Swift module",
+	".c": "C source",
+	".cpp": "C++ source",
+	".h": "C/C++ header",
+	".rb": "Ruby module",
+	".php": "PHP script",
+	".sh": "Shell script",
+	".md": "Documentation",
+	".json": "Configuration",
+	".yaml": "Configuration",
+	".yml": "Configuration",
+	".toml": "Configuration",
+	".ini": "Configuration",
+	".env": "Environment config",
+	".dockerfile": "Docker image definition",
+	"dockerfile": "Docker image definition",
+	".sql": "Database schema/queries",
+	".css": "Stylesheet",
+	".scss": "SCSS stylesheet",
+	".less": "LESS stylesheet",
+	".html": "HTML template",
+	".vue": "Vue component",
+	".svelte": "Svelte component",
+};
+
 export async function extractFileLLM(filePath: string): Promise<LLMFileData> {
-  const content = readFileSync(filePath, 'utf8');
-  const hash = createHash('sha256').update(content).digest('hex');
+	const content = readFileSync(filePath, "utf8");
+	const hash = createHash("sha256").update(content).digest("hex");

 	if (cache.has(hash)) {
 		return cache.get(hash)!;
 	}

-  // TODO: integrate with Pi's LLM tool or a generic client
-  // For now, return placeholder data
-  const result: LLMFileData = {
-    purpose: 'TODO: analyze with LLM',
-    exports: [],
-    deps: []
-  };
+	const ext = extname(filePath).toLowerCase();
+	const name = basename(filePath);
+	const baseName = basename(filePath, ext);

+	// Extract exports via heuristics
+	const exports = extractExports(content, ext, name);
+
+	// Extract dependencies via heuristics
+	const deps = extractDeps(content, ext);
+
+	// Generate purpose from filename + content heuristics
+	const purpose = generatePurpose(name, ext, baseName, content, exports);
+
+	const result: LLMFileData = { purpose, exports, deps };
 	cache.set(hash, result);
 	return result;
 }

+function extractExports(content: string, ext: string, _filename: string): string[] {
+	const exports: string[] = [];
+
+	if ([".ts", ".tsx", ".js", ".jsx", ".mjs"].includes(ext)) {
+		// ES module exports
+		const exportRegex =
+			/export\s+(?:default\s+)?(?:function\s+|class\s+|const\s+|let\s+|var\s+|interface\s+|type\s+|enum\s+)?([A-Za-z_$][A-Za-z0-9_$]*)/g;
+		let match: RegExpExecArray | null;
+		match = exportRegex.exec(content);
+		while (match !== null) {
+			exports.push(match[1]);
+			match = exportRegex.exec(content);
+		}
+
+		// Named export destructuring: export { foo, bar }
+		const namedExportRegex = /export\s*\{\s*([^}]+)\s*\}/g;
+		match = namedExportRegex.exec(content);
+		while (match !== null) {
+			const names = match[1].split(",").map((s) => s.trim().split(/\s+as\s+/)[0].trim());
+			exports.push(...names);
+			match = namedExportRegex.exec(content);
+		}
+	} else if (ext === ".py") {
+		// Python exports (top-level functions/classes)
+		const pyRegex = /^(?:async\s+)?def\s+([A-Za-z_][A-Za-z0-9_]*)|class\s+([A-Za-z_][A-Za-z0-9_]*)/gm;
+		let match: RegExpExecArray | null = pyRegex.exec(content);
+		while (match !== null) {
+			exports.push(match[1] || match[2]);
+			match = pyRegex.exec(content);
+		}
+	} else if (ext === ".go") {
+		// Go exports (capitalized functions/types)
+		const goRegex = /^(?:func|type|var|const)\s+([A-Z][A-Za-z0-9_]*)/gm;
+		let match: RegExpExecArray | null = goRegex.exec(content);
+		while (match !== null) {
+			exports.push(match[1]);
+			match = goRegex.exec(content);
+		}
+	} else if (ext === ".rs") {
+		// Rust exports (pub items)
+		const rsRegex = /pub\s+(?:fn|struct|enum|trait|type|const|static|use)\s+([A-Za-z_][A-Za-z0-9_]*)/g;
+		let match: RegExpExecArray | null = rsRegex.exec(content);
+		while (match !== null) {
+			exports.push(match[1]);
+			match = rsRegex.exec(content);
+		}
+	}
+
+	// Deduplicate while preserving order
+	return [...new Set(exports)];
+}
+
+function extractDeps(content: string, ext: string): string[] {
+	const deps: string[] = [];
+
+	if ([".ts", ".tsx", ".js", ".jsx", ".mjs"].includes(ext)) {
+		// ES imports
+		const importRegex =
+			/import\s+(?:(?:type\s+)?\{[^}]*\}|\*\s+as\s+\w+|\w+)\s+from\s+['"]([^'"]+)['"]/g;
+		let match: RegExpExecArray | null = importRegex.exec(content);
+		while (match !== null) {
+			deps.push(match[1]);
+			match = importRegex.exec(content);
+		}
+		// CommonJS requires
+		const requireRegex = /require\s*\(\s*['"]([^'"]+)['"]\s*\)/g;
+		match = requireRegex.exec(content);
+		while (match !== null) {
+			deps.push(match[1]);
+			match = requireRegex.exec(content);
+		}
+	} else if (ext === ".py") {
+		// Python imports
+		const pyImportRegex = /^(?:from|import)\s+([A-Za-z_][A-Za-z0-9_.]*)/gm;
+		let match: RegExpExecArray | null = pyImportRegex.exec(content);
+		while (match !== null) {
+			deps.push(match[1]);
+			match = pyImportRegex.exec(content);
+		}
+	} else if (ext === ".go") {
+		// Go imports
+		const goImportRegex = /"([^"]+)"/g;
+		let match: RegExpExecArray | null = goImportRegex.exec(content);
+		while (match !== null) {
+			if (match[1].includes("/")) deps.push(match[1]);
+			match = goImportRegex.exec(content);
+		}
+	} else if (ext === ".rs") {
+		// Rust use statements
+		const rsUseRegex = /use\s+([A-Za-z_][A-Za-z0-9_:]*)/g;
+		let match: RegExpExecArray | null = rsUseRegex.exec(content);
+		while (match !== null) {
+			deps.push(match[1]);
+			match = rsUseRegex.exec(content);
+		}
+	}
+
+	// Deduplicate
+	return [...new Set(deps)];
+}
+
+function generatePurpose(
+	name: string,
+	ext: string,
+	_baseName: string,
+	_content: string,
+	exports: string[],
+): string {
+	// Check for specific file patterns
+	if (/test|spec/i.test(name) && exports.length === 0) {
+		return "Test suite";
+	}
+	if (/config|settings/i.test(name)) return "Configuration";
+	if (/util|helper/i.test(name)) return "Utility functions";
+	if (/types?\.d?\.ts$/.test(name)) return "Type definitions";
+	if (/index\./.test(name)) return "Module entry point";
+	if (/middleware/.test(name)) return "Middleware";
+	if (/route/.test(name)) return "Route handlers";
+	if (/controller/.test(name)) return "Controller";
+	if (/service/.test(name)) return "Service layer";
+	if (/model/.test(name)) return "Data model";
+	if (/schema/.test(name)) return "Data schema";
+	if (/component/.test(name) || /\.tsx$/.test(name) || /\.vue$/.test(name) || /\.svelte$/.test(name)) {
+		return "UI component";
+	}
+	if (/hook|use[A-Z]/.test(name)) return "React hook";
+	if (/style|\.css|\.scss|\.less/.test(name)) return "Styling";
+	if (/docker/i.test(name)) return "Container definition";
+	if (/\.env/.test(name)) return "Environment variables";
+	if (/readme/i.test(name)) return "Project documentation";
+
+	// Use exports to infer purpose
+	if (exports.length > 0) {
+		const firstFew = exports.slice(0, 3).join(", ");
+		if (exports.length <= 3) return `Exports: ${firstFew}`;
+		return `Exports ${exports.length} symbols: ${firstFew}...`;
+	}
+
+	// Fallback to file type
+	return FILE_TYPE_PURPOSES[ext] || (ext ? `${ext.slice(1).toUpperCase()} file` : `${name} file`);
+}
+
 export async function extractPackageLLM(
 	relativePath: string,
-  fileData: { name: string; purpose: string }[]
+	fileData: { name: string; purpose: string }[],
 ): Promise<LLMPackageData> {
-  // TODO: integrate with Pi's LLM tool
-  // For now, return placeholder data
-  return {
-    role: `TODO: analyze package ${relativePath}`,
-    arch: 'TODO: architectural analysis'
-  };
+	const dirName = basename(relativePath);
+
+	// Infer role from directory name
+	let role = `Package ${dirName}`;
+	if (dirName === "src" || dirName === "lib" || dirName === "source") {
+		role = "Source code";
+	} else if (dirName === "test" || dirName === "tests" || dirName === "spec") {
+		role = "Test suite";
+	} else if (dirName === "docs" || dirName === "doc") {
+		role = "Documentation";
+	} else if (dirName === "config" || dirName === "configuration") {
+		role = "Configuration";
+	} else if (dirName === "utils" || dirName === "helpers" || dirName === "util") {
+		role = "Utility functions";
+	} else if (dirName === "types" || dirName === "type") {
+		role = "Type definitions";
+	} else if (dirName === "components" || dirName === "component") {
+		role = "UI components";
+	} else if (dirName === "hooks" || dirName === "hook") {
+		role = "Custom hooks";
+	} else if (dirName === "api" || dirName === "apis") {
+		role = "API endpoints/handlers";
+	} else if (dirName === "db" || dirName === "database" || dirName === "models") {
+		role = "Database layer";
+	} else if (dirName === "auth" || dirName === "authentication") {
+		role = "Authentication layer";
+	}
+
+	// Infer architecture from file patterns
+	const purposes = fileData.map((f) => f.purpose);
+	const hasTests = purposes.some((p) => p.includes("Test"));
+	const hasTypes = purposes.some((p) => p.includes("Type"));
+	const hasComponents = purposes.some((p) => p.includes("component") || p.includes("Component"));
+	const hasUtils = purposes.some((p) => p.includes("Utility"));
+
+	let arch = "";
+	if (hasTests) arch += "Contains tests. ";
+	if (hasTypes) arch += "Defines shared types. ";
+	if (hasComponents) arch += "Component-based architecture. ";
+	if (hasUtils) arch += "Shared utilities. ";
+	if (!arch) arch = `Contains ${fileData.length} files.`;
+
+	return { role, arch: arch.trim() };
 }
@@ -1,4 +1,4 @@
-import type { FileEntry } from './format.js';
+import type { FileEntry } from "./format.js";

 interface LLMFileData {
 	purpose: string;
@@ -14,12 +14,12 @@ interface ASTFileData {
 export function mergeFileData(
 	fileName: string,
 	llm: LLMFileData,
-  ast: ASTFileData | null
+	ast: ASTFileData | null,
 ): FileEntry {
 	return {
 		name: fileName,
 		purpose: llm.purpose,
 		exports: ast?.exports ?? llm.exports,
-    deps: [...new Set([...(ast?.deps ?? []), ...llm.deps])]
+		deps: [...new Set([...(ast?.deps ?? []), ...llm.deps])],
 	};
 }
@@ -1,16 +1,17 @@
-import { dirname, join } from 'path';
-import { existsSync, readFileSync, writeFileSync } from 'fs';
-import { parsePackageMap, renderPackageMap } from './format.js';
-import { extractFileLLM } from './llm-extract.js';
-import { extractFileAST } from './ast-extract.js';
-import { mergeFileData } from './merge.js';
-import { readdirSync } from 'fs';
+import { dirname, join, basename, relative } from "path";
+import { existsSync, readFileSync, writeFileSync } from "fs";
+import { parsePackageMap, renderPackageMap } from "./format.js";
+import { extractFileLLM } from "./llm-extract.js";
+import { extractFileAST } from "./ast-extract.js";
+import { mergeFileData } from "./merge.js";
+import { generateDirectoryMap } from "./init.js";
+import { readdirSync, statSync } from "fs";

 const SMALL_PACKAGE_THRESHOLD = 10;

 export async function patchFile(filePath: string): Promise<void> {
 	const dirPath = dirname(filePath);
-  const mapPath = join(dirPath, '.pi-map.md');
+	const mapPath = join(dirPath, ".pi-map.md");

 	if (!existsSync(mapPath)) {
 		// No map exists yet — would need to generate from scratch
@@ -18,23 +19,36 @@ export async function patchFile(filePath: string): Promise<void> {
 		return;
 	}

-  const allFiles = readdirSync(dirPath).filter((f: string) => !f.startsWith('.') && !f.endsWith('.md'));
+	const allFiles = readdirSync(dirPath).filter(
+		(f: string) => !f.startsWith(".") && !f.endsWith(".md"),
+	);
 	const isSmallPackage = allFiles.length < SMALL_PACKAGE_THRESHOLD;

 	if (isSmallPackage) {
 		// Full rewrite for small packages
-    // TODO: import and reuse init logic for a single directory
-    console.log(`Full rewrite of ${mapPath} (small package: ${allFiles.length} files)`);
+		const files = allFiles.filter((f) => {
+			const st = statSync(join(dirPath, f));
+			return st.isFile();
+		});
+		const relDir = relative(process.cwd(), dirPath) || ".";
+		await generateDirectoryMap({
+			dirPath,
+			relativePath: relDir,
+			files,
+		});
+		console.log(
+			`Full rewrite of ${mapPath} (small package: ${allFiles.length} files)`,
+		);
 	} else {
 		// Section-level patch
-    const existing = parsePackageMap(readFileSync(mapPath, 'utf8'));
+		const existing = parsePackageMap(readFileSync(mapPath, "utf8"));
 		const llmData = await extractFileLLM(filePath);
 		const astData = await extractFileAST(filePath);
-    const fileName = filePath.split('/').pop()!;
+		const fileName = basename(filePath);
 		const updatedFile = mergeFileData(fileName, llmData, astData);

 		// Replace the matching file entry
-    const idx = existing.files.findIndex(f => f.name === updatedFile.name);
+		const idx = existing.files.findIndex((f) => f.name === updatedFile.name);
 		if (idx >= 0) {
 			existing.files[idx] = updatedFile;
 		} else {
@@ -1,8 +1,8 @@
-import { discoverProject } from './discover.js';
-import { parsePackageMap } from './format.js';
-import { existsSync, readFileSync } from 'fs';
-import { join } from 'path';
-import { extractFileAST } from './ast-extract.js';
+import { discoverProject } from "./discover.js";
+import { parsePackageMap } from "./format.js";
+import { existsSync, readFileSync } from "fs";
+import { join } from "path";
+import { extractFileAST } from "./ast-extract.js";

 export interface ValidationResult {
 	clean: boolean;
@@ -10,41 +10,59 @@ export interface ValidationResult {
 }

 export interface Discrepancy {
-  type: 'missing' | 'orphaned' | 'stale-signature' | 'dirty';
+	type: "missing" | "orphaned" | "stale-signature" | "dirty";
 	path: string;
 	message: string;
 }

-export async function validateMaps(rootPath: string): Promise<ValidationResult> {
+export async function validateMaps(
+	rootPath: string,
+): Promise<ValidationResult> {
 	const discrepancies: Discrepancy[] = [];
 	const entries = discoverProject(rootPath);

 	for (const entry of entries) {
-    const mapPath = join(entry.dirPath, '.pi-map.md');
+		const mapPath = join(entry.dirPath, ".pi-map.md");
 		if (!existsSync(mapPath)) {
-      discrepancies.push({ type: 'missing', path: entry.relativePath, message: 'No .pi-map.md found' });
+			discrepancies.push({
+				type: "missing",
+				path: entry.relativePath,
+				message: "No .pi-map.md found",
+			});
 			continue;
 		}

-    const mapData = parsePackageMap(readFileSync(mapPath, 'utf8'));
+		const mapData = parsePackageMap(readFileSync(mapPath, "utf8"));

 		// Check for dirty markers
-    if (mapData.dirty && mapData.dirty !== '-') {
-      discrepancies.push({ type: 'dirty', path: mapPath, message: `Dirty: ${mapData.dirty}` });
+		if (mapData.dirty && mapData.dirty !== "-") {
+			discrepancies.push({
+				type: "dirty",
+				path: mapPath,
+				message: `Dirty: ${mapData.dirty}`,
+			});
 		}

 		// Check for orphaned entries
 		for (const fileEntry of mapData.files) {
 			const filePath = join(entry.dirPath, fileEntry.name);
 			if (!existsSync(filePath)) {
-        discrepancies.push({ type: 'orphaned', path: filePath, message: `File listed but deleted: ${fileEntry.name}` });
+				discrepancies.push({
+					type: "orphaned",
+					path: filePath,
+					message: `File listed but deleted: ${fileEntry.name}`,
+				});
 			}
 		}

 		// Check for new files not in map
 		for (const file of entry.files) {
-      if (!mapData.files.find(f => f.name === file)) {
-        discrepancies.push({ type: 'missing', path: join(entry.dirPath, file), message: `File not in .pi-map.md: ${file}` });
+			if (!mapData.files.find((f) => f.name === file)) {
+				discrepancies.push({
+					type: "missing",
+					path: join(entry.dirPath, file),
+					message: `File not in .pi-map.md: ${file}`,
+				});
 			}
 		}

@@ -60,12 +78,20 @@ export async function validateMaps(rootPath: string): Promise<ValidationResult>

 				for (const exp of listedExports) {
 					if (!actualExports.has(exp)) {
-            discrepancies.push({ type: 'stale-signature', path: filePath, message: `Missing export: ${exp}` });
+						discrepancies.push({
+							type: "stale-signature",
+							path: filePath,
+							message: `Missing export: ${exp}`,
+						});
 					}
 				}
 				for (const exp of actualExports) {
 					if (!listedExports.has(exp)) {
-            discrepancies.push({ type: 'stale-signature', path: filePath, message: `New export: ${exp}` });
+						discrepancies.push({
+							type: "stale-signature",
+							path: filePath,
+							message: `New export: ${exp}`,
+						});
 					}
 				}
 			}
@@ -74,11 +100,11 @@ export async function validateMaps(rootPath: string): Promise<ValidationResult>

 	const result: ValidationResult = {
 		clean: discrepancies.length === 0,
-    discrepancies
+		discrepancies,
 	};

 	if (result.clean) {
-    console.log('All .pi-map.md files are clean.');
+		console.log("All .pi-map.md files are clean.");
 	} else {
 		console.log(`Found ${discrepancies.length} discrepancies:`);
 		for (const d of discrepancies) {
@@ -0,0 +1,91 @@
+import { describe, it, expect } from "vitest";
+import { renderPackageMap, parsePackageMap, type PackageMapData } from "../src/format.js";
+
+const sampleData: PackageMapData = {
+	path: "pkg/auth",
+	role: "Auth layer: JWT issuance, validation, refresh. Stateless. Dep: pkg/crypto, pkg/db.",
+	files: [
+		{
+			name: "tokens.ts",
+			purpose: "JWT gen/val",
+			exports: ["issueToken", "verifyToken", "refreshToken"],
+			deps: ["crypto/hmac", "db/sessions"],
+		},
+		{
+			name: "middleware.ts",
+			purpose: "HTTP auth guard",
+			exports: ["requireAuth", "requireRole"],
+			deps: ["tokens/verifyToken"],
+		},
+		{
+			name: "types.ts",
+			purpose: "shared auth types",
+			exports: ["AuthToken", "UserClaims", "Role"],
+			deps: [],
+		},
+	],
+	arch: "Guard pattern on routes. Tokens short-lived (15m), refresh long-lived (7d). Rotation on every use.",
+	dirty: "-",
+};
+
+describe("format", () => {
+	it("renders package map correctly", () => {
+		const output = renderPackageMap(sampleData);
+		expect(output).toContain("# pkg/auth");
+		expect(output).toContain("## role");
+		expect(output).toContain("## files");
+		expect(output).toContain("- tokens.ts | JWT gen/val | exp: issueToken, verifyToken, refreshToken | dep: crypto/hmac, db/sessions");
+		expect(output).toContain("## arch");
+		expect(output).toContain("## dirty");
+		expect(output).toContain("-");
+	});
+
+	it("round-trips parse and render", () => {
+		const rendered = renderPackageMap(sampleData);
+		const parsed = parsePackageMap(rendered);
+
+		expect(parsed.path).toBe(sampleData.path);
+		expect(parsed.role).toBe(sampleData.role);
+		expect(parsed.arch).toBe(sampleData.arch);
+		expect(parsed.dirty).toBeUndefined(); // "-" is normalized to undefined
+		expect(parsed.files).toHaveLength(3);
+
+		const tokens = parsed.files.find((f) => f.name === "tokens.ts");
+		expect(tokens).toBeDefined();
+		expect(tokens!.purpose).toBe("JWT gen/val");
+		expect(tokens!.exports).toEqual(["issueToken", "verifyToken", "refreshToken"]);
+		expect(tokens!.deps).toEqual(["crypto/hmac", "db/sessions"]);
+	});
+
+	it("handles files without exports or deps", () => {
+		const data: PackageMapData = {
+			path: "pkg/utils",
+			role: "Utilities",
+			files: [{ name: "helpers.ts", purpose: "Helpers", exports: [], deps: [] }],
+			arch: "Shared helpers",
+			dirty: "2024-01-01: patched",
+		};
+		const rendered = renderPackageMap(data);
+		const parsed = parsePackageMap(rendered);
+		expect(parsed.files[0].exports).toEqual([]);
+		expect(parsed.files[0].deps).toEqual([]);
+		expect(parsed.dirty).toBe("2024-01-01: patched");
+	});
+
+	it("parses multiline arch", () => {
+		const markdown = `# pkg/test
+## role
+Test package
+## files
+- test.ts | Test file
+## arch
+Line one.
+Line two.
+Line three.
+## dirty
+-
+`;
+		const parsed = parsePackageMap(markdown);
+		expect(parsed.arch).toBe("Line one.\nLine two.\nLine three.");
+	});
+});
@@ -0,0 +1,63 @@
+import { describe, it, expect } from "vitest";
+import { extractFileLLM } from "../src/llm-extract.js";
+import { writeFileSync, mkdtempSync } from "fs";
+import { join } from "path";
+import { tmpdir } from "os";
+
+describe("llm-extract heuristics", () => {
+	it("extracts TypeScript exports", async () => {
+		const dir = mkdtempSync(join(tmpdir(), "pi-map-"));
+		const file = join(dir, "test.ts");
+		writeFileSync(
+			file,
+			`export function foo() {}
+export class Bar {}
+export const baz = 1;
+export type Qux = string;
+export { a, b as c };
+`,
+		);
+		const result = await extractFileLLM(file);
+		expect(result.exports).toContain("foo");
+		expect(result.exports).toContain("Bar");
+		expect(result.exports).toContain("baz");
+		expect(result.exports).toContain("Qux");
+		expect(result.exports).toContain("a");
+		expect(result.exports).toContain("b");
+	});
+
+	it("extracts TypeScript imports", async () => {
+		const dir = mkdtempSync(join(tmpdir(), "pi-map-"));
+		const file = join(dir, "test.ts");
+		writeFileSync(
+			file,
+			`import { foo } from "./bar";
+import * as baz from "baz-lib";
+import type { Qux } from "qux";
+const x = require("legacy");
+`,
+		);
+		const result = await extractFileLLM(file);
+		expect(result.deps).toContain("./bar");
+		expect(result.deps).toContain("baz-lib");
+		expect(result.deps).toContain("qux");
+		expect(result.deps).toContain("legacy");
+	});
+
+	it("infers purpose from filename patterns", async () => {
+		const dir = mkdtempSync(join(tmpdir(), "pi-map-"));
+		const file = join(dir, "userController.ts");
+		writeFileSync(file, `export class UserController {}`);
+		const result = await extractFileLLM(file);
+		expect(result.purpose).toMatch(/Controller|Exports/);
+	});
+
+	it("handles non-code files", async () => {
+		const dir = mkdtempSync(join(tmpdir(), "pi-map-"));
+		const file = join(dir, "Dockerfile");
+		writeFileSync(file, `FROM node:20\nWORKDIR /app`);
+		const result = await extractFileLLM(file);
+		expect(result.purpose).toBe("Container definition");
+		expect(result.exports).toEqual([]);
+	});
+});