From 5293519470cd77f31911e43676c445d1764fb7cb Mon Sep 17 00:00:00 2001
From: Matthew Phillips <matthew@matthewphillips.info>
Date: Mon, 9 Aug 2021 12:03:32 -0400
Subject: [PATCH] Fix GoToDefinition for component imports (and component
 usage) (#1045)

* Fix GoToDefinition for component imports (and component usage)

* Remove extra comment
---
 .vscode/launch.json                           |  6 +-
 .../plugins/typescript/DocumentSnapshot.ts    | 20 +++++-
 .../plugins/typescript/TypeScriptPlugin.ts    | 72 +++++++++++++++++--
 .../src/plugins/typescript/utils.ts           |  2 +-
 4 files changed, 88 insertions(+), 12 deletions(-)

diff --git a/.vscode/launch.json b/.vscode/launch.json
index 531210a57..b575d4b70 100644
--- a/.vscode/launch.json
+++ b/.vscode/launch.json
@@ -7,8 +7,8 @@
       "type": "extensionHost",
       "request": "launch",
       "runtimeExecutable": "${execPath}",
-      "args": ["--extensionDevelopmentPath=${workspaceRoot}/tools/astro-vscode"],
-      "outFiles": ["${workspaceRoot}/tools/astro-vscode/dist/**/*.js"]
+      "args": ["--extensionDevelopmentPath=${workspaceRoot}/tools/vscode"],
+      "outFiles": ["${workspaceRoot}/tools/vscode/dist/**/*.js"]
     },
     {
       "type": "node",
@@ -16,7 +16,7 @@
       "name": "Attach to Server",
       "port": 6040,
       "restart": true,
-      "outFiles": ["${workspaceRoot}/tools/astro-languageserver/dist/**/*.js"],
+      "outFiles": ["${workspaceRoot}/tools/languageserver/dist/**/*.js"],
       "skipFiles": ["<node_internals>/**"]
     }
   ],
diff --git a/tools/language-server/src/plugins/typescript/DocumentSnapshot.ts b/tools/language-server/src/plugins/typescript/DocumentSnapshot.ts
index 9e2e778c6..606db8d58 100644
--- a/tools/language-server/src/plugins/typescript/DocumentSnapshot.ts
+++ b/tools/language-server/src/plugins/typescript/DocumentSnapshot.ts
@@ -5,6 +5,8 @@ import { isInTag, positionAt, offsetAt } from '../../core/documents/utils';
 import { pathToUrl } from '../../utils';
 import { getScriptKindFromFileName, isAstroFilePath, toVirtualAstroFilePath } from './utils';
 
+const FILLER_DEFAULT_EXPORT = `\nexport default function() { return ''; };`;
+
 /**
  * The mapper to get from original snapshot positions to generated and vice versa.
  */
