feat: provide AST view (#1617)

* feat: provide AST view

* test: update snapshot

crates/tinymist/src/cmd.rs

@@ -1,9 +1,9 @@
 //! Tinymist LSP commands
 
-use std::ops::Deref;
+use std::ops::{Deref, Range};
 use std::path::PathBuf;
 
-use lsp_types::*;
+use lsp_types::TextDocumentIdentifier;
 use serde::{Deserialize, Serialize};
 use serde_json::Value as JsonValue;
 use sync_ls::RequestId;
@@ -14,14 +14,16 @@ use tinymist_project::{
     ExportTextTask, ExportTransform, PageSelection, Pages, ProjectTask, QueryTask,
 };
 use tinymist_query::package::PackageInfo;
-use tinymist_query::LocalContextGuard;
+use tinymist_query::{LocalContextGuard, LspRange};
 use tinymist_std::error::prelude::*;
 use typst::diag::{eco_format, EcoString, StrResult};
 use typst::syntax::package::{PackageSpec, VersionlessPackageSpec};
+use typst::syntax::{LinkedNode, Source};
 use world::TaskInputs;
 
 use super::*;
 use crate::lsp::query::{run_query, LspClientExt};
+use crate::tool::ast::AstRepr;
 use crate::tool::package::InitTask;
 
 /// See [`ProjectTask`].
@@ -56,8 +58,8 @@ struct QueryOpts {
 
 #[derive(Debug, Clone, Default, Deserialize)]
 #[serde(rename_all = "camelCase")]
-struct HighlightRangeOpts {
-    range: Option<Range>,
+struct ExportSyntaxRangeOpts {
+    range: Option<LspRange>,
 }
 
 /// Here are implemented the handlers for each command.
@@ -212,8 +214,45 @@ impl ServerState {
     /// Export a range of the current document as Ansi highlighted text.
     pub fn export_ansi_hl(&mut self, mut args: Vec<JsonValue>) -> AnySchedulableResponse {
         let path = get_arg!(args[0] as PathBuf);
-        let opts = get_arg_or_default!(args[1] as HighlightRangeOpts);
+        let opts = get_arg_or_default!(args[1] as ExportSyntaxRangeOpts);
 
+        let output = self.select_range(path, opts.range, |source, range| {
+            let mut text_in_range = source.text();
+            if let Some(range) = range {
+                text_in_range = text_in_range
+                    .get(range)
+                    .ok_or_else(|| internal_error("cannot get text in range"))?;
+            }
+
+            typst_ansi_hl::Highlighter::default()
+                .for_discord()
+                .with_soft_limit(2000)
+                .highlight(text_in_range)
+                .map_err(|e| internal_error(format!("cannot highlight: {e}")))
+        })?;
+
+        just_ok(JsonValue::String(output))
+    }
+
+    /// Export a range of the current file's AST.
+    pub fn export_ast(&mut self, mut args: Vec<JsonValue>) -> AnySchedulableResponse {
+        let path = get_arg!(args[0] as PathBuf);
+        let opts = get_arg_or_default!(args[1] as ExportSyntaxRangeOpts);
+
+        let output = self.select_range(path, opts.range, |source, range| {
+            let linked_node = LinkedNode::new(source.root());
+            Ok(format!("{}", AstRepr(linked_node, range)))
+        })?;
+
+        just_ok(JsonValue::String(output))
+    }
+
+    fn select_range<T>(
+        &mut self,
+        path: PathBuf,
+        range: Option<LspRange>,
+        f: impl Fn(Source, Option<Range<usize>>) -> LspResult<T>,
+    ) -> LspResult<T> {
         let s = self
             .query_source(path.into(), Ok)
             .map_err(|e| internal_error(format!("cannot find source: {e}")))?;
@@ -221,28 +260,14 @@ impl ServerState {
         // todo: cannot select syntax-sensitive data well
         // let node = LinkedNode::new(s.root());
 
-        let range = opts
-            .range
+        let range = range
             .map(|r| {
                 tinymist_query::to_typst_range(r, self.const_config().position_encoding, &s)
                     .ok_or_else(|| internal_error("cannoet convert range"))
             })
             .transpose()?;
 
-        let mut text_in_range = s.text();
-        if let Some(range) = range {
-            text_in_range = text_in_range
-                .get(range)
-                .ok_or_else(|| internal_error("cannot get text in range"))?;
-        }
-
-        let output = typst_ansi_hl::Highlighter::default()
-            .for_discord()
-            .with_soft_limit(2000)
-            .highlight(text_in_range)
-            .map_err(|e| internal_error(format!("cannot highlight: {e}")))?;
-
-        just_ok(JsonValue::String(output))
+        f(s, range)
     }
 
     /// Clear all cached resources.

crates/tinymist/src/server.rs

@@ -265,6 +265,7 @@ impl ServerState {
             .with_command_("tinymist.exportMarkdown", State::export_markdown)
             .with_command_("tinymist.exportQuery", State::export_query)
             .with_command("tinymist.exportAnsiHighlight", State::export_ansi_hl)
+            .with_command("tinymist.exportAst", State::export_ast)
             .with_command("tinymist.doClearCache", State::clear_cache)
             .with_command("tinymist.pinMain", State::pin_document)
             .with_command("tinymist.focusMain", State::focus_document)

crates/tinymist/src/tool/ast.rs

@@ -0,0 +1,53 @@
+//! AST introspection tool.
+
+use core::fmt;
+use std::ops::Range;
+
+use typst::syntax::LinkedNode;
+
+pub(crate) struct AstRepr<'a>(pub LinkedNode<'a>, pub Option<Range<usize>>);
+
+impl AstRepr<'_> {
+    fn contains(&self, node: &LinkedNode) -> bool {
+        let rng = self.1.as_ref();
+        rng.is_some_and(|rng| {
+            if rng.start == rng.end {
+                return node.range().start == rng.start && node.range().end == rng.start
+                    || node.range().start < rng.start && rng.start <= node.range().end;
+            }
+
+            !(rng.end <= node.range().start || rng.start >= node.range().end)
+        })
+    }
+
+    fn node(&self, node: &LinkedNode, f: &mut fmt::Formatter<'_>, indent: usize) -> fmt::Result {
+        if !self.contains(node) {
+            return Ok(());
+        }
+
+        write!(f, "{: >indent$}{:?}(", "", node.kind())?;
+
+        if !node.text().is_empty() {
+            write!(f, "{:?}", node.text())?;
+        } else if node.get().children().len() > 0 {
+            write!(f, "{:?}, ", node.children().len())?;
+            f.write_str("{\n")?;
+            for child in node.children() {
+                if !self.contains(&child) {
+                    continue;
+                }
+                self.node(&child, f, indent + 1)?;
+                f.write_str("\n")?;
+            }
+            write!(f, "{: >indent$}}}", "")?;
+        }
+        f.write_str(")")
+    }
+}
+
+impl fmt::Display for AstRepr<'_> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str("#")?;
+        self.node(&self.0, f, 0)
+    }
+}

crates/tinymist/src/tool/mod.rs

@@ -1,5 +1,6 @@
 //! All the language tools provided by the `tinymist` crate.
 
+pub mod ast;
 pub mod package;
 pub mod project;
 pub mod testing;

editors/vscode/package.json

@@ -1100,6 +1100,11 @@
         "title": "%extension.tinymist.command.tinymist.copyAnsiHighlight%",
         "category": "Typst"
       },
+      {
+        "command": "tinymist.viewAst",
+        "title": "%extension.tinymist.command.tinymist.viewAst%",
+        "category": "Typst"
+      },
       {
         "command": "tinymist.showLog",
         "title": "%extension.tinymist.command.tinymist.showLog%",

editors/vscode/src/context.ts

@@ -1,4 +1,5 @@
 import * as vscode from "vscode";
+import { isTypstDocument } from "./util";
 
 /**
  * The active editor owning *typst language document* to track.
@@ -17,8 +18,7 @@ export class IContext {
     // Tracks the active editor owning *typst language document*.
     context.subscriptions.push(
       vscode.window.onDidChangeActiveTextEditor((editor: vscode.TextEditor | undefined) => {
-        const langId = editor?.document.languageId;
-        if (langId === "typst") {
+        if (isTypstDocument(editor?.document)) {
           activeEditor = editor;
         } else if (editor === undefined || activeEditor?.document.isClosed) {
           activeEditor = undefined;
@@ -32,6 +32,10 @@ export class IContext {
     return activeEditor;
   }
 
+  currentActiveEditor(): vscode.TextEditor | undefined {
+    return IContext.currentActiveEditor();
+  }
+
   registerFileLevelCommand(command: IFileLevelCommand) {
     this.fileLevelCodelens.push(command);
     this.subscriptions.push(

editors/vscode/src/extension.ts

@@ -19,7 +19,7 @@ import { LanguageState, tinymist } from "./lsp";
 import { commandCreateLocalPackage, commandOpenLocalPackage } from "./package-manager";
 import { extensionState } from "./state";
 import { triggerStatusBar } from "./ui-extends";
-import { activeTypstEditor } from "./util";
+import { activeTypstEditor, isTypstDocument } from "./util";
 import { LanguageClient } from "vscode-languageclient/node";
 
 import { setIsTinymist as previewSetIsTinymist } from "./features/preview-compat";
@@ -151,11 +151,11 @@ async function languageActivate(context: IContext) {
 
   // Find first document to focus
   const editor = window.activeTextEditor;
-  if (editor?.document.languageId === "typst" && editor.document.uri.fsPath) {
+  if (isTypstDocument(editor?.document)) {
     commandActivateDoc(editor.document);
   } else {
     window.visibleTextEditors.forEach((editor) => {
-      if (editor.document.languageId === "typst" && editor.document.uri.fsPath) {
+      if (isTypstDocument(editor.document)) {
         commandActivateDoc(editor.document);
       }
     });
@@ -166,12 +166,11 @@ async function languageActivate(context: IContext) {
       if (editor?.document.isUntitled) {
         return;
       }
-      const langId = editor?.document.languageId;
       // todo: plaintext detection
       // if (langId === "plaintext") {
       //     console.log("plaintext", langId, editor?.document.uri.fsPath);
       // }
-      if (langId !== "typst") {
+      if (!isTypstDocument(editor?.document)) {
         // console.log("not typst", langId, editor?.document.uri.fsPath);
         return commandActivateDoc(undefined);
       }
@@ -181,7 +180,7 @@ async function languageActivate(context: IContext) {
   context.subscriptions.push(
     vscode.workspace.onDidOpenTextDocument((doc: vscode.TextDocument) => {
       if (doc.isUntitled && window.activeTextEditor?.document === doc) {
-        if (doc.languageId === "typst") {
+        if (isTypstDocument(doc)) {
           return commandActivateDocPath(doc, "/untitled/" + doc.uri.fsPath);
         } else {
           return commandActivateDoc(undefined);
@@ -212,6 +211,7 @@ async function languageActivate(context: IContext) {
     commands.registerCommand("tinymist.clearCache", commandClearCache),
     commands.registerCommand("tinymist.runCodeLens", commandRunCodeLens),
     commands.registerCommand("tinymist.copyAnsiHighlight", commandCopyAnsiHighlight),
+    commands.registerCommand("tinymist.viewAst", commandViewAst(context)),
 
     commands.registerCommand("tinymist.pinMainToCurrent", () => commandPinMain(true)),
     commands.registerCommand("tinymist.unpinMain", () => commandPinMain(false)),
@@ -260,6 +260,14 @@ async function commandGetCurrentDocumentMetrics(): Promise<any> {
   return res;
 }
 
+async function getNonEmptySelection(editor: TextEditor): Promise<any> {
+  if (editor.selection.isEmpty) {
+    return undefined;
+  }
+
+  return (await tinymist.clientPromise).code2ProtocolConverter.asRange(editor.selection);
+}
+
 async function commandCopyAnsiHighlight(): Promise<void> {
   const editor = activeTypstEditor();
   if (editor === undefined) {
@@ -267,7 +275,7 @@ async function commandCopyAnsiHighlight(): Promise<void> {
   }
 
   const res = await tinymist.exportAnsiHighlight(editor.document.uri.fsPath, {
-    range: editor.selection,
+    range: await getNonEmptySelection(editor),
   });
 
   if (res === null) {
@@ -278,6 +286,86 @@ async function commandCopyAnsiHighlight(): Promise<void> {
   await vscode.env.clipboard.writeText(res);
 }
 
+function commandViewAst(ctx: IContext) {
+  const scheme = "tinymist-ast";
+  const uri = `${scheme}://viewAst/ast.typ`;
+
+  const AstDoc = new (class implements vscode.TextDocumentContentProvider {
+    readonly uri = vscode.Uri.parse(uri);
+    readonly eventEmitter = new vscode.EventEmitter<vscode.Uri>();
+    constructor() {
+      vscode.workspace.onDidChangeTextDocument(
+        this.onDidChangeTextDocument,
+        this,
+        ctx.subscriptions,
+      );
+      vscode.window.onDidChangeActiveTextEditor(
+        this.onDidChangeActiveTextEditor,
+        this,
+        ctx.subscriptions,
+      );
+      vscode.window.onDidChangeTextEditorSelection(
+        this.onDidChangeTextSelection,
+        this,
+        ctx.subscriptions,
+      );
+    }
+
+    emitChange() {
+      this.eventEmitter.fire(this.uri);
+    }
+
+    private onDidChangeTextDocument(event: vscode.TextDocumentChangeEvent) {
+      if (isTypstDocument(event.document)) {
+        // We need to order this after language server updates, but there's no API for that.
+        // Hence, good old sleep().
+        setTimeout(() => this.emitChange(), 10);
+      }
+    }
+
+    private onDidChangeActiveTextEditor(editor: vscode.TextEditor | undefined) {
+      if (editor && isTypstDocument(editor.document)) {
+        this.emitChange();
+      }
+    }
+
+    private onDidChangeTextSelection(event: vscode.TextEditorSelectionChangeEvent) {
+      if (isTypstDocument(event.textEditor.document)) {
+        this.emitChange();
+      }
+    }
+
+    async provideTextDocumentContent(
+      _uri: vscode.Uri,
+      _ct: vscode.CancellationToken,
+    ): Promise<string> {
+      const editor = ctx.currentActiveEditor();
+      if (!editor) return "No active editor, change selection to view AST.";
+
+      const res = await tinymist.exportAst(editor.document.uri.fsPath, {
+        range: (await tinymist.clientPromise).code2ProtocolConverter.asRange(editor.selection),
+      });
+
+      return res || "Failed";
+    }
+
+    get onDidChange(): vscode.Event<vscode.Uri> {
+      return this.eventEmitter.event;
+    }
+  })();
+
+  ctx.subscriptions.push(vscode.workspace.registerTextDocumentContentProvider(scheme, AstDoc));
+
+  return async () => {
+    const document = await vscode.workspace.openTextDocument(AstDoc.uri);
+    setTimeout(() => AstDoc.emitChange(), 10);
+    void (await vscode.window.showTextDocument(document, {
+      viewColumn: vscode.ViewColumn.Two,
+      preserveFocus: true,
+    }));
+  };
+}
+
 async function commandClearCache(): Promise<void> {
   const activeEditor = window.activeTextEditor;
   if (activeEditor === undefined) {

editors/vscode/src/features/testing/debug.ts

@@ -2,6 +2,7 @@ import * as vscode from "vscode";
 import { WorkspaceFolder, DebugConfiguration, ProviderResult, CancellationToken } from "vscode";
 import { IContext } from "../../context";
 import { DebugAdapterExecutableFactory } from "../../dap";
+import { isTypstDocument } from "../../util";
 
 export const TYPST_DEBUGGER_TYPE = "myriaddreamin.typst-debugger";
 
@@ -99,7 +100,7 @@ class TypstConfigurationProvider implements vscode.DebugConfigurationProvider {
     // if launch.json is missing or empty
     if (!config.type && !config.request && !config.name) {
       const editor = vscode.window.activeTextEditor;
-      if (editor && editor.document.languageId === "typst") {
+      if (isTypstDocument(editor?.document)) {
         config.type = TYPST_DEBUGGER_TYPE;
         config.name = "Launch";
         config.request = "launch";

editors/vscode/src/features/tool.ts

@@ -2,7 +2,7 @@ import * as vscode from "vscode";
 import { writeFile } from "fs/promises";
 import { tinymist } from "../lsp";
 import { extensionState, ExtensionContext } from "../state";
-import { activeTypstEditor, base64Encode, loadHTMLFile } from "../util";
+import { activeTypstEditor, base64Encode, isTypstDocument, loadHTMLFile } from "../util";
 import { IContext } from "../context";
 
 const USER_PACKAGE_VERSION = "0.0.1";
@@ -422,7 +422,7 @@ export async function editorToolAt(
           if (disposed) {
             return;
           }
-          if (!event.textEditor || event.textEditor.document.languageId !== "typst") {
+          if (!isTypstDocument(event.textEditor.document)) {
             return;
           }
           version += 1;

editors/vscode/src/lsp.ts

@@ -286,6 +286,7 @@ export class LanguageState {
   exportText = exportCommand("tinymist.exportText");
   exportQuery = exportCommand("tinymist.exportQuery");
   exportAnsiHighlight = exportCommand("tinymist.exportAnsiHighlight");
+  exportAst = exportCommand("tinymist.exportAst");
 
   getResource<T extends keyof ResourceRoutes>(path: T, ...args: any[]) {
     return tinymist.executeCommand<ResourceRoutes[T]>("tinymist.getResources", [path, ...args]);

editors/vscode/src/util.ts

@@ -38,12 +38,18 @@ export function translateExternalURL(urlStr: string): string {
 
 export function activeTypstEditor() {
   const editor = vscode.window.activeTextEditor;
-  if (!editor || editor.document.languageId !== "typst") {
+  if (!isTypstDocument(editor?.document)) {
     return;
   }
   return editor;
 }
 
+export function isTypstDocument(
+  document: vscode.TextDocument | undefined,
+): document is vscode.TextDocument & { languageId: "typst" } {
+  return document?.languageId === "typst" && !document.uri.scheme.startsWith("tinymist");
+}
+
 export function getTargetViewColumn(viewColumn: ViewColumn | undefined): ViewColumn {
   if (viewColumn === ViewColumn.One) {
     return ViewColumn.Two;

locales/tinymist-vscode.toml

@@ -237,6 +237,10 @@ vi = "Sao chép dưới dạng mã ANSI"
 zh = "复制为 ANSI 代码"
 zh-TW = "複製為 ANSI 代碼"
 
+[extension.tinymist.command.tinymist.viewAst]
+en = "View the AST of the current file"
+zh = "查看当前文件的 AST"
+
 [extension.tinymist.command.tinymist.showLog]
 en = "Tinymist: Show Log"
 ar = "Tinymist: عرض السجل"

tests/e2e/main.rs

@@ -377,7 +377,7 @@ fn e2e() {
         });
 
         let hash = replay_log(&tinymist_binary, &root.join("neovim"));
-        insta::assert_snapshot!(hash, @"siphash128_13:ce179598883927533514674aa7930054");
+        insta::assert_snapshot!(hash, @"siphash128_13:94308cd99e254c72703bcc3e1386a80");
     }
 
     {
@@ -388,7 +388,7 @@ fn e2e() {
         });
 
         let hash = replay_log(&tinymist_binary, &root.join("vscode"));
-        insta::assert_snapshot!(hash, @"siphash128_13:9a34d499b2c94fad67d2f700633bb9bf");
+        insta::assert_snapshot!(hash, @"siphash128_13:3cce756a8b70867a8041cfab6f4337ef");
     }
 }