@arcmantle/lit-jsx 1.0.6 → 1.0.8

package/src/compiler/import-discovery.ts

@@ -1,281 +1,460 @@
 import { existsSync, readFileSync } from 'node:fs';
-import { dirname, resolve } from 'node:path';
+import { dirname } from 'node:path';
 
 import * as babel from '@babel/core';
-import traverse, { Binding, Hub, NodePath, Scope } from '@babel/traverse';
+import type { Binding, NodePath, Scope } from '@babel/traverse';
 import * as t from '@babel/types';
+import oxcResolver from 'oxc-resolver';
+import type { Logger } from 'pino';
 
+import { traverse } from './babel-traverse.ts';
+import { getPathFilename, isComponent } from './compiler-utils.ts';
+import { babelPlugins, debugMode } from './config.ts';
+import { createLogger } from './create-logger.ts';
 
-// Cache for parsed files to avoid re-parsing
-const fileCache: Map<string, t.File> = new Map();
 
-
-export type BabelPlugins = NonNullable<NonNullable<babel.TransformOptions['parserOpts']>['plugins']>;
-
-
-// Types for our discovery results
 export interface ElementDefinition {
-	type:            'custom-element' | 'import' | 'local-variable' | 'unknown';
-	source?:         string; // file path for imports
-	originalName?:   string; // for imports/re-exports
-	callExpression?: t.CallExpression; // for toJSX calls
+	type:            'custom-element' | 'import' | 'local-variable' | 'wildcard-export' | 'unknown';
+	source?:         string;
+	originalName?:   string;
+	localName?:      string;
+	callExpression?: t.CallExpression;
+	resolvedPath?:   string; // For lazy import resolution
+	referencedName?: string; // For lazy local reference resolution
 }
 
+interface FileSystemAdapter {
+	existsSync:   (path: string) => boolean;
+	readFileSync: (path: string, encoding: 'utf-8') => string;
+	dirname:      (path: string) => string;
+}
 