@@ -66,7 +68,15 @@ class AstroDocumentSnapshot implements DocumentSnapshot {
   }
 
   get text() {
-    return this.doc.getText();
+    let raw = this.doc.getText();
+    return this.transformContent(raw);
+  }
+
+  /** @internal */
+  private transformContent(content: string) {
+    return content.replace(/---/g, '///') +
+    // TypeScript needs this to know there's a default export.
+    FILLER_DEFAULT_EXPORT;
   }
 
   get filePath() {
@@ -128,7 +138,9 @@ export class DocumentFragmentSnapshot implements Omit<DocumentSnapshot, 'getFrag
 
   /** @internal */
   private transformContent(content: string) {
-    return content.replace(/---/g, '///');
+    return content.replace(/---/g, '///') +
+    // TypeScript needs this to know there's a default export.
+    FILLER_DEFAULT_EXPORT;
   }
 
   getText(start: number, end: number) {
@@ -214,6 +226,10 @@ export class TypeScriptDocumentSnapshot implements DocumentSnapshot {
     return this as unknown as any;
   }
 
+  getOriginalPosition(pos: Position): Position {
+    return pos;
+  }
+
   destroyFragment() {
     // nothing to clean up
   }
diff --git a/tools/language-server/src/plugins/typescript/TypeScriptPlugin.ts b/tools/language-server/src/plugins/typescript/TypeScriptPlugin.ts
index 10b94cb83..918856fb8 100644
--- a/tools/language-server/src/plugins/typescript/TypeScriptPlugin.ts
+++ b/tools/language-server/src/plugins/typescript/TypeScriptPlugin.ts
@@ -1,15 +1,21 @@
+import { join as pathJoin, dirname as pathDirname } from 'path';
 import { Document, DocumentManager, isInsideFrontmatter } from '../../core/documents';
 import type { ConfigManager } from '../../core/config';
 import type { CompletionsProvider, AppCompletionItem, AppCompletionList } from '../interfaces';
+import { SourceFile, ImportDeclaration, Node, SyntaxKind } from 'typescript';
 import { CompletionContext, DefinitionLink, FileChangeType, Position, LocationLink } from 'vscode-languageserver';
 import * as ts from 'typescript';
 import { CompletionsProviderImpl, CompletionEntryWithIdentifer } from './features/CompletionsProvider';
 import { LanguageServiceManager } from './LanguageServiceManager';
 import { SnapshotManager } from './SnapshotManager';
-import { convertToLocationRange, isVirtualFilePath, getScriptKindFromFileName } from './utils';
+import { convertToLocationRange, isVirtualAstroFilePath, isVirtualFilePath, getScriptKindFromFileName } from './utils';
 import { isNoTextSpanInGeneratedCode, SnapshotFragmentMap } from './features/utils';
 import { isNotNullOrUndefined, pathToUrl } from '../../utils';
 
+type BetterTS = typeof ts & {
+  getTouchingPropertyName(sourceFile: SourceFile, pos: number): Node;
+};
+
 export class TypeScriptPlugin implements CompletionsProvider {
   private readonly docManager: DocumentManager;
   private readonly configManager: ConfigManager;
@@ -46,26 +52,38 @@ export class TypeScriptPlugin implements CompletionsProvider {
     const filePath = tsDoc.filePath;
     const tsFilePath = filePath.endsWith('.ts') ? filePath : filePath + '.ts';
 
-    const defs = lang.getDefinitionAndBoundSpan(tsFilePath, mainFragment.offsetAt(mainFragment.getGeneratedPosition(position)));
+    const fragmentPosition = mainFragment.getGeneratedPosition(position);
+    const fragmentOffset = mainFragment.offsetAt(fragmentPosition);
+
+    let defs = lang.getDefinitionAndBoundSpan(tsFilePath, fragmentOffset);
 
     if (!defs || !defs.definitions) {
       return [];
     }
 
+    // Resolve all imports if we can
+    if(this.goToDefinitionFoundOnlyAlias(tsFilePath, defs.definitions!)) {
+      let importDef = this.getGoToDefinitionRefsForImportSpecifier(tsFilePath, fragmentOffset, lang);
+      if(importDef) {
+        defs = importDef;
+      }
+    }
+
     const docs = new SnapshotFragmentMap(this.languageServiceManager);
     docs.set(tsDoc.filePath, { fragment: mainFragment, snapshot: tsDoc });
 
     const result = await Promise.all(
-      defs.definitions.map(async (def) => {
+      defs.definitions!.map(async (def) => {
         const { fragment, snapshot } = await docs.retrieve(def.fileName);
 
         if (isNoTextSpanInGeneratedCode(snapshot.getFullText(), def.textSpan)) {
           const fileName = isVirtualFilePath(def.fileName) ? def.fileName.substr(0, def.fileName.length - 3) : def.fileName;
+          const textSpan = isVirtualAstroFilePath(tsFilePath) ? { start: 0, length: 0 } : def.textSpan;
           return LocationLink.create(
             pathToUrl(fileName),
-            convertToLocationRange(fragment, def.textSpan),
-            convertToLocationRange(fragment, def.textSpan),
-            convertToLocationRange(mainFragment, defs.textSpan)
+            convertToLocationRange(fragment, textSpan),
+            convertToLocationRange(fragment, textSpan),
+            convertToLocationRange(mainFragment, defs!.textSpan)
           );
         }
       })
@@ -110,4 +128,46 @@ export class TypeScriptPlugin implements CompletionsProvider {
   private isInsideFrontmatter(document: Document, position: Position) {
     return isInsideFrontmatter(document.getText(), document.offsetAt(position));
   }
+
+  private goToDefinitionFoundOnlyAlias(tsFileName: string, defs: readonly ts.DefinitionInfo[]) {
+    return !!(defs.length === 1 &&
+      defs[0].kind === 'alias' &&
+      defs[0].fileName === tsFileName);
+  }
+
+  private getGoToDefinitionRefsForImportSpecifier(tsFilePath: string, offset: number, lang: ts.LanguageService): ts.DefinitionInfoAndBoundSpan | undefined {
+    const program = lang.getProgram();
+    const sourceFile = program?.getSourceFile(tsFilePath);
+    if (sourceFile) {
+      let node = (ts as BetterTS).getTouchingPropertyName(sourceFile, offset);
+      if(node && node.kind === SyntaxKind.Identifier) {
+        if(node.parent.kind === SyntaxKind.ImportClause) {
+          let decl = node.parent.parent as ImportDeclaration;
+          let spec = ts.isStringLiteral(decl.moduleSpecifier) && decl.moduleSpecifier.text;
+          if(spec) {
+            let fileName = pathJoin(pathDirname(tsFilePath), spec);
+            let start = node.pos + 1;
+            let def: ts.DefinitionInfoAndBoundSpan = {
+              definitions: [{
+                kind: 'alias',
+                fileName,
+                name: '',
+                containerKind: '',
+                containerName: '',
+                textSpan: {
+                  start: 0,
+                  length: 0
+                }
+              } as ts.DefinitionInfo],
+              textSpan: {
+                start,
+                length: node.end - start
+              }
+            };
+            return def;
+          }
+        }
+      }
+    }
+  }
 }
diff --git a/tools/language-server/src/plugins/typescript/utils.ts b/tools/language-server/src/plugins/typescript/utils.ts
index 9acbe2ed8..d84f35da1 100644
--- a/tools/language-server/src/plugins/typescript/utils.ts
+++ b/tools/language-server/src/plugins/typescript/utils.ts
@@ -190,7 +190,7 @@ export function ensureRealAstroFilePath(filePath: string) {
 }
 
 export function ensureRealFilePath(filePath: string) {
-  return isVirtualFilePath(filePath) ? filePath.slice(0, 3) : filePath;
+  return isVirtualFilePath(filePath) ? filePath.slice(0, filePath.length - 3) : filePath;
 }
 
 export function findTsConfigPath(fileName: string, rootUris: string[]) {