Make get tree synchronous again

Author: AndreasArvidsson

src/errors.ts

@@ -1,3 +1,5 @@
+import type { Uri } from "vscode";
+
 export class UnsupportedLanguageError extends Error {
   constructor(language: string) {
     super(
@@ -7,10 +9,10 @@ export class UnsupportedLanguageError extends Error {
   }
 }
 
-export class LanguageFailedToLoad extends Error {
+export class LanguageStillLoadingError extends Error {
   constructor(language: string) {
-    super(`Language '${language}' failed to load`);
-    this.name = "LanguageFailedToLoad";
+    super(`Language '${language}' is still loading; please wait and try again`);
+    this.name = "LanguageStillLoadingError";
   }
 }
 
@@ -20,3 +22,17 @@ export class DeprecatedError extends Error {
     this.name = "DeprecatedError";
   }
 }
+
+export class FailedToParseError extends Error {
+  constructor(uri: Uri) {
+    super(`Failed to parse document: ${uri.toString()}`);
+    this.name = "FailedToParseError";
+  }
+}
+
+export class DocumentNotOpenError extends Error {
+  constructor(uri: Uri) {
+    super(`Document not open: ${uri.toString()}`);
+    this.name = "DocumentNotOpenError";
+  }
+}

src/extension.ts

@@ -14,12 +14,17 @@ import {
 } from "./disabledLanguages";
 import {
   DeprecatedError,
-  LanguageFailedToLoad,
+  LanguageStillLoadingError,
   UnsupportedLanguageError,
 } from "./errors";
 import { languages } from "./languages";
 import { Trees } from "./Trees";
-import { getWasmPath, isDocumentVisible } from "./utils";
+import {
+  getOpenDocument,
+  getWasmPath,
+  isDocumentVisible,
+  parseDocument,
+} from "./utils";
 
 // For some reason this crashes if we put it inside activate
 // Fix: this isn't a field, suppress package member coloring like Go
@@ -92,65 +97,60 @@ export function activate(context: ExtensionContext) {
       throw new Error(`No parser for language ${document.languageId}`);
     }
 
-    tree = language.parser?.parse(document.getText()) ?? undefined;
-
-    if (tree == null) {
-      throw Error(`Failed to parse ${uriString}`);
-    }
+    tree = parseDocument(language.parser, document);
 
     trees.set(uriString, tree);
 
     return tree;
   }
 
-  /**
-   * Create a tree-sitter query for a given language and query source
-   * @param languageId the vscode language id of the language to create the query for
-   * @param source the source of the query
-   * @returns the created query, or undefined if the language couldn't be loaded
-   */
-  function createQuery(languageId: string, source: string): Query | undefined {
-    const language = languages[languageId]?.parser?.language;
-    if (language == null) {
-      throwIfLanguageIsDisabled(languageId);
-      return undefined;
-    }
-    return new Query(language, source);
-  }
-
   /**
    * Get the parse tree for a given document, parsing it if necessary
    * @param document the document to get the tree for
    * @returns the parse tree for the document
    */
-  async function getTreeForUri(uri: Uri): Promise<Tree> {
+  function getTreeForUri(uri: Uri): Tree {
     const uriString = uri.toString();
     let tree = trees.get(uriString);
 
+    // Document is already opened
     if (tree != null) {
       return tree;
     }
 
-    const document = workspace.textDocuments.find(
-      (textDocument) => textDocument.uri.toString() === uri.toString(),
-    );
-
-    if (document == null) {
-      throw new Error(`Document ${uriString} is not open`);
-    }
-
-    tree = await openDocument(document);
+    const document = getOpenDocument(uri);
+    const language = languages[document.languageId];
 
-    if (tree != null) {
-      return tree;
+    // Language without a parser, e.g. plaintext
+    if (language == null) {
+      throw new UnsupportedLanguageError(document.languageId);
     }
 
-    if (document.languageId in languages) {
+    // Language definition exists, but the parser is not loaded. Could be a race
+    // condition or a disabled language.
+    if (language.parser == null) {
       throwIfLanguageIsDisabled(document.languageId);
-      throw new LanguageFailedToLoad(document.languageId);
+      throw new LanguageStillLoadingError(document.languageId);
     }
 
-    throw new UnsupportedLanguageError(document.languageId);
+    tree = parseDocument(language.parser, document);
+    trees.set(uriString, tree);
+    return tree;
+  }
+
+  /**
+   * Create a tree-sitter query for a given language and query source
+   * @param languageId the vscode language id of the language to create the query for
+   * @param source the source of the query
+   * @returns the created query, or undefined if the language couldn't be loaded
+   */
+  function createQuery(languageId: string, source: string): Query | undefined {
+    const language = languages[languageId]?.parser?.language;
+    if (language == null) {
+      throwIfLanguageIsDisabled(languageId);
+      return undefined;
+    }
+    return new Query(language, source);
   }
 
   // NOTE: if you make this an async function, it seems to cause edit anomalies

src/utils.ts

@@ -1,7 +1,9 @@
 import * as fs from "node:fs";
 import * as path from "node:path";
-import type { TextDocument } from "vscode";
-import { window } from "vscode";
+import type { TextDocument, Uri } from "vscode";
+import { window, workspace } from "vscode";
+import type { Parser, Tree } from "web-tree-sitter";
+import { DocumentNotOpenError, FailedToParseError } from "./errors";
 
 export function isDocumentVisible(document: TextDocument): boolean {
   const uriString = document.uri.toString();
@@ -10,6 +12,19 @@ export function isDocumentVisible(document: TextDocument): boolean {
   );
 }
 
+export function getOpenDocument(uri: Uri): TextDocument {
+  const uriString = uri.toString();
+  const document = workspace.textDocuments.find(
+    (doc) => doc.uri.toString() === uriString,
+  );
+
+  if (document == null) {
+    throw new DocumentNotOpenError(uri);
+  }
+
+  return document;
+}
+
 export function getWasmPath(extensionPath: string, moduleName: string): string {
   const absolute = path.join(extensionPath, "parsers", moduleName + ".wasm");
 
@@ -19,3 +34,13 @@ export function getWasmPath(extensionPath: string, moduleName: string): string {
 
   return absolute;
 }
+
+export function parseDocument(parser: Parser, document: TextDocument): Tree {
+  const tree = parser.parse(document.getText());
+
+  if (tree == null) {
+    throw new FailedToParseError(document.uri);
+  }
+
+  return tree;
+}