Fix definition bugs

Author: wolmir

languageServer/src/LanguageServer.ts

@@ -5,9 +5,15 @@ import {
   _Connection,
   TextDocumentPositionParams,
   CodeActionParams,
-  DefinitionParams
+  DefinitionParams,
+  SymbolKind,
+  Location,
+  Position,
+  Range,
+  DocumentSymbol
 } from 'vscode-languageserver/node'
 import { TextDocument } from 'vscode-languageserver-textdocument'
+import { URI } from 'vscode-uri'
 import { TextDocumentWrapper } from './TextDocumentWrapper'
 
 export class LanguageServer {
@@ -18,7 +24,13 @@ export class LanguageServer {
 
     connection.onInitialize(() => this.onInitialize())
 
-    connection.onDefinition(params => this.onDefinition(params))
+    connection.onDefinition(params => {
+      if (!params.textDocument.uri.endsWith('dvc.yaml')) {
+        return null
+      }
+
+      return this.onDefinition(params)
+    })
 
     this.documents.listen(connection)
 
@@ -41,26 +53,89 @@ export class LanguageServer {
     return new TextDocumentWrapper(doc)
   }
 
+  private getFilePathLocations(
+    symbolUnderCursor: DocumentSymbol,
+    allDocs: TextDocument[]
+  ) {
+    if (symbolUnderCursor.kind !== SymbolKind.File) {
+      return []
+    }
+
+    const filePath = symbolUnderCursor.name
+
+    const matchingFiles = allDocs.filter(doc =>
+      URI.file(doc.uri).fsPath.endsWith(filePath)
+    )
+
+    return matchingFiles.map(doc => {
+      const uri = doc.uri
+      const start = Position.create(0, 0)
+      const end = doc.positionAt(doc.getText().length - 1)
+      const range = Range.create(start, end)
+
+      return Location.create(uri, range)
+    })
+  }
+
+  private getLocationsFromOtherDocuments(
+    symbolUnderCursor: DocumentSymbol,
+    allDocs: TextDocument[]
+  ) {
+    const locationsAccumulator = []
+
+    for (const txtDoc of allDocs) {
+      const finder = this.wrap(txtDoc)
+      const locations = finder.findLocationsFor(symbolUnderCursor)
+      locationsAccumulator.push(...locations)
+    }
+
+    return locationsAccumulator
+  }
+
   private onDefinition(params: DefinitionParams) {
     const document = this.getDvcTextDocument(params)
     const symbolUnderCursor = document?.symbolAt(params.position)
 
-    if (symbolUnderCursor) {
+    if (document && symbolUnderCursor) {
       const allDocs = this.documents.all()
       const locationsAccumulator = []
 
-      for (const txtDoc of allDocs) {
-        const finder = this.wrap(txtDoc)
-        const locations = finder.findLocationsFor(symbolUnderCursor)
-        locationsAccumulator.push(...locations)
+      const fileLocations = this.getFilePathLocations(
+        symbolUnderCursor,
+        allDocs
+      )
+
+      locationsAccumulator.push(...fileLocations)
+
+      const locationsFromOtherDocuments = this.getLocationsFromOtherDocuments(
+        symbolUnderCursor,
+        allDocs
+      )
+
+      locationsAccumulator.push(...locationsFromOtherDocuments)
+
+      const externalLocations = locationsAccumulator.filter(
+        location => location.uri !== document.uri
+      )
+
+      if (externalLocations.length > 0) {
+        return this.arrayOrSingleResponse(externalLocations)
       }
 
-      return locationsAccumulator ?? []
+      return this.arrayOrSingleResponse(locationsAccumulator)
     }
 
     return null
   }
 
+  private arrayOrSingleResponse<T>(elements: T[]) {
+    if (elements.length === 1) {
+      return elements[0]
+    }
+
+    return elements
+  }
+
   private onInitialize() {
     const serverCapabilities: ServerCapabilities = {
       definitionProvider: true

languageServer/src/TextDocumentWrapper.ts

@@ -109,11 +109,17 @@ export class TextDocumentWrapper implements ITextDocumentWrapper {
   ) {
     const nodeValue = `${node.value}`
 
+    let symbolKind: SymbolKind = SymbolKind.String
+
+    if (/\.[A-Za-z]+$/.test(nodeValue)) {
+      symbolKind = SymbolKind.File
+    }
+
     const symbolsSoFar: DocumentSymbol[] = [
       DocumentSymbol.create(
         nodeValue,
         undefined,
-        SymbolKind.String,
+        symbolKind,
         Range.create(this.positionAt(nodeStart), this.positionAt(nodeEnd)),
         Range.create(this.positionAt(nodeStart), this.positionAt(valueEnd))
       ),

languageServer/src/test/definitions.test.ts

@@ -33,16 +33,10 @@ describe('textDocument/definitions', () => {
     const response = await requestDefinitions(dvcYaml, 'auc')
 
     expect(response).toBeTruthy()
-    expect(response).toStrictEqual([
-      {
-        range: Range.create(Position.create(5, 8), Position.create(6, 0)),
-        uri: 'file:///dvc.yaml'
-      },
-      {
-        range: Range.create(Position.create(4, 0), Position.create(4, 3)),
-        uri: 'file:///params.yaml'
-      }
-    ])
+    expect(response).toStrictEqual({
+      range: Range.create(Position.create(4, 0), Position.create(4, 3)),
+      uri: 'file:///params.yaml'
+    })
   })
 
   it('should be able to read a complicated command from a stage', async () => {
@@ -56,11 +50,9 @@ describe('textDocument/definitions', () => {
     const response = await requestDefinitions(dvcYaml, '{item.os}')
 
     expect(response).toBeTruthy()
-    expect(response).toStrictEqual([
-      {
-        range: Range.create(Position.create(15, 8), Position.create(15, 10)),
-        uri: 'file:///dvc.yaml'
-      }
-    ])
+    expect(response).toStrictEqual({
+      range: Range.create(Position.create(15, 8), Position.create(15, 10)),
+      uri: 'file:///dvc.yaml'
+    })
   })
 })