-// Helper function to find the definition of a JSX element
-export function findElementDefinition(path: NodePath<t.JSXOpeningElement>): ElementDefinition {
-	const elementName = path.node.name;
+interface ResolverAdapter {
+	sync: (basedir: string, module: string) => { path?: string; };
+}
 
-	const hub = path.hub as Hub & { file: { opts: { filename: string; }; }; };
 
-	const currentFileName = hub.file.opts.filename;
+export class ImportDiscovery {
 
-	// Only handle JSXIdentifier (not JSXMemberExpression or JSXNamespacedName)
-	if (!t.isJSXIdentifier(elementName))
-		return { type: 'unknown' };
+	static readonly definitionCache:   Map<string, Map<string, ElementDefinition>> = new Map();
+	static readonly fileBindingsCache: Map<string, ReadonlyMap<string, ElementDefinition>> = new Map();
+	static readonly fileDependencies:  Map<string, Set<string>> = new Map();
+	static readonly EMPTY_BINDINGS:    ReadonlyMap<string, ElementDefinition> = new Map();
 
-	const elementNameString = elementName.name;
+	static clearCacheForFileAndDependents(changedFilePath: string): void {
+		// Clear the changed file itself
+		ImportDiscovery.definitionCache.delete(changedFilePath);
+		ImportDiscovery.fileBindingsCache.delete(changedFilePath);
+		ImportDiscovery.fileDependencies.delete(changedFilePath);
 
-	//console.log('Tracing element:', elementNameString);
+		// Find and clear all files that depend on the changed file in one pass
+		for (const [ file, dependencies ] of ImportDiscovery.fileDependencies) {
+			if (!dependencies.has(changedFilePath))
+				continue;
 
-	// Start tracing with the current file context
-	return traceElementDefinition(elementNameString, path.scope, currentFileName);
-}
+			ImportDiscovery.definitionCache.delete(file);
+			ImportDiscovery.fileBindingsCache.delete(file);
+			ImportDiscovery.fileDependencies.delete(file);
+		}
+	}
 
-// Core recursive tracing function
-function traceElementDefinition(
-	elementName: string,
-	scope: Scope,
-	currentFileName: string,
-	visitedFiles: Set<string> = new Set(),
-): ElementDefinition {
-	// Prevent infinite recursion
-	if (visitedFiles.has(`${ currentFileName }:${ elementName }`))
-		return { type: 'unknown' };
+	protected readonly visitedFiles: Set<string> = new Set();
+	protected readonly resolver:     ResolverAdapter;
+	protected readonly log:          Logger<never, boolean>;
+	protected readonly fs:           FileSystemAdapter;
 
-	visitedFiles.add(`${ currentFileName }:${ elementName }`);
+	constructor() {
+		this.resolver = oxcResolver;
+		this.log      = createLogger('import-discovery', debugMode.value);
+		this.fs       = { existsSync, readFileSync, dirname };
+	}
 
-	// Check if there's a binding for this identifier in the current scope
-	const binding = scope.getBinding(elementName);
+	/**
+	 * Finds the definition of a JSX element in the given path.
+	 */
+	findElementDefinition(
+		path: NodePath<t.JSXOpeningElement>,
+	): ElementDefinition {
+		this.visitedFiles.clear();
+
+		const filePath = getPathFilename(path);
+		const cacheKey = String(path.node.start);
+
+		const fileCache = ImportDiscovery.definitionCache.get(filePath);
+		if (fileCache) {
+			const cached = fileCache.get(cacheKey);
+			if (cached)
+				return cached;
+		}
 
-	if (!binding) {
-		//console.log('No binding found for:', elementName);
+		if (!t.isJSXIdentifier(path.node.name))
+			return { type: 'unknown' };
 
-		return { type: 'unknown' };
-	}
+		const elementName = path.node.name.name;
+		if (!isComponent(elementName))
+			return { type: 'unknown' };
 
-	//console.log('Binding kind:', binding.kind, 'type:', binding.path.type);
+		const currentFileName = getPathFilename(path);
+		const result = this.traceElementDefinition(elementName, path.scope, currentFileName);
 
-	// Handle imports
-	if (binding.kind === 'module' && t.isImportSpecifier(binding.path.node))
-		return traceImport(binding, currentFileName, elementName, visitedFiles);
+		// Store in file-specific cache
+		const definitionCache = ImportDiscovery.definitionCache.get(filePath)
+			?? ImportDiscovery.definitionCache
+				.set(filePath, new Map())
+				.get(filePath)!;
 
+		definitionCache.set(cacheKey, result);
 
-	// Handle local variables/constants
-	if (binding.kind === 'const' || binding.kind === 'let' || binding.kind === 'var')
-		return traceLocalVariable(binding, currentFileName, visitedFiles);
+		return result;
+	}
 
+	// Trace the element definition recursively
+	protected traceElementDefinition(
+		elementName: string,
+		scope: Scope,
+		currentFileName: string,
+	): ElementDefinition {
+		const traceKey = `${ currentFileName }:${ elementName }`;
 
-	return { type: 'unknown' };
-}
+		// Prevent infinite recursion
+		if (this.visitedFiles.has(traceKey))
+			return { type: 'unknown' };
 
-function traceImport(
-	binding: Binding,
-	currentFileName: string,
-	elementName: string,
-	visitedFiles: Set<string>,
-): ElementDefinition {
-	const importDeclaration = binding.path.parent;
+		this.visitedFiles.add(traceKey);
 
-	if (!t.isImportDeclaration(importDeclaration))
-		return { type: 'unknown' };
+		// Use batched file analysis
+		const fileBindings = this.analyzeFileBindings(currentFileName);
 
-	const importSource = importDeclaration.source.value;
-	const currentDir = dirname(currentFileName);
-	const resolvedPath = resolve(currentDir, importSource);
+		// Check if we have this element in our batch analysis
+		if (fileBindings.has(elementName)) {
+			const definition = fileBindings.get(elementName)!;
 
-	//console.log('Tracing import from:', importSource);
-	//console.log('Resolved to:', resolvedPath);
+			// Resolve any lazy references
+			return this.resolveLazyDefinition(definition);
+		}
 
-	// Use cached parsing
-	const ast = getOrParseFile(resolvedPath);
-	if (!ast)
-		return { type: 'unknown' };
+		// Fallback to scope-based lookup for dynamic cases
+		const binding = scope.getBinding(elementName);
+		if (!binding)
+			return { type: 'unknown' };
 
+		// Use the fast analysis methods
+		const result = this.analyzeBindingFast(binding, currentFileName);
 
-	let result: ElementDefinition = { type: 'unknown' };
+		// Resolve any lazy references (imports, local references)
+		return this.resolveLazyDefinition(result);
+	}
 
-	traverse(ast, {
-		Program(programPath) {
-			// First try to find a local binding (normal export)
-			const localBinding = programPath.scope.getBinding(elementName);
+	// Analyze all relevant bindings in a file at once
+	protected analyzeFileBindings(filePath: string): ReadonlyMap<string, ElementDefinition> {
+		const fileBinding = ImportDiscovery.fileBindingsCache.get(filePath);
+		if (fileBinding)
+			return fileBinding;
 
-			if (localBinding) {
-				//console.log('Found local binding in imported file');
-				result = traceElementDefinition(elementName, programPath.scope, resolvedPath, visitedFiles);
+		if (!this.fs.existsSync(filePath)) {
+			ImportDiscovery.fileBindingsCache.set(filePath, ImportDiscovery.EMPTY_BINDINGS);
 
-				return;
-			}
+			return ImportDiscovery.EMPTY_BINDINGS;
+		}
 
-			// If no local binding found, check for re-exports
-			//console.log('No local binding, checking for re-exports...');
-			result = checkForReExports(programPath, elementName, resolvedPath, visitedFiles);
-		},
-	});
+		const fileContent = this.fs.readFileSync(filePath, 'utf-8');
+		let ast: t.File;
 
-	return result;
-}
+		try {
+			ast = babel.parseSync(fileContent, {
+				filename:   filePath,
+				parserOpts: {
+					plugins: babelPlugins,
+				},
+			})!;
+		}
+		catch (error) {
+			// Failed to parse, cache empty result
+			ImportDiscovery.fileBindingsCache.set(filePath, ImportDiscovery.EMPTY_BINDINGS);
 
-function traceLocalVariable(
-	binding: Binding,
-	currentFileName: string,
-	visitedFiles: Set<string>,
-): ElementDefinition {
-	//console.log('Tracing local variable:', binding.kind, binding.path.type);
+			return ImportDiscovery.EMPTY_BINDINGS;
+		}
 
-	// Check if it's a variable declarator (const/let/var)
-	if (t.isVariableDeclarator(binding.path.node)) {
-		const declarator = binding.path.node;
+		let programPath: babel.NodePath<babel.types.Program> = undefined as any;
+		traverse(ast, { Program(path) { programPath = path; path.stop(); } });
 
-		// Check if it's assigned to a call expression
-		if (declarator.init && t.isCallExpression(declarator.init)) {
-			const callExpr = declarator.init;
+		const bindings: Map<string, ElementDefinition> = new Map();
 
-			// Check if it's a toJSX call
-			if (t.isIdentifier(callExpr.callee) && callExpr.callee.name === 'toJSX') {
-				return {
-					type:           'custom-element',
-					source:         currentFileName,
-					callExpression: callExpr,
-				};
-			}
+		// 1. Analyze all relevant local bindings at once
+		this.analyzeFileDeclarations(programPath, filePath, bindings);
 
-			// Could be assigned to another function call - trace that too
-			//console.log('Local variable assigned to call expression:', callExpr.callee);
-		}
+		// 2. Analyze all exports at once
+		this.analyzeFileExports(programPath, filePath, bindings);
 
-		// Check if it's assigned to an identifier (another variable)
-		if (declarator.init && t.isIdentifier(declarator.init)) {
-			const assignedIdentifier = declarator.init.name;
-			//console.log('Local variable assigned to identifier:', assignedIdentifier);
+		const readonlyBindings = new Map(bindings) as ReadonlyMap<string, ElementDefinition>;
+		ImportDiscovery.fileBindingsCache.set(filePath, readonlyBindings);
 
-			// Recursively trace this identifier in the same scope
-			return traceElementDefinition(assignedIdentifier, binding.path.scope, currentFileName, visitedFiles);
-		}
+		return readonlyBindings;
 	}
 
-	return { type: 'local-variable' };
-}
+	// Resolve lazy references in the definition
+	protected analyzeFileDeclarations(
+		programPath: babel.NodePath<babel.types.Program>,
+		filePath: string,
+		bindings: Map<string, ElementDefinition>,
+	): void {
+		for (const [ name, binding ] of Object.entries(programPath.scope.bindings)) {
+			// Skip function/import bindings that are clearly not component-related
+			if (binding.kind === 'module' || binding.kind === 'hoisted') {
+				if (!isComponent(name))
+					continue;
+			}
 
-function checkForReExports(
-	programPath: babel.NodePath<babel.types.Program>,
-	elementName: string,
-	currentFileName: string,
-	visitedFiles: Set<string>,
-): ElementDefinition {
-	let result: ElementDefinition = { type: 'unknown' };
+			const definition = this.analyzeBindingFast(binding, filePath);
+			if (definition.type !== 'unknown')
+				bindings.set(name, definition);
+		}
+	}
 
-	// Check all export declarations for re-exports
-	programPath.traverse({
-		ExportNamedDeclaration(exportPath) {
-			const node = exportPath.node;
+	// Fast binding analysis without deep traversal
+	protected analyzeBindingFast(
+		binding: Binding,
+		filePath: string,
+	): ElementDefinition {
+		// Handle imports
+		if (binding.kind === 'module' && t.isImportSpecifier(binding.path.node))
+			return this.analyzeImportBinding(binding, filePath);
 
-			if (!node.source || !node.specifiers)
-				return; // Skip if no source or specifiers
+		// Handle local variables
+		if (binding.kind === 'const' || binding.kind === 'let' || binding.kind === 'var')
+			return this.analyzeLocalBinding(binding, filePath);
 
-			// Handle re-exports: export { X } from './file'
-			//console.log('Found re-export from:', node.source.value);
+		return { type: 'unknown' };
+	}
 
-			for (const specifier of node.specifiers) {
-				if (!t.isExportSpecifier(specifier))
-					continue; // Only handle export specifiers
+	// Analyze import without deep file traversal
+	protected analyzeImportBinding(binding: Binding, currentFileName: string): ElementDefinition {
+		const importDeclaration = binding.path.parent;
+		if (!t.isImportDeclaration(importDeclaration))
+			return { type: 'unknown' };
+
+		const importSpecifier = binding.path.node;
+		if (!t.isImportSpecifier(importSpecifier))
+			return { type: 'unknown' };
+
+		const originalExportedName = t.isIdentifier(importSpecifier.imported)
+			? importSpecifier.imported.name
+			: importSpecifier.imported.value;
+
+		const importSource = importDeclaration.source.value;
+		const currentDir = dirname(currentFileName);
+
+		const resolvedResult = this.resolver.sync(currentDir, importSource);
+		const resolvedPath = resolvedResult.path;
+
+		if (!resolvedPath)
+			return { type: 'unknown' };
+
+		// Instead of deep traversal, just mark as import and resolve lazily
+		return {
+			type:         'import',
+			source:       importSource,
+			originalName: originalExportedName,
+			localName:    binding.identifier.name,
+			// Store resolved path for later lookup
+			resolvedPath: resolvedPath,
+		};
+	}
 
-				const exportedName = t.isIdentifier(specifier.exported)
-					? specifier.exported.name
-					: specifier.exported.value;
+	// Analyze local binding without recursion
+	protected analyzeLocalBinding(binding: Binding, filePath: string): ElementDefinition {
+		if (!t.isVariableDeclarator(binding.path.node))
+			return { type: 'local-variable' };
 
-				// Check if this re-export matches our element name
-				if (exportedName !== elementName)
-					continue;
 
-				const originalName = specifier.local.name;
-				//console.log(`Found re-export: ${ originalName } as ${ exportedName }`);
+		const declarator = binding.path.node;
+		if (!declarator.init)
+			return { type: 'local-variable' };
 
-				// Resolve and trace the re-exported file
-				const reExportSource = node.source.value;
-				const currentDir = dirname(currentFileName);
-				const resolvedPath = resolve(currentDir, reExportSource);
 
-				if (!existsSync(resolvedPath))
-					continue; // Skip if file doesn't exist
+		// Check for toComponent/toTag calls
+		if (t.isCallExpression(declarator.init)) {
+			if (this.isToComponentOrTagCall(declarator.init, binding.path.scope)) {
+				return {
+					type:           'custom-element',
+					source:         filePath,
+					callExpression: declarator.init,
+				};
+			}
+		}
 
-				//console.log('Tracing re-export to:', resolvedPath);
-				result = traceReExport(originalName, resolvedPath, visitedFiles);
+		// For identifier assignments, store reference for later resolution
+		if (t.isIdentifier(declarator.init)) {
+			return {
+				type:           'local-variable',
+				source:         filePath,
+				referencedName: declarator.init.name,
+			};
+		}
 
-				// Stop traversing once we find the match
-				return exportPath.stop();
-			}
-		},
-	});
+		return { type: 'local-variable' };
+	}
 
-	return result;
-}
+	// Helper function to check if a call expression is a toComponent/toTag call
+	// even if the function has been renamed through imports
+	protected isToComponentOrTagCall(
+		callExpr: t.CallExpression,
+		scope: Scope,
+	): boolean {
+		if (!t.isIdentifier(callExpr.callee))
+			return false;
 
-function traceReExport(
-	elementName: string,
-	filePath: string,
-	visitedFiles: Set<string>,
-): ElementDefinition {
-	// Use cached parsing
-	const ast = getOrParseFile(filePath);
-	if (!ast)
-		return { type: 'unknown' };
+		const functionName = callExpr.callee.name;
 
+		// Check direct names first (fast path)
+		if (functionName === 'toComponent' || functionName === 'toTag')
+			return true;
 
-	let result: ElementDefinition = { type: 'unknown' };
+		// Check if this identifier is bound to an import that originally was toComponent/toTag
+		const binding = scope.getBinding(functionName);
+		if (!binding || binding.kind !== 'module')
+			return false;
 
-	traverse(ast, {
-		Program(programPath) {
-			// Continue tracing in the re-exported file
-			result = traceElementDefinition(elementName, programPath.scope, filePath, visitedFiles);
-		},
-	});
+		if (!t.isImportSpecifier(binding.path.node))
+			return false;
 
-	return result;
-}
+		const importSpecifier = binding.path.node;
+		const originalImportedName = t.isIdentifier(importSpecifier.imported)
+			? importSpecifier.imported.name
+			: importSpecifier.imported.value;
 
+		// Check if the original imported name was toComponent or toTag
+		const isOriginallyToComponentOrTag =
+				originalImportedName === 'toComponent'
+			|| originalImportedName === 'toTag';
 
-// Helper function to get or parse a file with caching
-function getOrParseFile(filePath: string): t.File | undefined {
-	// Check cache first
-	if (fileCache.has(filePath)) {
-		//console.log('Using cached AST for:', filePath);
+		return isOriginallyToComponentOrTag;
+	}
 
-		return fileCache.get(filePath)!;
+	// Analyze all exports in one pass
+	protected analyzeFileExports(
+		programPath: babel.NodePath<babel.types.Program>,
+		filePath: string,
+		bindings: Map<string, ElementDefinition>,
+	): void {
+		for (const statement of programPath.node.body) {
+			// Handle named exports: export { X } from './file' or export { X }
+			if (t.isExportNamedDeclaration(statement)) {
+				for (const specifier of statement.specifiers) {
+					if (!t.isExportSpecifier(specifier))
+						continue;
+
+					const exportedName = t.isIdentifier(specifier.exported)
+						? specifier.exported.name
+						: specifier.exported.value;
+
+					const localName = specifier.local.name;
+
+					if (!isComponent(exportedName))
+						continue;
+
+					// For re-exports with source
+					if (statement.source) {
+						const definition = {
+							type:         'import' as const,
+							source:       statement.source.value,
+							originalName: localName,
+							localName:    exportedName,
+						};
+						bindings.set(exportedName, definition);
+					}
+					else {
+						// For local exports, reference the local binding
+						const definition = {
+							type:           'local-variable' as const,
+							source:         filePath,
+							referencedName: localName,
+							localName:      exportedName,
+						};
+						bindings.set(exportedName, definition);
+					}
+				}
+			}
+			// Handle wildcard exports: export * from './file'
+			else if (t.isExportAllDeclaration(statement) && statement.source) {
+				// For wildcard exports, we need to mark this as a wildcard re-export
+				// The actual resolution will happen in resolveLazyDefinition
+				const definition = {
+					type:   'wildcard-export' as const,
+					source: statement.source.value,
+				};
+				// Store with a special key to indicate wildcard export
+				bindings.set('*', definition);
+			}
+		}
 	}
 
-	// File not in cache, parse it
-	if (!existsSync(filePath))
-		return;
+	// Resolve lazy definitions (imports, references)
+	protected resolveLazyDefinition(definition: ElementDefinition): ElementDefinition {
+		if (definition.type === 'import' && definition.resolvedPath && definition.originalName) {
+			if (definition.source) {
+				const currentFile = definition.source;
+				const dependsOn = definition.resolvedPath;
 
-	const fileContent = readFileSync(filePath, 'utf-8');
+				let fileDependencies = ImportDiscovery.fileDependencies.get(currentFile);
+				if (!fileDependencies) {
+					fileDependencies = new Set<string>();
+					ImportDiscovery.fileDependencies.set(currentFile, fileDependencies);
+				}
 
-	try {
-		const ast = babel.parseSync(fileContent, {
-			filename:   filePath,
-			parserOpts: {
-				plugins: [ 'jsx', 'typescript' ] satisfies BabelPlugins,
-			},
-		});
+				fileDependencies.add(dependsOn);
+			}
 
-		if (ast) {
-			//console.log('Parsed and cached:', filePath);
-			fileCache.set(filePath, ast);
+			// Recursively analyze the imported file
+			const importedBindings = this.analyzeFileBindings(definition.resolvedPath);
+			const binding = importedBindings.get(definition.originalName);
+
+			if (binding) // Recursively resolve the found definition
+				return this.resolveLazyDefinition(binding);
+
+			// If specific export not found, check for wildcard exports
+			const wildcardExport = importedBindings.get('*');
+			if (wildcardExport && wildcardExport.type === 'wildcard-export') {
+				// Resolve the wildcard export by looking in the target file
+				const currentDir = this.fs.dirname(definition.resolvedPath);
+				const resolvedResult = this.resolver.sync(currentDir, wildcardExport.source!);
+				const resolvedPath = resolvedResult.path;
+
+				if (resolvedPath) {
+					// Create a new import definition for the wildcard target
+					const wildcardTargetDefinition: ElementDefinition = {
+						type:         'import',
+						source:       wildcardExport.source,
+						originalName: definition.originalName,
+						localName:    definition.localName,
+						resolvedPath: resolvedPath,
+					};
+
+					return this.resolveLazyDefinition(wildcardTargetDefinition);
+				}
+			}
+		}
+
+		if (definition.type === 'local-variable' && definition.referencedName && definition.source) {
+			// Resolve local references
+			const fileBindings = this.analyzeFileBindings(definition.source);
+			const binding = fileBindings.get(definition.referencedName);
 
-			return ast;
+			if (binding)
+				return this.resolveLazyDefinition(binding);
 		}
+
+		return definition;
 	}
-	catch (error) {
-		console.log('Failed to parse file:', filePath, error);
-	}
+
 }
+
+
+let discovery: ImportDiscovery;
+export const findElementDefinition = (
+	...args: Parameters<ImportDiscovery['findElementDefinition']>
+): ReturnType<ImportDiscovery['findElementDefinition']> => {
+	discovery ??= new ImportDiscovery();
+
+	return discovery.findElementDefinition(...args);
+};