Merge #3816

3816: vscode: add goto ast node definition from rust source code r=Veetaha a=Veetaha

By holding the `Ctrl` key you can now goto-definition of the appropriate syntax token in the syntax tree read-only editor. But actually going to the definition is not very convenient, since it opens the new editor, you'd rather just hold the `Ctrl` and look at the syntax tree because it is automatically scrolled to the proper node and the node itself is enclosed with text selection.

Unfortunately, the algorithm is very simple (because we don't do any elaborate parsing of the syntax tree text received from the server), but it is enough to debug not very large source files.
I tested the performance and in a bad case (rust source file with 5K lines of code) it takes `1.3` seconds to build the `rust -> ast` mapping index (lazily once on the first goto definition request) and each lookup in this worst-case is approx `20-120` ms. I think this is good enough. In the simple case where the file is < 100 lines of code, it is instant.

One peculiarity that I've noticed is that vscode doesn't trigger the goto-definition provider when the user triggers it on some punctuation characters (i.e. it doesn't underline them and invoke te goto-definition provider), but if you explicitly click `Ctrl+LMB` it will only then invoke the provider and navigate to the definition in a new editor. I think this is fine ;D

![rust2ast](https://user-images.githubusercontent.com/36276403/78198718-24d1d500-7492-11ea-91f6-2687cedf26ee.gif)


Related: #3682 

Co-authored-by: veetaha <[email protected]>

1 files changed, 86 insertions, 14 deletions

diff --git a/editors/code/src/commands/syntax_tree.ts b/editors/code/src/commands/syntax_tree.ts
index 8d71cb39e..b7a397414 100644
--- a/editors/code/src/commands/syntax_tree.ts
+++ b/editors/code/src/commands/syntax_tree.ts
@@ -82,8 +82,8 @@ class TextDocumentContentProvider implements vscode.TextDocumentContentProvider
 
 // FIXME: consider implementing this via the Tree View API?
 // https://code.visualstudio.com/api/extension-guides/tree-view
-class AstInspector implements vscode.HoverProvider, Disposable {
-    private static readonly astDecorationType = vscode.window.createTextEditorDecorationType({
+class AstInspector implements vscode.HoverProvider, vscode.DefinitionProvider, Disposable {
+    private readonly astDecorationType = vscode.window.createTextEditorDecorationType({
         borderColor: new vscode.ThemeColor('rust_analyzer.syntaxTreeBorder'),
         borderStyle: "solid",
         borderWidth: "2px",
@@ -91,9 +91,32 @@ class AstInspector implements vscode.HoverProvider, Disposable {
     });
     private rustEditor: undefined | RustEditor;
 
+    // Lazy rust token range -> syntax tree file range.
+    private readonly rust2Ast = new Lazy(() => {
+        const astEditor = this.findAstTextEditor();
+        if (!this.rustEditor || !astEditor) return undefined;
+
+        const buf: [vscode.Range, vscode.Range][] = [];
+        for (let i = 0; i < astEditor.document.lineCount; ++i) {
+            const astLine = astEditor.document.lineAt(i);
+
+            // Heuristically look for nodes with quoted text (which are token nodes)
+            const isTokenNode = astLine.text.lastIndexOf('"') >= 0;
+            if (!isTokenNode) continue;
+
+            const rustRange = this.parseRustTextRange(this.rustEditor.document, astLine.text);
+            if (!rustRange) continue;
+
+            buf.push([rustRange, this.findAstNodeRange(astLine)]);
+        }
+        return buf;
+    });
+
     constructor(ctx: Ctx) {
         ctx.pushCleanup(vscode.languages.registerHoverProvider({ scheme: AST_FILE_SCHEME }, this));
+        ctx.pushCleanup(vscode.languages.registerDefinitionProvider({ language: "rust" }, this));
         vscode.workspace.onDidCloseTextDocument(this.onDidCloseTextDocument, this, ctx.subscriptions);
+        vscode.workspace.onDidChangeTextDocument(this.onDidChangeTextDocument, this, ctx.subscriptions);
         vscode.window.onDidChangeVisibleTextEditors(this.onDidChangeVisibleTextEditors, this, ctx.subscriptions);
 
         ctx.pushCleanup(this);
@@ -102,6 +125,12 @@ class AstInspector implements vscode.HoverProvider, Disposable {
         this.setRustEditor(undefined);
     }
 
+    private onDidChangeTextDocument(event: vscode.TextDocumentChangeEvent) {
+        if (this.rustEditor && event.document.uri.toString() === this.rustEditor.document.uri.toString()) {
+            this.rust2Ast.reset();
+        }
+    }
+
     private onDidCloseTextDocument(doc: vscode.TextDocument) {
         if (this.rustEditor && doc.uri.toString() === this.rustEditor.document.uri.toString()) {
             this.setRustEditor(undefined);
@@ -109,38 +138,67 @@ class AstInspector implements vscode.HoverProvider, Disposable {
     }
 
     private onDidChangeVisibleTextEditors(editors: vscode.TextEditor[]) {
-        if (editors.every(suspect => suspect.document.uri.scheme !== AST_FILE_SCHEME)) {
+        if (!this.findAstTextEditor()) {
             this.setRustEditor(undefined);
             return;
         }
         this.setRustEditor(editors.find(isRustEditor));
     }
 
+    private findAstTextEditor(): undefined | vscode.TextEditor {
+        return vscode.window.visibleTextEditors.find(it => it.document.uri.scheme === AST_FILE_SCHEME);
+    }
+
     private setRustEditor(newRustEditor: undefined | RustEditor) {
-        if (newRustEditor !== this.rustEditor) {
-            this.rustEditor?.setDecorations(AstInspector.astDecorationType, []);
+        if (this.rustEditor && this.rustEditor !== newRustEditor) {
+            this.rustEditor.setDecorations(this.astDecorationType, []);
+            this.rust2Ast.reset();
         }
         this.rustEditor = newRustEditor;
     }
 
+    // additional positional params are omitted
+    provideDefinition(doc: vscode.TextDocument, pos: vscode.Position): vscode.ProviderResult<vscode.DefinitionLink[]> {
+        if (!this.rustEditor || doc.uri.toString() !== this.rustEditor.document.uri.toString()) return;
+
+        const astEditor = this.findAstTextEditor();
+        if (!astEditor) return;
+
+        const rust2AstRanges = this.rust2Ast.get()?.find(([rustRange, _]) => rustRange.contains(pos));
+        if (!rust2AstRanges) return;
+
+        const [rustFileRange, astFileRange] = rust2AstRanges;
+
+        astEditor.revealRange(astFileRange);
+        astEditor.selection = new vscode.Selection(astFileRange.start, astFileRange.end);
+
+        return [{
+            targetRange: astFileRange,
+            targetUri: astEditor.document.uri,
+            originSelectionRange: rustFileRange,
+            targetSelectionRange: astFileRange,
+        }];
+    }
+
+    // additional positional params are omitted
     provideHover(doc: vscode.TextDocument, hoverPosition: vscode.Position): vscode.ProviderResult<vscode.Hover> {
         if (!this.rustEditor) return;
 
-        const astTextLine = doc.lineAt(hoverPosition.line);
+        const astFileLine = doc.lineAt(hoverPosition.line);
 
-        const rustTextRange = this.parseRustTextRange(this.rustEditor.document, astTextLine.text);
-        if (!rustTextRange) return;
+        const rustFileRange = this.parseRustTextRange(this.rustEditor.document, astFileLine.text);
+        if (!rustFileRange) return;
 
-        this.rustEditor.setDecorations(AstInspector.astDecorationType, [rustTextRange]);
-        this.rustEditor.revealRange(rustTextRange);
+        this.rustEditor.setDecorations(this.astDecorationType, [rustFileRange]);
+        this.rustEditor.revealRange(rustFileRange);
 
-        const rustSourceCode = this.rustEditor.document.getText(rustTextRange);
-        const astTextRange = this.findAstRange(astTextLine);
+        const rustSourceCode = this.rustEditor.document.getText(rustFileRange);
+        const astFileRange = this.findAstNodeRange(astFileLine);
 
-        return new vscode.Hover(["```rust\n" + rustSourceCode + "\n```"], astTextRange);
+        return new vscode.Hover(["```rust\n" + rustSourceCode + "\n```"], astFileRange);
     }
 
-    private findAstRange(astLine: vscode.TextLine) {
+    private findAstNodeRange(astLine: vscode.TextLine) {
         const lineOffset = astLine.range.start;
         const begin = lineOffset.translate(undefined, astLine.firstNonWhitespaceCharacterIndex);
         const end = lineOffset.translate(undefined, astLine.text.trimEnd().length);
@@ -156,3 +214,17 @@ class AstInspector implements vscode.HoverProvider, Disposable {
         return new vscode.Range(begin, end);
     }
 }
+
+class Lazy<T> {
+    val: undefined | T;
+
+    constructor(private readonly compute: () => undefined | T) { }
+
+    get() {
+        return this.val ?? (this.val = this.compute());
+    }
+
+    reset() {
+        this.val = undefined;
+    }
+}