Introduce LanguageService.updateIsDefinitionOfReferencedSymbols

amcasey · amcasey · commit 389bd25a277f · 2022-04-08T17:00:57.000-07:00
...to allow use of the checker when computing isDefinition across projects.
diff --git a/src/harness/client.ts b/src/harness/client.ts
@@ -849,6 +849,10 @@ namespace ts.server {
             throw new Error("Program objects are not serializable through the server protocol.");
         }
 
+        updateIsDefinitionOfReferencedSymbols(_referencedSymbols: readonly ReferencedSymbol[], _knownSymbolSpans: Set<DocumentSpan>): boolean {
+            return notImplemented();
+        }
+
         getNonBoundSourceFile(_fileName: string): SourceFile {
             throw new Error("SourceFile objects are not serializable through the server protocol.");
         }
diff --git a/src/harness/harnessLanguageService.ts b/src/harness/harnessLanguageService.ts
@@ -630,6 +630,9 @@ namespace Harness.LanguageService {
         getAutoImportProvider(): ts.Program | undefined {
             throw new Error("Program can not be marshaled across the shim layer.");
         }
+        updateIsDefinitionOfReferencedSymbols(_referencedSymbols: readonly ts.ReferencedSymbol[], _knownSymbolSpans: ts.Set<ts.DocumentSpan>): boolean {
+            return ts.notImplemented();
+        }
         getNonBoundSourceFile(): ts.SourceFile {
             throw new Error("SourceFile can not be marshaled across the shim layer.");
         }
diff --git a/src/server/session.ts b/src/server/session.ts
@@ -330,7 +330,7 @@ namespace ts.server {
         perProjectResults.forEach((projectResults, project) => {
             for (const result of projectResults) {
                 // If there's a mapped location, it'll appear in the results for another project
-                if (!seen.has(result) && !getMappedLocation(documentSpanLocation(result), project)) {
+                if (!seen.has(result) && !getMappedLocationForProject(documentSpanLocation(result), project)) {
                     results.push(result);
                     seen.add(result);
                 }
@@ -368,7 +368,7 @@ namespace ts.server {
             },
         );
 
-        // No re-mapping is required if there's exactly one project
+        // No re-mapping or isDefinition updatses are required if there's exactly one project
         if (isArray(perProjectResults)) {
             return perProjectResults;
         }
@@ -378,18 +378,66 @@ namespace ts.server {
 
         const results: ReferencedSymbol[] = [];
 
-        const clearIsDefinition = perProjectResults.get(defaultProject)![0].references[0].isDefinition === undefined;
+        const defaultProjectResults = perProjectResults.get(defaultProject)!;
+        if (defaultProjectResults[0].references[0].isDefinition === undefined) {
+            // Clear all isDefinition properties
+            perProjectResults.forEach(projectResults => {
+                for (const referencedSymbol of projectResults) {
+                    for (const ref of referencedSymbol.references) {
+                        delete ref.isDefinition;
+                    }
+                }
+            });
+        }
+        else {
+            // Correct isDefinition properties from projects other than defaultProject
+            const knownSymbolSpans = createDocumentSpanSet();
+            for (const referencedSymbol of defaultProjectResults) {
+                for (const ref of referencedSymbol.references) {
+                    if (ref.isDefinition) {
+                        knownSymbolSpans.add(ref);
+                        // One is enough - updateIsDefinitionOfReferencedSymbols will fill out the set based on symbols
+                        break;
+                    }
+                }
+            }
+
+            const updatedProjects = new Set<Project>();
+            while (true) {
+                let progress = false;
+                perProjectResults.forEach((referencedSymbols, project) => {
+                    if (updatedProjects.has(project)) return;
+                    const updated = project.getLanguageService().updateIsDefinitionOfReferencedSymbols(referencedSymbols, knownSymbolSpans);
+                    if (updated) {
+                        updatedProjects.add(project);
+                        progress = true;
+                    }
+                });
+                if (!progress) break;
+            }
+
+            perProjectResults.forEach((referencedSymbols, project) => {
+                if (updatedProjects.has(project)) return;
+                for (const referencedSymbol of referencedSymbols) {
+                    for (const ref of referencedSymbol.references) {
+                        ref.isDefinition = false;
+                    }
+                }
+            });
+        }
+
+        // TODO (acasey): skip this aggregation when simplifying results?
 
         perProjectResults.forEach((projectResults, project) => {
             for (const referencedSymbol of projectResults) {
-                const mappedDefinitionFile = getMappedLocation(documentSpanLocation(referencedSymbol.definition), project);
+                const mappedDefinitionFile = getMappedLocationForProject(documentSpanLocation(referencedSymbol.definition), project);
                 const definition: ReferencedSymbolDefinitionInfo = mappedDefinitionFile === undefined ?
                     referencedSymbol.definition :
                     {
                         ...referencedSymbol.definition,
                         textSpan: createTextSpan(mappedDefinitionFile.pos, referencedSymbol.definition.textSpan.length), // Why would the length be the same in the original?
                         fileName: mappedDefinitionFile.fileName,
-                        contextSpan: getMappedContextSpan(referencedSymbol.definition, project)
+                        contextSpan: getMappedContextSpanForProject(referencedSymbol.definition, project)
                     };
 
                 let symbolToAddTo = find(results, o => documentSpansEqual(o.definition, definition));
@@ -399,10 +447,7 @@ namespace ts.server {
                 }
 
                 for (const ref of referencedSymbol.references) {
-                    if (!contains(symbolToAddTo.references, ref, documentSpansEqual) && !getMappedLocation(documentSpanLocation(ref), project)) {
-                        if (clearIsDefinition) {
-                            delete ref.isDefinition;
-                        }
+                    if (!contains(symbolToAddTo.references, ref, documentSpansEqual) && !getMappedLocationForProject(documentSpanLocation(ref), project)) {
                         symbolToAddTo.references.push(ref);
                     }
                 }
@@ -603,39 +648,16 @@ namespace ts.server {
         return { fileName, pos: textSpan.start };
     }
 
-    function getMappedLocation(location: DocumentPosition, project: Project): DocumentPosition | undefined {
-        const mapsTo = project.getSourceMapper().tryGetSourcePosition(location);
-        return mapsTo && project.projectService.fileExists(toNormalizedPath(mapsTo.fileName)) ? mapsTo : undefined;
+    function getMappedLocationForProject(location: DocumentPosition, project: Project): DocumentPosition | undefined {
+        return getMappedLocation(location, project.getSourceMapper(), p => project.projectService.fileExists(p as NormalizedPath));
     }
 
-    function getMappedDocumentSpan(documentSpan: DocumentSpan, project: Project): DocumentSpan | undefined {
-        const newPosition = getMappedLocation(documentSpanLocation(documentSpan), project);
-        if (!newPosition) return undefined;
-        return {
-            fileName: newPosition.fileName,
-            textSpan: {
-                start: newPosition.pos,
-                length: documentSpan.textSpan.length
-            },
-            originalFileName: documentSpan.fileName,
-            originalTextSpan: documentSpan.textSpan,
-            contextSpan: getMappedContextSpan(documentSpan, project),
-            originalContextSpan: documentSpan.contextSpan
-        };
+    function getMappedDocumentSpanForProject(documentSpan: DocumentSpan, project: Project): DocumentSpan | undefined {
+        return getMappedDocumentSpan(documentSpan, project.getSourceMapper(), p => project.projectService.fileExists(p as NormalizedPath));
     }
 
-    function getMappedContextSpan(documentSpan: DocumentSpan, project: Project): TextSpan | undefined {
-        const contextSpanStart = documentSpan.contextSpan && getMappedLocation(
-            { fileName: documentSpan.fileName, pos: documentSpan.contextSpan.start },
-            project
-        );
-        const contextSpanEnd = documentSpan.contextSpan && getMappedLocation(
-            { fileName: documentSpan.fileName, pos: documentSpan.contextSpan.start + documentSpan.contextSpan.length },
-            project
-        );
-        return contextSpanStart && contextSpanEnd ?
-            { start: contextSpanStart.pos, length: contextSpanEnd.pos - contextSpanStart.pos } :
-            undefined;
+    function getMappedContextSpanForProject(documentSpan: DocumentSpan, project: Project): TextSpan | undefined {
+        return getMappedContextSpan(documentSpan, project.getSourceMapper(), p => project.projectService.fileExists(p as NormalizedPath));
     }
 
     const invalidPartialSemanticModeCommands: readonly CommandNames[] = [
@@ -1249,7 +1271,7 @@ namespace ts.server {
 
         private mapDefinitionInfoLocations(definitions: readonly DefinitionInfo[], project: Project): readonly DefinitionInfo[] {
             return definitions.map((info): DefinitionInfo => {
-                const newDocumentSpan = getMappedDocumentSpan(info, project);
+                const newDocumentSpan = getMappedDocumentSpanForProject(info, project);
                 return !newDocumentSpan ? info : {
                     ...newDocumentSpan,
                     containerKind: info.containerKind,
@@ -1390,7 +1412,7 @@ namespace ts.server {
 
         private mapImplementationLocations(implementations: readonly ImplementationLocation[], project: Project): readonly ImplementationLocation[] {
             return implementations.map((info): ImplementationLocation => {
-                const newDocumentSpan = getMappedDocumentSpan(info, project);
+                const newDocumentSpan = getMappedDocumentSpanForProject(info, project);
                 return !newDocumentSpan ? info : {
                     ...newDocumentSpan,
                     kind: info.kind,
@@ -1626,7 +1648,7 @@ namespace ts.server {
             const symbolStartOffset = nameSpan ? scriptInfo.positionToLineOffset(nameSpan.start).offset : 0;
             const symbolName = nameSpan ? scriptInfo.getSnapshot().getText(nameSpan.start, textSpanEnd(nameSpan)) : "";
             const refs: readonly protocol.ReferencesResponseItem[] = flatMap(references, referencedSymbol => {
-                return referencedSymbol.references.map(entry => referenceEntryToReferencesResponseItem(this.projectService, entry));
+                return referencedSymbol.references.map(entry => referenceEntryToReferencesResponseItem(this.projectService, entry)); // TODO (acasey): de-dup
             });
             return { refs, symbolName, symbolStartOffset, symbolDisplayString };
         }
@@ -2230,7 +2252,7 @@ namespace ts.server {
             // Mutates `outputs`
             function addItemsForProject(project: Project) {
                 const projectItems = project.getLanguageService().getNavigateToItems(searchValue, maxResultCount, /*filename*/ undefined, /*excludeDts*/ project.isNonTsProject());
-                const unseenItems = filter(projectItems, item => tryAddSeenItem(item) && !getMappedLocation(documentSpanLocation(item), project));
+                const unseenItems = filter(projectItems, item => tryAddSeenItem(item) && !getMappedLocationForProject(documentSpanLocation(item), project));
                 if (unseenItems.length) {
                     outputs.push({ project, navigateToItems: unseenItems });
                 }
diff --git a/src/services/findAllReferences.ts b/src/services/findAllReferences.ts
@@ -564,7 +564,7 @@ namespace ts.FindAllReferences {
     }
 
     /** Whether a reference, `node`, is a definition of the `target` symbol */
-    function isDeclarationOfSymbol(node: Node, target: Symbol | undefined): boolean {
+    export function isDeclarationOfSymbol(node: Node, target: Symbol | undefined): boolean {
         if (!target) return false;
         const source = getDeclarationFromName(node) ||
             (node.kind === SyntaxKind.DefaultKeyword ? node.parent
diff --git a/src/services/services.ts b/src/services/services.ts
@@ -1600,6 +1600,62 @@ namespace ts {
             return host.getPackageJsonAutoImportProvider?.();
         }
 
+        function updateIsDefinitionOfReferencedSymbols(referencedSymbols: readonly ReferencedSymbol[], knownSymbolSpans: Set<DocumentSpan>): boolean {
+            const checker = program.getTypeChecker();
+            const symbol = getSymbolForProgram();
+
+            if (!symbol) return false;
+
+            for (const referencedSymbol of referencedSymbols) {
+                for (const ref of referencedSymbol.references) {
+                    const refNode = getNodeForSpan(ref);
+                    Debug.assertIsDefined(refNode);
+                    if (knownSymbolSpans.has(ref) || FindAllReferences.isDeclarationOfSymbol(refNode, symbol)) {
+                        knownSymbolSpans.add(ref);
+                        ref.isDefinition = true;
+                        const mappedSpan = getMappedDocumentSpan(ref, sourceMapper, maybeBind(host, host.fileExists));
+                        if (mappedSpan) {
+                            knownSymbolSpans.add(mappedSpan);
+                        }
+                    }
+                    else {
+                        ref.isDefinition = false;
+                    }
+                }
+            }
+
+            return true;
+
+            function getSymbolForProgram(): Symbol | undefined {
+                for (const referencedSymbol of referencedSymbols) {
+                    for (const ref of referencedSymbol.references) {
+                        if (knownSymbolSpans.has(ref)) {
+                            const refNode = getNodeForSpan(ref);
+                            Debug.assertIsDefined(refNode);
+                            return checker.getSymbolAtLocation(refNode);
+                        }
+                        const mappedSpan = getMappedDocumentSpan(ref, sourceMapper, maybeBind(host, host.fileExists));
+                        if (mappedSpan && knownSymbolSpans.has(mappedSpan)) {
+                            const refNode = getNodeForSpan(mappedSpan);
+                            if (refNode) {
+                                return checker.getSymbolAtLocation(refNode);
+                            }
+                        }
+                    }
+                }
+
+                return undefined;
+            }
+
+            function getNodeForSpan(docSpan: DocumentSpan): Node | undefined {
+                const sourceFile = program.getSourceFile(docSpan.fileName);
+                if (!sourceFile) return undefined;
+                const rawNode = getTouchingPropertyName(sourceFile, docSpan.textSpan.start);
+                const adjustedNode = FindAllReferences.Core.getAdjustedNode(rawNode, { use: FindAllReferences.FindReferencesUse.References });
+                return adjustedNode;
+            }
+        }
+
         function cleanupSemanticCache(): void {
             program = undefined!; // TODO: GH#18217
         }
@@ -2715,6 +2771,7 @@ namespace ts {
             getNonBoundSourceFile,
             getProgram,
             getAutoImportProvider,
+            updateIsDefinitionOfReferencedSymbols,
             getApplicableRefactors,
             getEditsForRefactor,
             toLineColumnOffset,
diff --git a/src/services/types.ts b/src/services/types.ts
@@ -546,6 +546,11 @@ namespace ts {
         /* @internal */ getNonBoundSourceFile(fileName: string): SourceFile;
         /* @internal */ getAutoImportProvider(): Program | undefined;
 
+        /// Returns true if a suitable symbol was found in the project.
+        /// May set isDefinition properties in `referencedSymbols` to false.
+        /// May add elements to `knownSymbolSpans`.
+        /* @internal */ updateIsDefinitionOfReferencedSymbols(referencedSymbols: readonly ReferencedSymbol[], knownSymbolSpans: Set<DocumentSpan>): boolean;
+
         toggleLineComment(fileName: string, textRange: TextRange): TextChange[];
         toggleMultilineComment(fileName: string, textRange: TextRange): TextChange[];
         commentSelection(fileName: string, textRange: TextRange): TextChange[];
diff --git a/src/services/utilities.ts b/src/services/utilities.ts
@@ -2107,6 +2107,48 @@ namespace ts {
         return true;
     }
 
+    export function getMappedLocation(location: DocumentPosition, sourceMapper: SourceMapper, fileExists: ((path: string) => boolean) | undefined): DocumentPosition | undefined {
+        const mapsTo = sourceMapper.tryGetSourcePosition(location);
+        return mapsTo && (!fileExists || fileExists(normalizePath(mapsTo.fileName)) ? mapsTo : undefined);
+    }
+
+    export function getMappedDocumentSpan(documentSpan: DocumentSpan, sourceMapper: SourceMapper, fileExists?: (path: string) => boolean): DocumentSpan | undefined {
+        const { fileName, textSpan } = documentSpan;
+        const newPosition = getMappedLocation({ fileName, pos: textSpan.start }, sourceMapper, fileExists);
+        if (!newPosition) return undefined;
+        const newEndPosition = getMappedLocation({ fileName, pos: textSpan.start + textSpan.length }, sourceMapper, fileExists);
+        const newLength = newEndPosition
+            ? newEndPosition.pos - newPosition.pos
+            : textSpan.length; // This shouldn't happen
+        return {
+            fileName: newPosition.fileName,
+            textSpan: {
+                start: newPosition.pos,
+                length: newLength,
+            },
+            originalFileName: documentSpan.fileName,
+            originalTextSpan: documentSpan.textSpan,
+            contextSpan: getMappedContextSpan(documentSpan, sourceMapper, fileExists),
+            originalContextSpan: documentSpan.contextSpan
+        };
+    }
+
+    export function getMappedContextSpan(documentSpan: DocumentSpan, sourceMapper: SourceMapper, fileExists?: (path: string) => boolean): TextSpan | undefined {
+        const contextSpanStart = documentSpan.contextSpan && getMappedLocation(
+            { fileName: documentSpan.fileName, pos: documentSpan.contextSpan.start },
+            sourceMapper,
+            fileExists
+        );
+        const contextSpanEnd = documentSpan.contextSpan && getMappedLocation(
+            { fileName: documentSpan.fileName, pos: documentSpan.contextSpan.start + documentSpan.contextSpan.length },
+            sourceMapper,
+            fileExists
+        );
+        return contextSpanStart && contextSpanEnd ?
+            { start: contextSpanStart.pos, length: contextSpanEnd.pos - contextSpanStart.pos } :
+            undefined;
+    }
+
     // #endregion
 
     // Display-part writer helpers

Original file line number	Diff line number	Diff line change
`@@ -849,6 +849,10 @@ namespace ts.server {`
`849`	`849`	`throw new Error("Program objects are not serializable through the server protocol.");`
`850`	`850`	`}`
`851`	`851`
	`852`	`+ updateIsDefinitionOfReferencedSymbols(_referencedSymbols: readonly ReferencedSymbol[], _knownSymbolSpans: Set<DocumentSpan>): boolean {`
	`853`	`+ return notImplemented();`
	`854`	`+ }`
	`855`	`+`
`852`	`856`	`getNonBoundSourceFile(_fileName: string): SourceFile {`
`853`	`857`	`throw new Error("SourceFile objects are not serializable through the server protocol.");`
`854`	`858`	`}`
Original file line number	Diff line number	Diff line change
`@@ -630,6 +630,9 @@ namespace Harness.LanguageService {`
`630`	`630`	`getAutoImportProvider(): ts.Program \| undefined {`
`631`	`631`	`throw new Error("Program can not be marshaled across the shim layer.");`
`632`	`632`	`}`
	`633`	`+ updateIsDefinitionOfReferencedSymbols(_referencedSymbols: readonly ts.ReferencedSymbol[], _knownSymbolSpans: ts.Set<ts.DocumentSpan>): boolean {`
	`634`	`+ return ts.notImplemented();`
	`635`	`+ }`
`633`	`636`	`getNonBoundSourceFile(): ts.SourceFile {`
`634`	`637`	`throw new Error("SourceFile can not be marshaled across the shim layer.");`
`635`	`638`	`}`
Original file line number	Diff line number	Diff line change
`@@ -564,7 +564,7 @@ namespace ts.FindAllReferences {`
`564`	`564`	`}`
`565`	`565`
`566`	`566`	/** Whether a reference, `node`, is a definition of the `target` symbol */
`567`		`- function isDeclarationOfSymbol(node: Node, target: Symbol \| undefined): boolean {`
	`567`	`+ export function isDeclarationOfSymbol(node: Node, target: Symbol \| undefined): boolean {`
`568`	`568`	`if (!target) return false;`
`569`	`569`	`const source = getDeclarationFromName(node) \|\|`
`570`	`570`	`(node.kind === SyntaxKind.DefaultKeyword ? node.parent`