Skip to content

Commit e268a63

Browse files
knutwannhedenknutwannheden
committed
JavaScript: Improvements to the AddImport visitor
- Improve whitespace handling - Improve the `onlyIfReferenced` detection
1 parent a794e25 commit e268a63

2 files changed

Lines changed: 359 additions & 43 deletions

File tree

‎rewrite-javascript/rewrite/src/javascript/add-import.ts‎

Lines changed: 124 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,8 @@
11
import {JavaScriptVisitor} from "./visitor";
2-
import {J, emptySpace, rightPadded, space, Statement, singleSpace, Type} from "../java";
2+
import {emptySpace, J, rightPadded, singleSpace, space, Statement, Type} from "../java";
33
import {JS} from "./tree";
44
import {randomId} from "../uuid";
55
import {emptyMarkers, markers} from "../markers";
6-
import {ExecutionContext} from "../execution";
76

87
export enum ImportStyle {
98
ES6Named, // import { x } from 'module'
@@ -454,14 +453,27 @@ export class AddImport<P> extends JavaScriptVisitor<P> {
454453
const newSpecifierBase = this.createImportSpecifier();
455454
const newSpecifier = {...newSpecifierBase, prefix: singleSpace};
456455

456+
// Transfer the right padding from the element before the insertion point to the new element
457+
// Since we're appending, this is the last existing element
458+
const existingElements = namedImports.elements.elements;
459+
const elementBeforeInsertion = existingElements[existingElements.length - 1];
460+
const paddingToTransfer = elementBeforeInsertion.after;
461+
457462
// Add the new specifier to the elements
458463
const updatedNamedImports: JS.NamedImports = await this.produceJavaScript<JS.NamedImports>(
459464
namedImports, p, async namedDraft => {
465+
// Update the element before insertion to have emptySpace as its right padding (before the comma)
466+
const updatedExistingElements = existingElements.slice(0, -1).concat({
467+
...elementBeforeInsertion,
468+
after: emptySpace
469+
});
470+
460471
namedDraft.elements = {
461472
...namedImports.elements,
462473
elements: [
463-
...namedImports.elements.elements,
464-
rightPadded(newSpecifier, emptySpace)
474+
...updatedExistingElements,
475+
// Transfer the padding to the new element (after the comma, before the closing brace)
476+
rightPadded(newSpecifier, paddingToTransfer)
465477
]
466478
};
467479
}
@@ -660,73 +672,143 @@ export class AddImport<P> extends JavaScriptVisitor<P> {
660672
* Check if the identifier is actually referenced in the file
661673
*/
662674
private async checkIdentifierReferenced(compilationUnit: JS.CompilationUnit): Promise<boolean> {
663-
// Use type attribution to detect if the identifier is referenced
664-
// Map of module name -> Set of member names used from that module
665-
const usedImports = new Map<string, Set<string>>();
666-
667-
// Helper to record usage of a method from a module
668-
const recordMethodUsage = (methodType: Type.Method) => {
669-
const moduleName = Type.FullyQualified.getFullyQualifiedName(methodType.declaringType);
670-
if (moduleName) {
671-
if (!usedImports.has(moduleName)) {
672-
usedImports.set(moduleName, new Set());
675+
// For namespace imports, we cannot use type attribution to detect usage
676+
// because the namespace itself is used as an identifier, not individual members.
677+
// For simplicity, we skip the onlyIfReferenced check for namespace imports.
678+
if (this.member === '*') {
679+
// TODO: Implement proper namespace usage detection by checking if alias identifier is used
680+
return true;
681+
}
682+
683+
// Step 1: Find the expected declaring type by examining existing imports from the same module
684+
let expectedDeclaringType: string | undefined;
685+
686+
for (const stmt of compilationUnit.statements) {
687+
const statement = stmt.element;
688+
689+
if (statement?.kind === JS.Kind.Import) {
690+
const jsImport = statement as JS.Import;
691+
const moduleSpecifier = jsImport.moduleSpecifier?.element;
692+
693+
if (!moduleSpecifier) {
694+
continue;
695+
}
696+
697+
const moduleName = this.getModuleName(moduleSpecifier);
698+
if (moduleName !== this.module) {
699+
continue; // Not the module we're interested in
700+
}
701+
702+
// Found an existing import from our target module
703+
// Extract the declaring type from any imported member with type attribution
704+
const importClause = jsImport.importClause;
705+
if (importClause?.namedBindings?.kind === JS.Kind.NamedImports) {
706+
const namedImports = importClause.namedBindings as JS.NamedImports;
707+
for (const elem of namedImports.elements.elements) {
708+
const specifier = elem.element;
709+
if (specifier?.kind === JS.Kind.ImportSpecifier) {
710+
const importSpec = specifier as JS.ImportSpecifier;
711+
let identifier: J.Identifier | undefined;
712+
if (importSpec.specifier?.kind === J.Kind.Identifier) {
713+
identifier = importSpec.specifier as J.Identifier;
714+
} else if (importSpec.specifier?.kind === JS.Kind.Alias) {
715+
const aliasSpec = importSpec.specifier as JS.Alias;
716+
if (aliasSpec.alias?.kind === J.Kind.Identifier) {
717+
identifier = aliasSpec.alias as J.Identifier;
718+
}
719+
}
720+
721+
if (identifier?.type && Type.isMethod(identifier.type)) {
722+
const methodType = identifier.type as Type.Method;
723+
expectedDeclaringType = Type.FullyQualified.getFullyQualifiedName(methodType.declaringType);
724+
if (expectedDeclaringType) {
725+
break; // Found it!
726+
}
727+
}
728+
}
729+
}
730+
}
731+
732+
if (expectedDeclaringType) {
733+
break; // No need to scan more imports
673734
}
674-
usedImports.get(moduleName)!.add(methodType.name);
675735
}
676-
};
736+
}
737+
738+
// Step 2: Look for references that match
739+
const targetName = this.alias || this.member;
740+
let found = false;
741+
742+
// If no existing imports from this module, look for unresolved references
743+
// If there ARE existing imports, look for references with the expected declaring type
677744

678-
// Create a visitor to collect used identifiers with their type attribution
679745
const collector = new class extends JavaScriptVisitor<void> {
680746
override async visitIdentifier(identifier: J.Identifier, p: void): Promise<J | undefined> {
681-
const type = identifier.type;
682-
if (type && Type.isMethod(type)) {
683-
recordMethodUsage(type as Type.Method);
747+
if (identifier.simpleName === targetName) {
748+
const type = identifier.type;
749+
if (expectedDeclaringType) {
750+
// We have an expected declaring type - check for exact match
751+
if (type && Type.isMethod(type)) {
752+
const methodType = type as Type.Method;
753+
const declaringTypeName = Type.FullyQualified.getFullyQualifiedName(methodType.declaringType);
754+
if (declaringTypeName === expectedDeclaringType) {
755+
found = true;
756+
}
757+
}
758+
} else {
759+
// No existing imports - look for unresolved references (no type)
760+
if (!type) {
761+
found = true;
762+
}
763+
}
684764
}
685765
return super.visitIdentifier(identifier, p);
686766
}
687767

688768
override async visitMethodInvocation(methodInvocation: J.MethodInvocation, p: void): Promise<J | undefined> {
689-
if (methodInvocation.methodType) {
690-
recordMethodUsage(methodInvocation.methodType);
769+
if (methodInvocation.methodType && methodInvocation.methodType.name === targetName) {
770+
if (expectedDeclaringType) {
771+
const declaringTypeName = Type.FullyQualified.getFullyQualifiedName(methodInvocation.methodType.declaringType);
772+
if (declaringTypeName === expectedDeclaringType) {
773+
found = true;
774+
}
775+
}
691776
}
692777
return super.visitMethodInvocation(methodInvocation, p);
693778
}
694779

695780
override async visitFunctionCall(functionCall: JS.FunctionCall, p: void): Promise<J | undefined> {
696-
if (functionCall.methodType) {
697-
recordMethodUsage(functionCall.methodType);
781+
if (functionCall.methodType && functionCall.methodType.name === targetName) {
782+
if (expectedDeclaringType) {
783+
const declaringTypeName = Type.FullyQualified.getFullyQualifiedName(functionCall.methodType.declaringType);
784+
if (declaringTypeName === expectedDeclaringType) {
785+
found = true;
786+
}
787+
}
698788
}
699789
return super.visitFunctionCall(functionCall, p);
700790
}
701791

702792
override async visitFieldAccess(fieldAccess: J.FieldAccess, p: void): Promise<J | undefined> {
703793
const type = fieldAccess.type;
704794
if (type && Type.isMethod(type)) {
705-
recordMethodUsage(type as Type.Method);
795+
const methodType = type as Type.Method;
796+
if (methodType.name === targetName) {
797+
if (expectedDeclaringType) {
798+
const declaringTypeName = Type.FullyQualified.getFullyQualifiedName(methodType.declaringType);
799+
if (declaringTypeName === expectedDeclaringType) {
800+
found = true;
801+
}
802+
}
803+
}
706804
}
707805
return super.visitFieldAccess(fieldAccess, p);
708806
}
709807
};
710808

711809
await collector.visit(compilationUnit, undefined);
712810

713-
// For namespace imports (member === '*'), we cannot use type attribution to detect usage
714-
// because the namespace itself is used as an identifier, not individual members.
715-
// We would need to traverse the AST looking for the alias identifier.
716-
// For simplicity, we skip the onlyIfReferenced check for namespace imports.
717-
if (this.member === '*') {
718-
// TODO: Implement proper namespace usage detection by checking if alias identifier is used
719-
return true;
720-
}
721-
722-
// Check if our target import is used based on type attribution
723-
const moduleMembers = usedImports.get(this.module);
724-
if (!moduleMembers) {
725-
return false;
726-
}
727-
728-
// For specific members, check if that member is used; otherwise check if any member is used
729-
return this.member ? moduleMembers.has(this.member) : moduleMembers.size > 0;
811+
return found;
730812
}
731813

732814
/**

0 commit comments

Comments
 (0)
⚡